"""
MaramaRoute live demo (v3) — fronts the LumynaX gateway when configured;
otherwise runs router logic locally against the live registry.

Env:
  LUMYNAX_GATEWAY_URL   e.g. https://gateway.lumynax.com (optional)
  LUMYNAX_GATEWAY_KEY   Bearer key for the gateway (optional)
"""
import json
import os
from dataclasses import dataclass, field
from typing import Any, Optional

import gradio as gr
import httpx
from huggingface_hub import hf_hub_download

REGISTRY_REPO = "AbteeXAILab/marama-route"
REGISTRY_PATH = "configs/lumynax_model_registry.json"
GATEWAY_URL   = os.environ.get("LUMYNAX_GATEWAY_URL", "").rstrip("/")
GATEWAY_KEY   = os.environ.get("LUMYNAX_GATEWAY_KEY", "lumynax-local-dev")

BRAND_CSS = """
:root { --lx-paper:#fffefa; --lx-ink:#0a0a0b; --lx-amber:#e08a2c; --lx-amber-dark:#9a5416; --lx-muted:#726b62; --lx-soft:#f6f0e8; }
body, .gradio-container { background: var(--lx-paper) !important; color: var(--lx-ink) !important; }
h1, h2, h3 { font-family: 'Cormorant Garamond','EB Garamond',Georgia,serif; }
.lx-hero { background: linear-gradient(135deg,#fffefa 0%,#f6f0e8 100%); border: 1px solid rgba(10,10,11,0.08); border-radius: 12px; padding: 18px 22px; margin-bottom: 14px; }
.lx-pick { background: #f0f7ee; border-left: 4px solid #4caf50; padding: 12px 14px; border-radius: 6px; }
.lx-runner-up { background: var(--lx-soft); border-left: 3px solid var(--lx-amber); padding: 10px 12px; border-radius: 6px; margin-top: 8px; }
table.lx-runners { border-collapse: collapse; width: 100%; font-size: 0.9em; }
table.lx-runners th, table.lx-runners td { border-bottom: 1px solid rgba(10,10,11,0.08); padding: 6px 8px; text-align: left; }
.lx-banner-nz { background: linear-gradient(90deg, #e08a2c22, #0a0a0b11); padding: 6px 12px; border-radius: 6px; font-size: 0.85em; margin-bottom: 8px; }
"""

HERO_HTML = """
<div class="lx-banner-nz">🇳🇿 <b>Made in Aotearoa New Zealand</b> · AbteeX AI Labs · <a href="https://abteex.com" target="_blank">abteex.com</a> · <a href="https://lumynax.com" target="_blank">lumynax.com</a></div>
<div class="lx-hero">
  <h1 style="margin:0 0 6px 0;">🧭 MaramaRoute — sovereign router</h1>
  <p style="margin:0; color:var(--lx-muted);"><em>Ko te marama te tuapapa.</em> Pick the right model from the live LumynaX family of 98 models — gated on residency, modality, tools, and sovereignty tier.</p>
</div>
"""


@dataclass
class RouteDecision:
    pick: Optional[dict] = None
    score: float = 0.0
    runners_up: list = field(default_factory=list)
    reason: str = ""


def _load_registry() -> dict:
    try:
        p = hf_hub_download(repo_id=REGISTRY_REPO, filename=REGISTRY_PATH,
                            repo_type="model", token=os.environ.get("HF_TOKEN"))
        return json.loads(open(p, encoding="utf-8").read())
    except Exception as e:
        return {"models": [], "_error": str(e)}


REGISTRY = _load_registry()


def _route_via_gateway(modalities, requires_local, requires_tools, requires_json, jurisdiction):
    """Hit the gateway's /v1/route endpoint."""
    try:
        with httpx.Client(timeout=30) as c:
            r = c.get(f"{GATEWAY_URL}/v1/route", params={
                "modalities": ",".join(modalities), "requires_local": requires_local,
                "requires_tools": requires_tools, "requires_json": requires_json,
                "jurisdiction": jurisdiction,
            }, headers={"Authorization": f"Bearer {GATEWAY_KEY}"})
            r.raise_for_status()
            return r.json()
    except Exception as e:
        return {"_error": f"gateway unreachable: {e}"}


def route_local(modalities, requires_local, requires_tools, requires_json,
                jurisdiction, max_params_b, task_hint) -> RouteDecision:
    d = RouteDecision()
    candidates = []
    for m in REGISTRY.get("models", []):
        if any(mod not in (m.get("modalities") or []) for mod in modalities): continue
        if requires_local and (m.get("sovereignty_tier") or 5) < 3: continue
        if jurisdiction and jurisdiction not in (m.get("residency") or []): continue
        if requires_tools and not m.get("supports_tools"): continue
        if requires_json and not m.get("supports_json"): continue
        if max_params_b > 0 and (m.get("total_params_b") or 0) > max_params_b: continue
        q = int(m.get("quality_rank") or 5); s = int(m.get("sovereignty_tier") or 3); c = int(m.get("cost_rank") or 5)
        score = (6 - q) * 2 + s * 1.5 + (6 - c) * 0.5
        if task_hint and (task_hint.lower() in " ".join(m.get("tags") or []).lower() + " " + m["model_id"].lower()):
            score += 3
        if max_params_b > 0 and (m.get("total_params_b") or 0) <= max_params_b * 0.5: score += 1
        candidates.append((score, m))
    candidates.sort(key=lambda x: -x[0])
    if candidates:
        d.score, d.pick = candidates[0]
        d.runners_up = candidates[1:5]
        d.reason = f"top-of-{len(candidates)} candidates"
    return d


