workoutMakerMCP/workout_server.py

from typing import List, Optional
from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field, ConfigDict
import random

# Initialize FastMCP server
mcp = FastMCP("workout_mcp")

# --- Data Models ---

class ExerciseType(str):
    PUSH = "push"
    PULL = "pull"
    KNEE_HIP = "knee_hip" # Combined for simplicity or split if needed. Prompt said "knee/hip dominant".
    CORE = "core"
    OTHER = "other"

class Exercise(BaseModel):
    name: str
    type: str
    steps: List[str]
    image_prompt: str

# Sample Database (in-memory)
# In a real app, this might come from a DB or be generated.
EXERCISE_DB = {
    "pushup": Exercise(
        name="Push Up",
        type=ExerciseType.PUSH,
        steps=["Start in plank position", "Lower body until chest nearly touches floor", "Push back up"],
        image_prompt="A fit character performing a pushup, gym background, side view, correct form"
    ),
    "pullup": Exercise(
        name="Pull Up",
        type=ExerciseType.PULL,
        steps=["Hang from bar", "Pull chin over bar", "Lower down"],
        image_prompt="A fit character performing a pullup, gym background, back view, correct form"
    ),
    "squat": Exercise(
        name="Bodyweight Squat",
        type=ExerciseType.KNEE_HIP,
        steps=["Stand feet shoulder width", "Lower hips back and down", "Stand back up"],
        image_prompt="A fit character performing a squat, gym background, side view, correct form"
    ),
    "plank": Exercise(
        name="Plank",
        type=ExerciseType.CORE,
        steps=["Forearms on ground", "Body straight", "Hold"],
        image_prompt="A fit character holding a plank position, gym background, side view"
    ),
     "lunges": Exercise(
        name="Walking Lunges",
        type=ExerciseType.KNEE_HIP,
        steps=["Step forward with one leg", "Lower hips until both knees are bent at a 90-degree angle", "Push off with your front foot to return to starting position", "Repeat with the other leg"],
        image_prompt="A fit character performing lunges, gym background, side view, correct form"
    ),
    "overhead_press": Exercise(
        name="Overhead Press",
        type=ExerciseType.PUSH,
        steps=["Stand with feet shoulder-width apart", "Hold dumbbells at shoulder height", "Press weights overhead until arms are fully extended", "Lower back to starting position"],
        image_prompt="A fit character performing an overhead press with dumbbells, gym background, front view, correct form"
    ),
    "dumbbell_row": Exercise(
        name="Dumbbell Row",
        type=ExerciseType.PULL,
        steps=["Place one knee and hand on a bench", "Hold a dumbbell in the other hand", "Pull the weight up towards your hip", "Lower back down"],
        image_prompt="A fit character performing a dumbbell row, gym background, side view, correct form"
    ),
    "russian_twist": Exercise(
        name="Russian Twist",
        type=ExerciseType.CORE,
        steps=["Sit on the floor with knees bent", "Lean back slightly", "Twist torso to one side, then the other", "Hold a weight for added resistance if desired"],
        image_prompt="A fit character performing russian twists, gym background, front view, correct form"
    ),
    "bench_press": Exercise(
        name="Bench Press",
        type=ExerciseType.PUSH,
        steps=["Lie on a flat bench", "Grip barbell at shoulder width", "Lower barbell to chest", "Press barbell up until arms are fully extended"],
        image_prompt="A fit character performing a bench press with barbell, gym background, side view, correct form"
    ),
    "incline_pushup": Exercise(
        name="Incline Push Up",
        type=ExerciseType.PUSH,
        steps=["Place hands on an elevated surface", "Assume plank position with elevated hands", "Lower body until chest nearly touches surface", "Push back up"],
        image_prompt="A fit character performing an incline pushup on a bench, gym background, side view, correct form"
    ),
    "dips": Exercise(
        name="Dips",
        type=ExerciseType.PUSH,
        steps=["Grip parallel bars with arms extended", "Lower body by bending elbows", "Push back up to starting position"],
        image_prompt="A fit character performing dips on parallel bars, gym background, side view, correct form"
    ),
    "lat_pulldown": Exercise(
        name="Lat Pulldown",
        type=ExerciseType.PULL,
        steps=["Sit at lat pulldown machine", "Grip the bar with hands slightly wider than shoulder width", "Pull the bar down towards your chest", "Slowly return to starting position"],
        image_prompt="A fit character performing a lat pulldown on a machine, gym background, front view, correct form"
    ),
    "barbell_row": Exercise(
        name="Barbell Row",
        type=ExerciseType.PULL,
        steps=["Stand with feet shoulder-width apart", "Hinge at hips and grip barbell", "Row barbell towards your torso", "Lower barbell back down"],
        image_prompt="A fit character performing a barbell row, gym background, side view, correct form"
    ),
    "assisted_pullup": Exercise(
        name="Assisted Pull Up",
        type=ExerciseType.PULL,
        steps=["Use an assisted pull-up machine or resistance band", "Place knees or feet on the assistance pad", "Pull chin over the bar", "Lower down with control"],
        image_prompt="A fit character performing an assisted pullup, gym background, back view, correct form"
    ),
    "deadlift": Exercise(
        name="Deadlift",
        type=ExerciseType.KNEE_HIP,
        steps=["Stand with feet hip-width apart, barbell over mid-foot", "Bend hips and knees to grip the bar", "Drive through heels to stand up", "Lower back down with control"],
        image_prompt="A fit character performing a deadlift with barbell, gym background, side view, correct form"
    ),
    "leg_press": Exercise(
        name="Leg Press",
        type=ExerciseType.KNEE_HIP,
        steps=["Sit on leg press machine with back against pad", "Place feet on platform shoulder-width apart", "Push platform away by extending legs", "Lower platform back down"],
        image_prompt="A fit character performing a leg press on a machine, gym background, side view, correct form"
    ),
    "calf_raises": Exercise(
        name="Calf Raises",
        type=ExerciseType.KNEE_HIP,
        steps=["Stand with feet hip-width apart", "Rise up onto the balls of your feet", "Hold the top position briefly", "Lower heels back down"],
        image_prompt="A fit character performing calf raises, gym background, side view, correct form"
    ),
    "bicycle_crunch": Exercise(
        name="Bicycle Crunch",
        type=ExerciseType.CORE,
        steps=["Lie on your back with hands behind head", "Bring one elbow towards the opposite knee", "Extend the other leg", "Alternate in a cycling motion"],
        image_prompt="A fit character performing bicycle crunches, gym background, front view, correct form"
    ),
    "dead_bug": Exercise(
        name="Dead Bug",
        type=ExerciseType.CORE,
        steps=["Lie on your back with arms extended towards ceiling", "Raise legs with knees bent at 90 degrees", "Lower opposite arm and leg while keeping back on floor", "Return and alternate sides"],
        image_prompt="A fit character performing a dead bug exercise, gym background, front view, correct form"
    ),
    "mountain_climbers": Exercise(
        name="Mountain Climbers",
        type=ExerciseType.CORE,
        steps=["Start in plank position", "Bring one knee towards your chest", "Quickly switch legs in a running motion", "Maintain a steady pace"],
        image_prompt="A fit character performing mountain climbers, gym background, side view, correct form"
    )
}

