import json
from pydantic import BaseModel, Field
from APIs.aichat import generate


# 1. Define the output schema as a Pydantic model
class DishSafetyResult(BaseModel):
    is_safe: bool = Field(
        description="True if no microwave hazards presents."
    )
    warning_message: str = Field(
        description="Short explanation of any hazard found (no more than 50 characters), or an empty string if safe."
    )
    detected_hazards: list[str] = Field(
        description="List of specific hazard items detected."
    )


def check_dish_safety(image_path: str) -> dict:
    prompt = (
        "Analyze this top-down photo of a dish prepared for microwave cooking. "
        "Inspect the area for metal utensils, aluminum foil, metallic dish patterns, or unvented plastic wraps that could be dangerous for the dish microwave cooking process or pose a safety risk."
    )

    # 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)}")