def render(prompt, modalities, requires_local, requires_tools, requires_json,
           jurisdiction, max_params_b, task_hint):
    backend_note = ""
    if GATEWAY_URL:
        gw = _route_via_gateway(modalities, requires_local, requires_tools, requires_json, jurisdiction)
        if "_error" not in gw:
            backend_note = f"<div class='lx-runner-up' style='border-left-color:#4caf50;'>Routed via gateway <code>{GATEWAY_URL}</code></div>"
            pick = next((m for m in REGISTRY.get("models", []) if m["repo_id"].endswith(gw["model"])), None)
            if pick:
                d = RouteDecision(pick=pick, score=gw.get("score", 0.0),
                                   reason=f"gateway top-of-{len(gw.get('alternatives', [])) + 1}")
                # enrich with full info
                return backend_note + _format(d)
    d = route_local(modalities, requires_local, requires_tools, requires_json,
                    jurisdiction, max_params_b, task_hint)
    if not d.pick:
        return f'<div class="lx-pick" style="background:#fde8e8;border-left-color:#c1351a;"><b>No candidate matches your filters.</b></div>'
    return _format(d)


def _format(d: RouteDecision) -> str:
    p = d.pick
    runners_html = "<table class='lx-runners'><tr><th>Runner-up</th><th>Score</th><th>Params</th></tr>"
    for s, m in d.runners_up:
        tp = m.get("total_params_b"); ap = m.get("active_params_b")
        runners_html += f"<tr><td><code>{m['repo_id'].split('/')[-1]}</code></td><td>{s:.2f}</td><td>{tp}B{f'/{ap}Ba' if ap else ''}</td></tr>"
    runners_html += "</table>"
    tp = p.get("total_params_b"); ap = p.get("active_params_b")
    return f"""
<div class='lx-pick'>
  <div style='display:flex; justify-content:space-between; align-items:baseline;'>
    <h2 style='margin:0;'>🎯 {p['title']}</h2>
    <span style='font-family:monospace;'>score {d.score:.2f}</span>
  </div>
  <p><a href='https://huggingface.co/{p['repo_id']}' target='_blank'><code>{p['repo_id']}</code></a></p>
  <div style='display:grid; grid-template-columns:1fr 1fr 1fr; gap:8px; font-size:0.95em;'>
    <div><b>Params:</b> {tp}B{f' / {ap}B active' if ap else ''}</div>
    <div><b>Context:</b> {p.get('context_tokens','—')}</div>
    <div><b>Runtime:</b> {p.get('runtime','—')}</div>
    <div><b>Modalities:</b> {', '.join(p.get('modalities') or [])}</div>
    <div><b>Sovereignty:</b> tier {p.get('sovereignty_tier')}</div>
    <div><b>License:</b> {p.get('license_id','—')}</div>
  </div>
  <p style='margin:10px 0 0 0; font-size:0.9em; color:var(--lx-muted);'>{d.reason}. Run with <code>lumynax run {p['repo_id'].split('/')[-1]}</code></p>
</div>
{('<div class="lx-runner-up"><b>Runner-ups:</b>' + runners_html + '</div>') if d.runners_up else ''}
"""


with gr.Blocks(css=BRAND_CSS, title="LumynaX · MaramaRoute") as demo:
    gr.HTML(HERO_HTML)
    note = f"Loaded {len(REGISTRY.get('models', []))} models" + (f" · gateway at {GATEWAY_URL}" if GATEWAY_URL else " · local routing")
    gr.Markdown(f"_{note}_")
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Prompt or task", lines=2)
            with gr.Row():
                modalities = gr.CheckboxGroup(["text","vision","audio"], value=["text"], label="Modalities")
                jurisdiction = gr.Dropdown(["NZ","AU","global",""], value="NZ", label="Jurisdiction")
            with gr.Row():
                requires_local = gr.Checkbox(False, label="Sovereignty tier ≥ 3")
                requires_tools = gr.Checkbox(False, label="Tool-calling")
                requires_json = gr.Checkbox(False, label="JSON-mode")
            with gr.Row():
                max_params_b = gr.Slider(0, 700, value=0, step=1, label="Max params (B), 0=unlimited")
                task_hint = gr.Textbox(label="Task hint", placeholder="coder, reasoning, vision, translate…")
            go = gr.Button("Route", variant="primary")
        with gr.Column(scale=3):
            out = gr.HTML()
    gr.Examples(
        examples=[
            ["Fix this Python bug", ["text"], True, False, False, "NZ", 100, "coder"],
            ["Describe this image", ["text","vision"], False, False, False, "NZ", 0, "vision"],
            ["Mathematical proof", ["text"], False, False, False, "NZ", 0, "reasoning"],
            ["Translate to te reo Maori", ["text"], True, False, False, "NZ", 5, "translate"],
            ["1M token document summary", ["text"], False, False, False, "NZ", 50, "long-context"],
        ],
        inputs=[prompt, modalities, requires_local, requires_tools, requires_json, jurisdiction, max_params_b, task_hint],
    )
    go.click(render, [prompt, modalities, requires_local, requires_tools, requires_json, jurisdiction, max_params_b, task_hint], out)
    gr.Markdown("---\n*Made in Aotearoa New Zealand · [abteex.com](https://abteex.com) · [lumynax.com](https://lumynax.com) · [GitHub](https://github.com/Aimaghsoodi/lumynax-release) · Ko te marama te tuapapa.*")


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