import logging
logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

from typing import Any

import gradio as gr
from PIL import Image

from src.agents.mise_en_place import identify_ingredients
from src.agents.progress_validator import validate
from src.agents.recipe_planner import plan_recipe, propose_dishes
from src.agents.step_illustrator import illustrate_recipe
from src.data.nutrition import compute_nutrition
from src.ui.components import (
    DishOptions,
    IngredientChips,
    NutritionGrid,
    RecipeHero,
    StepCard,
    VerdictBadge,
)
from src.ui.theme import CSS, theme


# ---------------------------------------------------------------------------
# Callbacks
# ---------------------------------------------------------------------------

def _clean_ingredients(items: list | None) -> list[str]:
    """Normalize a raw ingredient list (dedup, lowercase, strip empties)."""
    out, seen = [], set()
    for it in (items or []):
        name = str(it).strip().lower()
        if name and name not in seen:
            seen.add(name)
            out.append(name)
    return out


def on_propose(fridge_image: Image.Image | None, state: dict | None):
    """Photo → ingredients → 3 dish options (and fill the editable list)."""
    state = state or {}
    if fridge_image is None:
        return (
            IngredientChips.render({}),
            DishOptions.render({}),
            gr.update(choices=[], value=None),
            state,
            gr.update(choices=[], value=[]),
        )

    ingredients = identify_ingredients(fridge_image)
    options = propose_dishes(ingredients)

    state.update({
        "ingredients_have": ingredients,
        "options": [o.model_dump() for o in options],
    })

    radio_choices = [o.name for o in options]
    return (
        IngredientChips.render({"have": ingredients, "missing": []}),
        DishOptions.render({"options": state["options"]}),
        gr.update(choices=radio_choices, value=radio_choices[0] if radio_choices else None),
        state,
        gr.update(choices=ingredients, value=ingredients),
    )


def on_update_ingredients(state: dict | None, ingredients: list | None):
    """Manual edit of the ingredient list → refresh chips + re-propose dishes."""
    state = state or {}
    ingredients = _clean_ingredients(ingredients)
    state["ingredients_have"] = ingredients

    if not ingredients:
        state["options"] = []
        return (
            IngredientChips.render({}),
            DishOptions.render({}),
            gr.update(choices=[], value=None),
            state,
        )

    options = propose_dishes(ingredients)
    state["options"] = [o.model_dump() for o in options]
    radio_choices = [o.name for o in options]
    return (
        IngredientChips.render({"have": ingredients, "missing": []}),
        DishOptions.render({"options": state["options"]}),
        gr.update(choices=radio_choices, value=radio_choices[0] if radio_choices else None),
        state,
    )


def on_cook(state: dict | None, dish_name: str | None, illustrate: bool, ingredients: list | None):
    """Chosen dish → full recipe + nutrition (+ FLUX images if requested)."""
    state = state or {}
    if not dish_name:
        return (
            RecipeHero.render({}),
            StepCard.render({}),
            NutritionGrid.render({"nutrition": {}}),
            state,
        )

    # Prefer the (possibly hand-edited) ingredient list from the editor.
    ingredients = _clean_ingredients(ingredients) or state.get("ingredients_have", [])
    state["ingredients_have"] = ingredients
    recipe = plan_recipe(dish_name, ingredients)

    nutrition = compute_nutrition(ingredients, recipe.servings)
    recipe.nutrition = nutrition
    state["recipe"] = recipe.model_dump()

    if illustrate:
        log.info("Generating FLUX step images via Modal...")
        recipe = illustrate_recipe(recipe)
        state["recipe"] = recipe.model_dump()

    return (
        RecipeHero.render(recipe.model_dump()),
        StepCard.render({"steps": [s.model_dump() for s in recipe.steps]}),
        NutritionGrid.render({"nutrition": nutrition}),
        state,
    )


def on_validate(state: dict | None, step_idx: float, progress_image: Image.Image | None):
    """Progress photo + step number → verdict badge."""
    state = state or {}
    recipe = state.get("recipe", {})
    steps = recipe.get("steps", [])
    idx = max(0, int(step_idx) - 1)
    instruction = steps[idx]["instruction"] if idx < len(steps) else "Cook the dish properly."
    result = validate(progress_image, instruction)
    return VerdictBadge.render(result)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

