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 dominent".
    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"
    )
}

# --- 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()