# --- Inputs ---

class CreateExerciseInput(BaseModel):
    name: str = Field(..., description="Name of the exercise")
    type: str = Field(..., description="Type: push, pull, knee_hip, core")
    steps: List[str] = Field(..., description="List of step-by-step instructions")
    visual_description: str = Field(..., description="Visual description for image generation")

class CreateSessionInput(BaseModel):
    name: str = Field(..., description="Name of the session (e.g. 'Morning Blast')")
    difficulty: str = Field("beginner", description="beginner, intermediate, advanced")
    duration_minutes: int = Field(30, description="Target duration in minutes")

class CreatePlanInput(BaseModel):
    goal: str = Field(..., description="Goal: strength, weight_loss, endurance")
    age_range: str = Field(..., description="Age range (e.g. '20-30', '50+')")
    length_days: int = Field(..., description="Length of plan in days")


# --- Helpers ---

def generate_nanobanana_prompt(description: str, style: str = "photorealistic") -> str:
    """Creates a prompt optimized for NanoBanana/Gemini image generation."""
    # Enforcing the requirement: 'consistent characters' and 'simple'
    base_prompt = "A consistent character design, athletic wear, clean background. "
    return f"{base_prompt} {description}. Style: {style}, high quality."

# --- Tools ---

@mcp.tool()
def create_exercise_page(params: CreateExerciseInput) -> str:
    """Creates a Markdown page for a single exercise with Nanobanana image prompts."""
    
    image_prompt = generate_nanobanana_prompt(params.visual_description)
    
    md = f"# {params.name}\n\n"
    md += f"**Type:** {params.type.capitalize()}\n\n"
    md += "## Steps\n"
    for i, step in enumerate(params.steps, 1):
        md += f"{i}. {step}\n"
    
    md += "\n## Visual Reference\n"
    md += f"> **NanoBanana Image Prompt:**\n"
    md += f"> `/image prompt=\"{image_prompt}\" --count=1`\n\n"
    md += "*(Use the above command with your Gemini agent to generate the visual)*\n"
    
    return md