def build_ui() -> gr.Blocks:
    initial_state: dict[str, Any] = {}

    with gr.Blocks(title="Cook With Me", theme=theme, css=CSS) as demo:
        gr.Markdown(
            "# 🍲 Cook With Me\n"
            "_Snap your fridge · Pick a dish · Cook step by step · Check your progress._"
        )

        state = gr.State(initial_state)

        with gr.Tabs():
            # ----------------------------------------------------------------
            # Tab 1 — Cook
            # ----------------------------------------------------------------
            with gr.Tab("🍳 Cook"):
                with gr.Row():
                    # Left — inputs
                    with gr.Column(scale=1):
                        fridge_input = gr.Image(
                            label="📸 Photo of your fridge or pantry",
                            type="pil",
                            height=300,
                        )
                        propose_btn = gr.Button("🔍 What can I cook?", variant="primary")

                        gr.Markdown("### Ingredients I see")
                        chips = gr.HTML(IngredientChips.render({}))

                        ingredient_editor = gr.Dropdown(
                            choices=[],
                            value=[],
                            multiselect=True,
                            allow_custom_value=True,
                            label="✏️ Add or remove ingredients (type + Enter to add, ✕ to remove)",
                            interactive=True,
                        )
                        update_btn = gr.Button("🔄 Update ingredients & dishes")

                        gr.Markdown("### Pick a dish")
                        dish_options_html = gr.HTML(DishOptions.render({}))
                        dish_radio = gr.Radio(
                            choices=[],
                            label="Choose one",
                            interactive=True,
                        )

                        with gr.Accordion("⚙️ Generation options", open=False):
                            illustrate_chk = gr.Checkbox(
                                value=False,
                                label="🎨 Generate step images with FLUX.2 (requires Modal deployment)",
                            )

                        cook_btn = gr.Button("👨‍🍳 Build my recipe", variant="primary")

                    # Right — recipe output
                    with gr.Column(scale=2):
                        hero = gr.HTML(RecipeHero.render({}))
                        steps_panel = gr.HTML(StepCard.render({}))
                        nutrition_panel = gr.HTML(NutritionGrid.render({"nutrition": {}}))

            # ----------------------------------------------------------------
            # Tab 2 — Check Progress
            # ----------------------------------------------------------------
            with gr.Tab("📷 Check Progress"):
                gr.Markdown(
                    "Upload a photo of your pan or plate. The vision model compares it "
                    "against the current recipe step and tells you if you can move on."
                )
                with gr.Row():
                    with gr.Column():
                        step_idx = gr.Number(value=1, precision=0, label="Active step #")
                        progress_input = gr.Image(
                            label="📸 Your pan / plate",
                            type="pil",
                            height=300,
                        )
                        validate_btn = gr.Button("✅ How am I doing?", variant="primary")
                    with gr.Column():
                        verdict_panel = gr.HTML(VerdictBadge.render({}))

            # ----------------------------------------------------------------
            # Tab 3 — About
            # ----------------------------------------------------------------
            with gr.Tab("ℹ️ About"):
                gr.Markdown(
                    """
### How it works
1. **Snap** your fridge — the fine-tuned vision model (MiniCPM-V-4.6) identifies every ingredient.
2. **Pick** one of three AI-suggested dishes tailored to what you have.
3. **Cook** step by step with a generated recipe, per-serving nutrition, and optional FLUX.2 step images.
4. **Check** your progress — upload a photo of your pan and get a *go / wait / fix* verdict.

### Models
| Role | Model | Params |
|---|---|---|
| Vision (ingredients + validator) | `openbmb/MiniCPM-V-4.6` (fine-tuned) | ~4.6B |
| Recipe Planner | `openbmb/MiniCPM4.1-8B` (fine-tuned on Kaggle recipes) | ~8B |
| Step Illustrator | `FLUX.2-klein-9B` via Modal | ~9B |

**Total ≤ 21.6B params** (cap: 32B ✓)

### Badges targeted
✓ Well-Tuned · ✓ Off-Brand · ✓ Sharing is Caring · ✓ Field Notes

### Hackathon
Hugging Face Small Models / Big Adventures · June 2026 · Track: Backyard AI
                    """
                )

        # --------------------------------------------------------------------
        # Wire callbacks
        # --------------------------------------------------------------------
        propose_btn.click(
            fn=on_propose,
            inputs=[fridge_input, state],
            outputs=[chips, dish_options_html, dish_radio, state, ingredient_editor],
        )

        update_btn.click(
            fn=on_update_ingredients,
            inputs=[state, ingredient_editor],
            outputs=[chips, dish_options_html, dish_radio, state],
        )

        cook_btn.click(
            fn=on_cook,
            inputs=[state, dish_radio, illustrate_chk, ingredient_editor],
            outputs=[hero, steps_panel, nutrition_panel, state],
        )

        validate_btn.click(
            fn=on_validate,
            inputs=[state, step_idx, progress_input],
            outputs=[verdict_panel],
        )

    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(__import__("os").environ.get("PORT", 7860)),
        show_error=True,
        inbrowser=True,
    )