from __future__ import annotations

import base64
import io
import json
import os
import tempfile
import zipfile
from pathlib import Path

import gradio as gr
import httpx
from PIL import Image

try:
    import spaces
except ImportError:  # Local CPU development; Hugging Face provides this module in the Space.
    class _LocalSpaces:
        @staticmethod
        def GPU(*_args, **_kwargs):
            return lambda function: function

    spaces = _LocalSpaces()

API_URL = os.getenv("MODAL_API_URL", "http://127.0.0.1:8000").rstrip("/")
API_TOKEN = os.getenv("OPENCOMIC_API_TOKEN", "")


@spaces.GPU(duration=60)
def zerogpu_compatibility_probe() -> str:
    """Satisfy the existing ZeroGPU Space hardware contract; normal UI calls use Modal."""
    return "OpenComic uses its authenticated Modal renderer; no Hugging Face GPU is required."


def _headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {API_TOKEN}"} if API_TOKEN else {}


def _upload_payload(files: list[str] | None) -> tuple[bytes, str, str]:
    paths = [Path(item) for item in files or []]
    if not paths:
        raise gr.Error("Upload at least one image, PDF, CBZ, or ZIP file.")
    if len(paths) == 1:
        return paths[0].read_bytes(), paths[0].name, "application/octet-stream"
    stream = io.BytesIO()
    with zipfile.ZipFile(stream, "w", zipfile.ZIP_DEFLATED) as archive:
        for index, path in enumerate(paths):
            archive.writestr(f"{index:04d}{path.suffix.lower()}", path.read_bytes())
    return stream.getvalue(), "uploaded-pages.cbz", "application/vnd.comicbook+zip"


def analyze(files: list[str] | None, reading_direction: str, continuity_notes: str):
    payload, name, mime = _upload_payload(files)
    with httpx.Client(timeout=900) as client:
        response = client.post(
            f"{API_URL}/analyze-comic",
            params={
                "reading_direction": reading_direction,
                "continuity_notes": continuity_notes,
            },
            headers=_headers(),
            files={"file": (name, payload, mime)},
        )
    if response.is_error:
        raise gr.Error(f"Analysis failed ({response.status_code}): {response.text[:500]}")
    result = response.json()
    summary = (
        f"Analyzed {result['pages']} pages and {result['panels']} panels. "
        "Qwen semantic analysis built a concrete cast, event chain, setting, and unresolved thread."
    )
    session = {
        "memory": result["memory"],
        "context_images": result.get("context_images", []),
        "reference_images": result.get("reference_images", []),
        "semantic_analysis": result.get("semantic_analysis", {}),
        "layout_analysis": result.get("layout_analysis", {}),
        "reading_direction": reading_direction,
    }
    return session, result["memory"], summary


def _data_url_to_image(value: str) -> Image.Image:
    payload = base64.b64decode(value.split(",", 1)[1])
    return Image.open(io.BytesIO(payload)).convert("RGB")


def _continuation_request(
    session: dict,
    pages: int,
    creativity: float,
    dialogue_density: float,
    style_fidelity: float,
    character_fidelity: float,
    reading_direction: str,
    research_mode: bool,
) -> dict:
    resolved_direction = (
        session.get("reading_direction", "ltr")
        if reading_direction == "auto"
        else reading_direction
    )
    return {
        "memory": session.get("memory", session),
        "context_images": session.get("context_images", []),
        "reference_images": session.get("reference_images", []),
        "settings": {
            "pages": int(pages),
            "creativity": creativity,
            "dialogue_density": dialogue_density,
            "style_fidelity": style_fidelity,
            "character_fidelity": character_fidelity,
            "reading_direction": resolved_direction,
            "research_mode": research_mode,
            "seed": 20260823,
        },
    }


def generate(
    session: dict | None,
    pages: int,
    creativity: float,
    dialogue_density: float,
    style_fidelity: float,
    character_fidelity: float,
    reading_direction: str,
    research_mode: bool,
):
    if not session:
        raise gr.Error("Analyze a comic before requesting a continuation.")
    request = _continuation_request(
        session,
        pages,
        creativity,
        dialogue_density,
        style_fidelity,
        character_fidelity,
        reading_direction,
        research_mode,
    )
    reference_images = request["reference_images"]
    context_images = request["context_images"]
    with httpx.Client(timeout=3600) as client:
        planned_response = client.post(
            f"{API_URL}/plan-continuation", headers=_headers(), json=request
        )
        if planned_response.is_error:
            raise gr.Error(
                f"Planning failed ({planned_response.status_code}): {planned_response.text[:500]}"
            )
        planned = planned_response.json()
        render_request = {
            **request,
            "scripts": planned["scripts"],
            "planner_metadata": planned.get("planner", {}),
        }
        response = client.post(f"{API_URL}/render-script", headers=_headers(), json=render_request)
        if response.is_error:
            raise gr.Error(f"Rendering failed ({response.status_code}): {response.text[:500]}")
        result = response.json()
    images = [_data_url_to_image(item) for item in result.get("page_data_urls", [])]
    research = {
        "model_variant": [item.get("model_variant") for item in result.get("scripts", [])],
        "routing": planned.get("routing", result.get("routing", [])),
        "job_id": result.get("job_id"),
        "renderer": result.get(
            "renderer", {"status": "RENDERER NOT RUN (visible placeholder panels)"}
        ),
        "planner": result.get("planner", {}),
    }
    updated_session = {
        "memory": result.get("memory"),
        "context_images": context_images,
        "reference_images": result.get("next_reference_images", reference_images),
        "semantic_analysis": session.get("semantic_analysis", {}),
        "layout_analysis": session.get("layout_analysis", {}),
    }
    return (
        images,
        result.get("scripts", []),
        updated_session,
        result.get("memory"),
        research,
    )


