from __future__ import annotations

import html
import json
import sys
import traceback
from pathlib import Path

import gradio as gr

try:
    import spaces
except Exception:  # local dev outside HF Spaces
    class _SpacesShim:
        @staticmethod
        def GPU(*args, **kwargs):
            def deco(fn):
                return fn
            return deco
    spaces = _SpacesShim()

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from compose.merge import compose
from compose.provenance import provenance_sentence
from registry import lora_choices

SAGE = "#5f6f52"
CREAM = "#f6f1e8"
TERRACOTTA = "#b7653c"


def render_node(node: dict, depth: int = 0) -> str:
    margin = depth * 22
    children = "".join(render_node(child, depth + 1) for child in node.get("children", []))
    status = node.get("status", "unknown")
    checkpoint = node.get("checkpoint_step") or "not selected yet"
    cycle = "<div class='warn'>Cycle detected — recursion stopped.</div>" if node.get("cycle_detected") else ""
    return f"""
    <div class="lineage-card" style="margin-left:{margin}px">
      <div class="card-top">
        <strong>{html.escape(node.get('cultural_source') or node['lora_id'])}</strong>
        <span class="pill">{node['weight_pct']:g}%</span>
      </div>
      <div class="muted">{html.escape(node['lora_id'])} · {html.escape(str(node.get('type')))} · {html.escape(str(status))}</div>
      <div class="muted">checkpoint: {html.escape(str(checkpoint))}</div>
      <a href="https://huggingface.co/{html.escape(node.get('hf_repo') or '')}" target="_blank">{html.escape(node.get('hf_repo') or '')}</a>
      {cycle}
    </div>
    {children}
    """


def render_lineage(provenance: dict) -> str:
    ancestry = provenance.get("ancestry", [])
    bars = []
    colors = [SAGE, TERRACOTTA, "#3f5f8f", "#9a7b33", "#72517e"]
    for i, node in enumerate(ancestry):
        bars.append(
            f"<div title='{html.escape(node['lora_id'])}' style='width:{node['weight_pct']}%;background:{colors[i % len(colors)]}'>{node['weight_pct']:g}%</div>"
        )
    cards = "".join(render_node(node) for node in ancestry)
    raw = html.escape(json.dumps(provenance, indent=2))
    sentence = html.escape(provenance_sentence(provenance))
    full_prompt = html.escape(provenance.get("full_prompt", ""))
    model = html.escape(provenance.get("inference_model", provenance.get("base_model", "")))
    return f"""
    <style>
      .vl-wrap {{ background:{CREAM}; border:1px solid #d9cfbd; border-radius:18px; padding:18px; color:#263322; }}
      .stack {{ display:flex; height:38px; overflow:hidden; border-radius:999px; box-shadow: inset 0 0 0 1px rgba(0,0,0,.12); margin:12px 0 18px; }}
      .stack div {{ color:white; font-weight:700; display:flex; align-items:center; justify-content:center; min-width:48px; }}
      .lineage-card {{ background:white; border:1px solid #ddd2c0; border-radius:14px; padding:12px; margin-top:10px; }}
      .card-top {{ display:flex; justify-content:space-between; gap:12px; }}
      .pill {{ background:{SAGE}; color:white; padding:3px 10px; border-radius:999px; font-size:12px; }}
      .muted {{ color:#66705f; font-size:13px; margin-top:3px; }}
      .warn {{ color:#9b341f; font-weight:700; margin-top:6px; }}
      details {{ margin-top:16px; }}
      pre {{ white-space:pre-wrap; font-size:12px; background:#fff; padding:12px; border-radius:10px; }}
    </style>
    <div class="vl-wrap">
      <h3>Visual lineage</h3>
      <p>{sentence}</p>
      <div class="muted">Model: {model}</div>
      <div class="muted">Full prompt: {full_prompt}</div>
      <div class="stack">{''.join(bars)}</div>
      {cards}
      <details><summary>Raw provenance JSON</summary><pre>{raw}</pre></details>
    </div>
    """


@spaces.GPU(duration=180)
def generate(lora_a: str, lora_b: str, blend_a: int, prompt: str, seed: int, size: int):
    if not prompt.strip():
        raise gr.Error("Give the image a prompt first.")
    weight_a = blend_a / 100
    weight_b = 1 - weight_a
    try:
        result = compose(
            lora_ids=[lora_a, lora_b],
            weights=[weight_a, weight_b],
            prompt=prompt,
            seed=int(seed),
            registry_path=str(ROOT / "registry/loras.json"),
            output_dir=str(ROOT / "outputs"),
            width=int(size),
            height=int(size),
        )
        return result["image"], render_lineage(result["provenance"])
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Generation failed: {type(e).__name__}: {e}")


def build_demo():
    choices = lora_choices(ROOT / "registry/loras.json")
    default_a = "eritrean_krar_v1" if any(v == "eritrean_krar_v1" for _, v in choices) else choices[0][1]
    default_b = default_a
    with gr.Blocks(title="Visual Lineage", theme=gr.themes.Soft(primary_hue="green")) as demo:
        gr.Markdown(
            "# Visual Lineage\n"
            "*A 23andMe for imagined instruments.*\n\n"
            "Live FLUX.2 [klein] inference with published instrument LoRAs. "
            "Select any two from Eritrean krar, Korean gayageum, or Brazilian berimbau "
            "and trace the visual lineage of every generation."
        )
        with gr.Row():
            with gr.Column(scale=1):
                lora_a = gr.Dropdown(label="Instrument/source A", choices=choices, value=default_a)
                lora_b = gr.Dropdown(label="Instrument/source B", choices=choices, value=default_b)
                blend = gr.Slider(0, 100, value=100, step=1, label="Blend: A % / B %")
                prompt = gr.Textbox(label="Describe the imagined instrument image", value="a musician holding a newly invented bowl-shaped string instrument in a small warm room", lines=3)
                seed = gr.Number(label="Seed", value=42, precision=0)
                size = gr.Slider(512, 1024, value=768, step=256, label="Image size")
                btn = gr.Button("Generate with live LoRA", variant="primary")
            with gr.Column(scale=2):
                output_image = gr.Image(label="Generated image", type="pil")
                lineage_panel = gr.HTML(label="Visual lineage")
        btn.click(generate, [lora_a, lora_b, blend, prompt, seed, size], [output_image, lineage_panel], show_progress="full")
    return demo


if __name__ == "__main__":
    build_demo().launch()