import dspy
from pydantic import BaseModel, Field


class ExerciseBlock(BaseModel):
    order: int = Field(description="Position in workout sequence")
    name: str = Field(description="Exercise name")
    description: str = Field(description="Brief description of the movement")
    sets: int = Field(description="Number of sets")
    reps: int = Field(description="Number of reps per set (0 if timed)")
    duration_seconds: int = Field(description="Duration in seconds per set (0 if rep-based)")
    weight_kg: float = Field(description="Prescribed weight in kg")
    rest_seconds: int = Field(description="Rest time between sets in seconds")
    coaching_tip: str = Field(description="Key coaching cue for this exercise")


class KettlebellSessionOutput(BaseModel):
    reasoning: str = Field(description="Step-by-step reasoning for session design choices")
    title: str = Field(description="Session title")
    focus: str = Field(description="Session focus e.g. strength, conditioning, mobility")
    total_duration_min: int = Field(description="Estimated total workout duration in minutes")
    difficulty: str = Field(description="Difficulty level: beginner, intermediate, advanced")
    exercises: list[ExerciseBlock] = Field(description="Ordered list of exercises in the session")
    notes: str = Field(description="Coaching notes and any special instructions for the session")


class GenerateKettlebellSession(dspy.Signature):
    """Generate a personalized kettlebell workout session based on user profile and preferences.

    Think step-by-step: assess user fitness level, pick movements appropriate to the focus and
    difficulty, assign weights respecting progressive overload principles from available weights,
    sequence exercises for proper warm-up and fatigue management, and ensure total work time
    (sets × reps/duration + rest periods) fits within the requested duration.
    """

    user_profile: str = dspy.InputField(desc="User details including age, weight, fitness level, and goals")
    available_weights_kg: str = dspy.InputField(desc="Comma-separated list of available kettlebell weights in kg")
    focus: str = dspy.InputField(desc="Session focus: strength, conditioning, mobility, fat loss, etc.")
    duration_minutes: int = dspy.InputField(desc="Target session duration in minutes")
    session: KettlebellSessionOutput = dspy.OutputField(desc="Complete structured kettlebell session")


class KettlebellModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.generate = dspy.ChainOfThought(GenerateKettlebellSession)

    def forward(self, user_profile: str, available_weights_kg: str, focus: str, duration_minutes: int):
        return self.generate(
            user_profile=user_profile,
            available_weights_kg=available_weights_kg,
            focus=focus,
            duration_minutes=duration_minutes,
        )


kettlebell_module = KettlebellModule()