def draft_story_stage(
    session: dict | None,
    pages: int,
    creativity: float,
    dialogue_density: float,
    style_fidelity: float,
    character_fidelity: float,
    reading_direction: str,
    research_mode: bool,
):
    if not session:
        raise gr.Error("Analyze a comic before drafting its continuation.")
    request = _continuation_request(
        session,
        pages,
        creativity,
        dialogue_density,
        style_fidelity,
        character_fidelity,
        reading_direction,
        research_mode,
    )
    with httpx.Client(timeout=1800) as client:
        response = client.post(f"{API_URL}/draft-story", headers=_headers(), json=request)
    if response.is_error:
        raise gr.Error(f"Story drafting failed ({response.status_code}): {response.text[:500]}")
    drafted = response.json()
    state = {"request": request, "story": drafted, "scripts": None}
    return (
        state,
        json.dumps(drafted["story_brief"], indent=2, ensure_ascii=False),
        "Story drafted and editorially validated. Edit it if needed, then lock it into a storyboard.",
        {"story_stage": drafted},
    )


def storyboard_stage(planning_state: dict | None, approved_story: dict | str | None):
    if not planning_state or not approved_story:
        raise gr.Error("Draft and approve a story before storyboarding.")
    if isinstance(approved_story, str):
        try:
            approved_story = json.loads(approved_story)
        except json.JSONDecodeError as exc:
            raise gr.Error(f"The edited story is not valid JSON: {exc}") from exc
    request = {**planning_state["request"], "locked_story_brief": approved_story}
    with httpx.Client(timeout=1800) as client:
        response = client.post(
            f"{API_URL}/storyboard-continuation", headers=_headers(), json=request
        )
    if response.is_error:
        raise gr.Error(f"Storyboarding failed ({response.status_code}): {response.text[:500]}")
    planned = response.json()
    state = {**planning_state, "request": request, "storyboard": planned, "scripts": planned["scripts"]}
    return (
        state,
        planned.get("planner", {}).get("causal_plan", {}),
        planned["scripts"],
        "Storyboard locked. Review the page and panel beats, then start visual rendering.",
        {"story_stage": planning_state.get("story", {}), "storyboard_stage": planned},
    )


def render_locked_stage(session: dict | None, planning_state: dict | None):
    if not session or not planning_state or not planning_state.get("scripts"):
        raise gr.Error("Lock a storyboard before rendering.")
    request = planning_state["request"]
    render_request = {
        **request,
        "scripts": planning_state["scripts"],
        "planner_metadata": planning_state.get("storyboard", {}).get("planner", {}),
    }
    with httpx.Client(timeout=3600) as client:
        response = client.post(f"{API_URL}/render-script", headers=_headers(), json=render_request)
    if response.is_error:
        raise gr.Error(f"Rendering failed ({response.status_code}): {response.text[:500]}")
    result = response.json()
    images = [_data_url_to_image(item) for item in result.get("page_data_urls", [])]
    research = {
        "story_stage": planning_state.get("story", {}),
        "storyboard_stage": planning_state.get("storyboard", {}),
        "renderer": result.get("renderer", {}),
        "job_id": result.get("job_id"),
    }
    updated_session = {
        "memory": result.get("memory"),
        "context_images": request.get("context_images", []),
        "reference_images": result.get(
            "next_reference_images", request.get("reference_images", [])
        ),
        "semantic_analysis": session.get("semantic_analysis", {}),
        "layout_analysis": session.get("layout_analysis", {}),
        "reading_direction": request["settings"]["reading_direction"],
    }
    return images, result.get("scripts", []), updated_session, result.get("memory"), research


def record_preference(choice: str, notes: str) -> str:
    if not choice:
        return "Choose a preference before submitting."
    row = {"preference": choice, "notes": notes[:1000]}
    destination = Path(tempfile.gettempdir()) / "opencomic_pairwise.jsonl"
    with destination.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    return "Anonymous preference recorded for this research session."


