import base64

import dspy
from pydantic import BaseModel, Field

from app.config import settings


class NutritionalInfo(BaseModel):
    reasoning: str = Field(description="Step-by-step reasoning for the nutritional estimates")
    name: str = Field(description="Name of the food item")
    calories: float = Field(description="Estimated calories")
    protein: float = Field(description="Estimated protein in grams")
    carbs: float = Field(description="Estimated carbohydrates in grams")
    fats: float = Field(description="Estimated fats in grams")
    micros: dict | None = None


class ExtractNutrition(dspy.Signature):
    """Extract nutritional information from a food description.

    You must first provide a detailed step-by-step reasoning analysis of the ingredients,
    portions, AND preparation methods (cooking oils, butter, sauces) before estimating values.
    Verify if the caloric totals match the sum of macros (multiplying protein/carbs by 4, fats by 9).
    """

    description: str = dspy.InputField(desc="Description of the food or meal")
    nutritional_info: NutritionalInfo = dspy.OutputField(desc="Nutritional information with reasoning")


class AnalyzeFoodImage(dspy.Signature):
    """Analyze the food image to estimate nutritional content.

    1. Identify all food items and estimated portion sizes.
    2. CRITICAL: Account for hidden calories from cooking fats, oils, and sauces (searing, frying).
    3. Reason step-by-step about the total composition before summing macros.
    """

    image: dspy.Image = dspy.InputField(desc="The food image")
    description: str = dspy.InputField(desc="Additional user description", default="")
    nutritional_info: NutritionalInfo = dspy.OutputField(desc="Nutritional information with reasoning")


class NutritionModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.extract = dspy.ChainOfThought(ExtractNutrition)
        self.analyze_image = dspy.ChainOfThought(AnalyzeFoodImage)

        # Load optimized prompts if available
        import os

        compiled_path = os.path.join(os.path.dirname(__file__), "nutrition_compiled.json")
        if os.path.exists(compiled_path):
            self.load(compiled_path)
            print(f"Loaded optimized DSPy prompts from {compiled_path}")
        else:
            print("No optimized prompts found, using default zero-shot.")

    def forward(self, description: str):
        pred = self.extract(description=description)

        # Assertion: Check Macro Consistency
        calc_cals = (
            (pred.nutritional_info.protein * 4) + (pred.nutritional_info.carbs * 4) + (pred.nutritional_info.fats * 9)
        )

        # dspy.Suggest is not available in dspy>=3.1.0
        # dspy.Suggest(
        #     abs(calc_cals - pred.nutritional_info.calories) < (pred.nutritional_info.calories * 0.20),
        #     f"The sum of macros ({calc_cals:.1f}) should match the total calories "
        #     f"({pred.nutritional_info.calories}). Check your math.",
        # )
        return pred

    def forward_image(self, image_url: str, description: str = ""):
        image = dspy.Image(image_url)
        pred = self.analyze_image(image=image, description=description)

        # Assertion: Check Macro Consistency
        calc_cals = (
            (pred.nutritional_info.protein * 4) + (pred.nutritional_info.carbs * 4) + (pred.nutritional_info.fats * 9)
        )

        # dspy.Suggest is not available in dspy>=3.1.0
        # dspy.Suggest(
        #     abs(calc_cals - pred.nutritional_info.calories) < (pred.nutritional_info.calories * 0.20),
        #     f"The sum of macros ({calc_cals:.1f}) should match the total calories "
        #     f"({pred.nutritional_info.calories}). Check your math.",
        # )
        return pred


nutrition_module = NutritionModule()


def analyze_nutrition_from_image(image_bytes: bytes, description: str = "") -> NutritionalInfo:
    if not settings.OPENAI_API_KEY:
        raise ValueError("OpenAI API Key not set")

    # Convert to base64 data URI
    base64_image = base64.b64encode(image_bytes).decode("utf-8")
    image_url = f"data:image/jpeg;base64,{base64_image}"

    # Use DSPy module
    result = nutrition_module.forward_image(image_url=image_url, description=description)
    return result.nutritional_info