@mcp.tool()
def create_session(params: CreateSessionInput) -> str:
    """Creates a workout session grouping exercises. Ensures variety."""
    
    # Logic: Select one of each type for a balanced session
    selected_exercises = []
    
    # Simple selection logic from DB
    types_needed = [ExerciseType.PUSH, ExerciseType.PULL, ExerciseType.KNEE_HIP, ExerciseType.CORE]
    
    for t in types_needed:
        # Find candidates
        candidates = [e for e in EXERCISE_DB.values() if e.type == t]
        if candidates:
            # Randomly pick one or based on difficulty (mocking logic here)
            selected_exercises.append(random.choice(candidates))
    
    md = f"# Workout Session: {params.name}\n\n"
    md += f"**Difficulty:** {params.difficulty} | **Duration:** {params.duration_minutes} mins\n\n"
    md += "## Warmup\n- 5 mins light cardio\n\n"
    md += "## The Circuit\n"
    md += "*Perform 3 rounds of the following:*\n\n"
    
    for ex in selected_exercises:
        md += f"### {ex.name}\n"
        md += f"**Target:** {ex.type.capitalize()}\n"
        md += f"- Instructions: {ex.steps[0]}...\n"
        md += f"- [View Full Exercise Page via `create_exercise_page`]\n\n"
    
    md += "## Cooldown\n- 5 mins stretching\n"
    
    return md

@mcp.tool()
def create_workout_plan(params: CreatePlanInput) -> str:
    """Creates a multi-day workout plan based on goals."""
    
    md = f"# {params.length_days}-Day Workout Plan: {params.goal.capitalize()}\n"
    md += f"**Target Audience:** Age {params.age_range}\n\n"
    
    # Logic to adjust intensity based on age/goal
    rest_frequency = 3 # Default: rest every 3rd day
    if "50+" in params.age_range or "beginner" in params.goal:
        rest_frequency = 2
        
    for day in range(1, params.length_days + 1):
        md += f"## Day {day}\n"
        
        if day % rest_frequency == 0:
            md += "**Rest & Recovery**\n- Light walking or stretching.\n\n"
        else:
            session_name = f"Day {day} {params.goal} Focus"
            md += f"**Session:** {session_name}\n\n"
            
            # Select exercises for this day to meet requirements
            daily_exercises = []
            types_needed = [ExerciseType.PUSH, ExerciseType.PULL, ExerciseType.KNEE_HIP, ExerciseType.CORE]
            
            for t in types_needed:
                candidates = [e for e in EXERCISE_DB.values() if e.type == t]
                if candidates:
                    daily_exercises.append(random.choice(candidates))
            
            md += "| Type | Exercise | Sets/Reps |\n"
            md += "|---|---|---|\n"
            for ex in daily_exercises:
                 # Adjust reps based on goal
                reps = "3x10"
                if params.goal == "strength": reps = "5x5"
                elif params.goal == "endurance": reps = "3x15+"
                
                md += f"| {ex.type.capitalize()} | {ex.name} | {reps} |\n"
            
            md += "\n"

    return md

if __name__ == "__main__":
    mcp.run()