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

class UtensilMaterialCheck(BaseModel):
    object_name: str = Field(
        description="Name of utensil or container seen (e.g., spoon, fork, tray)."
    )
    is_metal: bool = Field(
        description="True ONLY if the object is made of metal, stainless steel, or aluminum."
    )
    
class DishSafetyResult(BaseModel):
    utensils_and_containers: list[UtensilMaterialCheck] = Field(
        description="List each visible non-food object and check if it is made of metal."
    )
    material_analysis: str = Field(
        description="Short description of the materials present."
    )
    detected_hazards: list[str] = Field(
        description="List objects from utensils_and_containers where is_metal is True."
    )
    is_safe: bool = Field(
        description="True if detected_hazards is completely empty []. False if ANY metal is found."
    )
    warning_message: str = Field(
        description="If is_safe is False, set to 'REMOVE METAL UTENSIL OR FOIL BEFORE MICROWAVING'. Otherwise empty ''."
    )

def preprocess_dish_image(image_path: str) -> str:
    """Brightens dark areas and sharpens food texture for small vision models."""
    img = Image.open(image_path)
    
    # 1. Boost brightness slightly
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(1.4)
    
    # 2. Boost contrast to separate food shapes from dark shadows
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(1.3)
    
    processed_path = "/tmp/processed_dish.jpg"
    img.save(processed_path, quality=85)
    return processed_path

def check_dish_safety(image_path: str) -> dict:
    prompt = (
        "Analyze this microwave dish image strictly for EQUIPMENT DAMAGE HAZARDS (Metal, Steel, Aluminum Foil).\n\n"
        "GOAL:\n"
        "Identify if any metallic item (metal spoon, fork, knife, aluminum foil, wire) is inside the dish.\n\n"
        "RULES:\n"
        "1. Examine all utensils carefully, even in dim lighting or dark spots.\n"
        "2. Spoons, forks, and knives are often STAINLESS STEEL / METAL. If a spoon or fork is visible, mark is_metal = True unless it is clearly bright colored plastic.\n"
        "3. Ignore chemical or health concerns (plastic cancer risks are IRRELEVANT). Focus ONLY on spark/fire hazards (metal, foil).\n"
        "4. If ANY utensil is metal: set is_metal=True, add it to detected_hazards, and set is_safe=False."
    )

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

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