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

class DishSafetyResult(BaseModel):
    step_1_visible_items: list[str] = Field(
        description="List ONLY 2-3 broad visible items (e.g., ['ceramic bowl', 'brown stew']). Do NOT guess objects if blurry."
    )
    step_2_has_metal_or_cutlery: bool = Field(
        description="Is a metal spoon, fork, knife, or foil clearly visible? True ONLY if distinct metal is seen, otherwise False."
    )
    detected_hazards: list[str] = Field(
        description="List ONLY physical metal items (e.g. ['metal spoon']). MUST be an empty list [] if step_2 is False."
    )
    is_safe: bool = Field(
        description="MUST be True if step_2 is False. Set to False ONLY if metal cutlery/foil is present."
    )
    warning_message: str = Field(
        description="Short warning if unsafe, otherwise empty string ''."
    )

def check_dish_safety(image_path: str) -> dict:
    prompt = (
        "You are a strict microwave safety vision inspector analyzing a top-down frame.\n\n"
        "RULES:\n"
        "1. First, list basic items in `step_1_visible_items`.\n"
        "2. Examine the dish for shiny metallic cutlery, forks, spoons, or aluminum foil.\n"
        "3. Set `step_2_has_metal_or_cutlery` to True ONLY if clear metallic cutlery is visible.\n"
        "4. If the image is blurry and no metal is clearly identified, assume NO metal is present (step_2 = False, is_safe = True)."
    )

    response_raw = generate(
        prompt=prompt,
        images=[image_path],
        output_format=DishSafetyResult,
        should_think=False,
    )
    
    if config.DEBUG_MESSAGES:
        print(f"[Debug] Raw response from AI generator: {response_raw}")

    # Handling response mapping...
    if isinstance(response_raw, str):
        try:
            validated_result = DishSafetyResult.model_validate_json(response_raw)
            return validated_result.model_dump()
        except Exception:
            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)}")