from __future__ import annotations

import json
from pathlib import Path

import gradio as gr

from decline import (
    OUTPUT_ROOT,
    analyze_with_nano_banana,
    bundle_outputs,
    ensure_dirs,
    image_from_upload,
    run_nano_banana_clean,
    save_image,
    to_overlay,
    write_json,
)

APP_TITLE = "Nano Banana 2 Studio"
DEFAULT_ANALYSIS_REQUEST = (
    "Lis tout le texte visible, identifie les blocs, explique la structure de la mise en page, "
    "et prépare un clean plate précis."
)
DEFAULT_CLEAN_INSTRUCTION = (
    "Retire tout le texte visible, les logos typographiques et les chiffres. "
    "Garde la composition, la lumière, les ombres et le fond naturels."
)



def _image_or_none(value):
    return image_from_upload(value)



def _json(text: object) -> str:
    return json.dumps(text, ensure_ascii=False, indent=2)



def analyze_ui(image, analysis_request):
    img = _image_or_none(image)
    if img is None:
        return None, "Aucune image fournie.", None

    analysis = analyze_with_nano_banana(img, user_request=analysis_request)
    overlay = to_overlay(img, analysis)
    ensure_dirs()
    overlay_path = save_image(overlay, OUTPUT_ROOT / "analysis_overlay.png")
    json_path = write_json(OUTPUT_ROOT / "analysis.json", analysis)
    bundle = bundle_outputs([overlay_path, json_path], "analysis_bundle.zip")
    status = f"Backend: {analysis.get('backend')} | {analysis.get('status')} | blocs: {len(analysis.get('detected_boxes', []))}"
    return overlay, _json(analysis), status, str(bundle)



def clean_ui(image, analysis_request, clean_instruction, seed):
    img = _image_or_none(image)
    if img is None:
        return None, None, "Aucune image fournie.", None, None

    analysis, clean_img, clean_status = run_nano_banana_clean(
        img,
        user_request=analysis_request,
        extra_instruction=clean_instruction,
        seed=int(seed) if seed not in (None, "") else None,
    )
    overlay = to_overlay(img, analysis)
    ensure_dirs()
    clean_path = save_image(clean_img, OUTPUT_ROOT / "clean_plate.png")
    overlay_path = save_image(overlay, OUTPUT_ROOT / "clean_overlay.png")
    json_path = write_json(OUTPUT_ROOT / "clean_analysis.json", analysis)
    bundle = bundle_outputs([clean_path, overlay_path, json_path], "clean_bundle.zip")
    status = f"{analysis.get('backend')} | {clean_status}"
    return overlay, clean_img, _json(analysis), status, str(bundle)



def full_ui(image, analysis_request, clean_instruction, seed):
    img = _image_or_none(image)
    if img is None:
        return None, None, None, "Aucune image fournie.", None

    analysis, clean_img, clean_status = run_nano_banana_clean(
        img,
        user_request=analysis_request,
        extra_instruction=clean_instruction,
        seed=int(seed) if seed not in (None, "") else None,
    )
    overlay = to_overlay(img, analysis)
    ensure_dirs()
    analysis_path = write_json(OUTPUT_ROOT / "full_analysis.json", analysis)
    clean_path = save_image(clean_img, OUTPUT_ROOT / "full_clean_plate.png")
    overlay_path = save_image(overlay, OUTPUT_ROOT / "full_overlay.png")
    bundle = bundle_outputs([analysis_path, clean_path, overlay_path], "full_bundle.zip")
    status = f"Backend: {analysis.get('backend')} | {analysis.get('status')} | {clean_status}"
    return overlay, clean_img, _json(analysis), status, str(bundle)



def main():
    ensure_dirs()
    with gr.Blocks(title=APP_TITLE, theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            "# Nano Banana 2 Studio\n"
            "Version de départ centrée sur 2 choses: la compréhension du texte et le clean plate.\n"
            "\n"
            "Backend attendu:\n"
            "- `NANO_BANANA_API_KEY` ou `GOOGLE_API_KEY`\n"
            "- modèle par défaut: `gemini-3.1-flash-image`\n"
            "\n"
            "Si la clé manque, l'app retombe sur les heuristiques locales et ComfyUI pour garder un résultat exploitable."
        )

        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="Image source", type="pil")
                analysis_request = gr.Textbox(
                    label="Consigne de compréhension",
                    value=DEFAULT_ANALYSIS_REQUEST,
                    lines=4,
                )
                clean_instruction = gr.Textbox(
                    label="Consigne clean plate",
                    value=DEFAULT_CLEAN_INSTRUCTION,
                    lines=3,
                )
                seed = gr.Number(label="Seed optionnelle", value=42, precision=0)
                btn_analyze = gr.Button("Comprendre le texte", variant="secondary")
                btn_clean = gr.Button("Générer le clean plate", variant="primary")
                btn_full = gr.Button("Faire les deux", variant="primary")

            with gr.Column(scale=1):
                status = gr.Textbox(label="Statut", lines=3)
                bundle = gr.File(label="Télécharger le bundle")
                with gr.Tabs():
                    with gr.Tab("Compréhension"):
                        overlay = gr.Image(label="Text overlay / zones détectées")
                        analysis_json = gr.Code(label="JSON d'analyse", language="json")
                    with gr.Tab("Clean plate"):
                        clean_image = gr.Image(label="Clean plate")
                        clean_json = gr.Code(label="JSON + notes", language="json")

        btn_analyze.click(
            fn=analyze_ui,
            inputs=[image, analysis_request],
            outputs=[overlay, status, analysis_json, bundle],
        )
        btn_clean.click(
            fn=clean_ui,
            inputs=[image, analysis_request, clean_instruction, seed],
            outputs=[overlay, clean_image, clean_json, status, bundle],
        )
        btn_full.click(
            fn=full_ui,
            inputs=[image, analysis_request, clean_instruction, seed],
            outputs=[overlay, clean_image, clean_json, status, bundle],
        )

        gr.Markdown(
            "## Notes\n"
            "- Nano Banana 2 correspond ici au modèle `gemini-3.1-flash-image`.\n"
            "- Le clean plate peut fallback sur ComfyUI si l'API Gemini n'est pas dispo.\n"
            "- On peut brancher la suite ensuite: recadrages, localisations, édition guidée par masque, etc."
        )

    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    main()