import dspy
from pydantic import BaseModel, Field


class PlanOutput(BaseModel):
    reasoning: str = Field(description="Reasoning behind the selected plan based on user goals")
    title: str = Field(description="Title of the plan")
    summary: str = Field(description="Brief summary of the plan")
    diet_plan: list[str] = Field(description="List of daily diet recommendations")
    exercise_plan: list[str] = Field(description="List of daily exercise routines")
    tips: list[str] = Field(description="Additional health tips")


class GeneratePlan(dspy.Signature):
    """Generate a personalized diet and exercise plan based on user goal and details.

    Analyze the user's profile and goal, explain your reasoning, and then generate the plan.
    """

    user_profile: str = dspy.InputField(desc="User details (age, weight, height, etc)")
    goal: str = dspy.InputField(desc="Specific user goal")
    plan: PlanOutput = dspy.OutputField(desc="Structured plan with reasoning")


class PlanModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.generate = dspy.ChainOfThought(GeneratePlan)

    def forward(self, user_profile: str, goal: str):
        return self.generate(user_profile=user_profile, goal=goal)


plan_module = PlanModule()