import json
from pydantic import BaseModel, Field
from APIs.aichat import generate
import shared.config as config


# Output schema as a Pydantic model
class DishSafetyResult(BaseModel):
    visible_objects: list[str] = Field(
        description="List all distinct physical objects visible in or around the dish (e.g., bowl, plastic tray, liquid, spoon, cover)."
    )
    material_analysis: str = Field(
        description="Analyze the physical material of each visible object (e.g., ceramic, polypropylene plastic, stainless steel, aluminum foil)."
    )
    detected_hazards: list[str] = Field(
        description=(
            "List ONLY explicit, high-risk microwave hazards: metal items, cutlery, aluminum foil, "
            "metallic trim/gilding, Styrofoam (expanded polystyrene), or tightly sealed/unvented foil lids. "
            "Do NOT include standard plastic containers, TV dinner trays, polypropylene (#5), or Tupperware as hazards. Empty if none."
        )
    )
    is_safe: bool = Field(
        description="Must be set to True if detected_hazards is empty. Otherwise False."
    )
    warning_message: str = Field(
        description="One short sentence explaining the hazard if detected_hazards is not empty, otherwise an empty string."
    )


def check_dish_safety(image_path: str) -> dict:
    prompt = (
        "Examine this photo of a meal intended for microwave heating.\n\n"
        "Inspect all objects to determine if any CRITICAL microwave hazards exist.\n\n"
        "STRICT HAZARDS TO DETECT:\n"
        "- Metal utensils, cutlery, or metal objects.\n"
        "- Aluminum foil, metallic packaging, or decorative metallic trim.\n"
        "- Expanded Polystyrene / Styrofoam containers.\n"
        "- Completely sealed non-vented foil or plastic film lids (explosion risk).\n\n"
        "SAFE MATERIALS (DO NOT FLAG AS HAZARDS):\n"
        "- Standard microwavable food containers, plastic meal prep trays, black plastic TV dinner trays, and polypropylene (PP / #5) plastics.\n"
        "- Glass, ceramic, or paper containers.\n\n"
        "Only mark is_safe as False if a clear, high-risk hazard from the strictly dangerous list above is present."
    )

    # 2. Pass the Pydantic class directly to generate()
    response_raw = generate(
        prompt=prompt,
        images=[image_path],
        output_format=DishSafetyResult,
        should_think=False,
    )
    print(f"[Debug] Raw response from AI generator: {response_raw}")

    # 3. Handle response parsing
    if isinstance(response_raw, str):
        # Parse and validate the JSON string into the Pydantic model, then return as a dict
        try:
            validated_result = DishSafetyResult.model_validate_json(response_raw)
            return validated_result.model_dump()
        except Exception:
            # Fallback to standard json.loads if raw parsing is needed
            return json.loads(response_raw)

    if isinstance(response_raw, DishSafetyResult):
        return response_raw.model_dump()

    if isinstance(response_raw, dict):
        return response_raw

    raise ValueError(f"Unexpected response type from AI generator: {type(response_raw)}")