with gr.Blocks(title="OpenComic-Continue") as demo:
    gr.Markdown(
        "# OpenComic-Continue\n"
        "A multi-page story-first research prototype: understand several context pages, lock a "
        "complete continuation arc, storyboard each page into natural four- or five-panel pacing, "
        "then render every box from the previous "
        "image plus immutable cast/prop references. "
        "Only upload material you own or have permission to transform."
    )
    memory_state = gr.State()
    planning_state = gr.State()
    with gr.Tab("Analyze"):
        uploads = gr.File(
            label="Comic pages, PDF, CBZ, or ZIP",
            file_count="multiple",
            type="filepath",
        )
        direction = gr.Radio(["ltr", "rtl"], value="ltr", label="Reading direction")
        continuity_notes = gr.Textbox(
            label="Optional continuity corrections",
            placeholder="Example: four physical cats; the round bottle is pink, not red",
            lines=2,
            max_lines=4,
        )
        analyze_button = gr.Button("Analyze comic", variant="primary")
        analysis_summary = gr.Markdown()
        memory_json = gr.JSON(label="StoryMemory")
        analyze_button.click(
            analyze,
            [uploads, direction, continuity_notes],
            [memory_state, memory_json, analysis_summary],
            api_name="analyze_quality",
        )
    with gr.Tab("Continue"):
        with gr.Row():
            page_count = gr.Slider(2, 4, value=2, step=1, label="Continuation pages")
            creativity = gr.Slider(0, 1, value=0.5, label="Creativity")
            dialogue = gr.Slider(
                0,
                1,
                value=0.0,
                label="Dialogue density (0 recommended unless source dialogue was extracted)",
            )
        with gr.Row():
            style = gr.Slider(0, 1, value=0.8, label="Style fidelity")
            character = gr.Slider(0, 1, value=0.9, label="Character fidelity")
            generation_direction = gr.Dropdown(
                ["auto", "ltr", "rtl"], value="auto", label="Reading direction"
            )
        research_mode = gr.Checkbox(label="Research mode (show routing and model metadata)")
        gr.Markdown(
            "1. Draft the causal prose story. 2. Edit/approve it and lock a storyboard. "
            "3. Render the locked panels sequentially with strict anatomy, cast, and prop checks."
        )
        with gr.Row():
            draft_button = gr.Button("1 · Draft story", variant="primary")
            storyboard_button = gr.Button("2 · Lock storyboard")
            render_button = gr.Button("3 · Render panels")
        stage_status = gr.Markdown()
        story_brief = gr.Code(label="Editable story brief", language="json", lines=24)
        storyboard_json = gr.JSON(label="Locked page/panel storyboard")
        gallery = gr.Gallery(label="Continuation pages", columns=2, object_fit="contain")
        scripts = gr.JSON(label="Structured scripts")
        updated_memory = gr.JSON(label="Updated StoryMemory")
        research_json = gr.JSON(label="Research metadata")
        draft_button.click(
            draft_story_stage,
            [
                memory_state,
                page_count,
                creativity,
                dialogue,
                style,
                character,
                generation_direction,
                research_mode,
            ],
            [planning_state, story_brief, stage_status, research_json],
            api_name="draft_story_quality",
        )
        storyboard_button.click(
            storyboard_stage,
            [planning_state, story_brief],
            [planning_state, storyboard_json, scripts, stage_status, research_json],
            api_name="storyboard_quality",
        )
        render_button.click(
            render_locked_stage,
            [memory_state, planning_state],
            [gallery, scripts, memory_state, updated_memory, research_json],
            api_name="render_quality",
        )
        # Keep the original combined API contract for scripted research clients.
        legacy_generate_button = gr.Button("Legacy combined generation", visible=False)
        legacy_generate_button.click(
            generate,
            [
                memory_state,
                page_count,
                creativity,
                dialogue,
                style,
                character,
                generation_direction,
                research_mode,
            ],
            [gallery, scripts, memory_state, updated_memory, research_json],
            api_name="generate_quality",
        )
    with gr.Tab("Pairwise evaluation"):
        gr.Markdown(
            "Use this form after comparing two model outputs supplied by a study administrator. "
            "No identity is collected."
        )
        preference = gr.Radio(["A", "B", "Tie", "Both invalid"], label="Preferred output")
        preference_notes = gr.Textbox(label="Optional reason", lines=3)
        preference_button = gr.Button("Submit preference")
        preference_status = gr.Markdown()
        preference_button.click(
            record_preference, [preference, preference_notes], preference_status
        )
    gr.Markdown(
        "Generated continuations may be inaccurate or legally restricted. The current demo "
        "labels unexecuted image rendering rather than presenting placeholders as model results."
    )


if __name__ == "__main__":
    demo.launch(show_error=True)