"""
SovereignCode live demo (v3) — policy gates + MaramaRoute over the live registry,
hitting the LumynaX gateway when configured.
"""
from __future__ import annotations
import hashlib, json, os
from dataclasses import dataclass, field
from datetime import datetime, timezone
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-result-allow { background: #f0f7ee; border-left: 4px solid #4caf50; padding: 14px 16px; border-radius: 6px; }
.lx-result-deny  { background: #fde8e8; border-left: 4px solid #c1351a; padding: 14px 16px; border-radius: 6px; }
.lx-routing { background: var(--lx-soft); border-left: 3px solid var(--lx-amber); padding: 12px 14px; border-radius: 6px; margin-top: 10px; }
.lx-audit { font-family: ui-monospace, SFMono-Regular, monospace; font-size: 0.85em; background: var(--lx-soft); padding: 10px; border-radius: 4px; overflow-x: auto; }
.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;">🛡️ SovereignCode — Data Capsule + MaramaRoute</h1>
  <p style="margin:0; color:var(--lx-muted);">Define a capsule policy, run a request through 6 gates, and if allowed route to the best model in the live 98-model LumynaX family.</p>
</div>
"""


@dataclass
class Decision:
    allow: bool
    gate: str
    reason: str
    obligations: list = field(default_factory=list)
    route_pick: Optional[dict] = None
    route_score: float = 0.0
    route_reason: str = ""


GATE_EXPLAIN = {
    "purpose": "Request's purpose isn't in the capsule's allowed_purposes whitelist.",
    "residency": "Capsule's data_residency tag doesn't include the request's jurisdiction.",
    "remote-model": "Capsule forbids any remote frontier model for this data class.",
    "training": "Capsule forbids using this data for model training/fine-tuning.",
    "export": "Capsule's allowed_exports doesn't include the requested egress destination.",
    "approval": "Request needs a named human approver per capsule policy.",
    "allow": "All policy gates passed. Request will be routed via MaramaRoute.",
}


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 evaluate(capsule, request) -> Decision:
    obs = []
    if request.get("purpose") not in capsule.get("allowed_purposes", []):
        return Decision(False, "purpose", f"purpose '{request.get('purpose')}' not allowed")
    if request.get("jurisdiction") not in capsule.get("data_residency", []):
        return Decision(False, "residency", f"jurisdiction '{request.get('jurisdiction')}' not in {capsule.get('data_residency')}")
    if request.get("remote_model_required") and not capsule.get("allow_remote_models"):
        return Decision(False, "remote-model", "capsule forbids remote frontier models")
    if request.get("for_training") and not capsule.get("allow_training"):
        return Decision(False, "training", "capsule forbids using this data for training")
    if request.get("export_to") and request.get("export_to") not in capsule.get("allowed_exports", []):
        return Decision(False, "export", f"export to '{request.get('export_to')}' not allowed")
    if capsule.get("require_human_approval") and not request.get("approver"):
        return Decision(False, "approval", "human approver required")
    if capsule.get("require_human_approval"):
        obs.append(f"audit notes approver={request.get('approver')}")
    obs.append("hash-chain audit append")
    pick, score, rr = _route(request)
    return Decision(True, "allow", "all 6 gates passed; routed via MaramaRoute", obs,
                    route_pick=pick, route_score=score, route_reason=rr)


def _route(request) -> tuple[Optional[dict], float, str]:
    if GATEWAY_URL:
        try:
            with httpx.Client(timeout=30) as c:
                r = c.get(f"{GATEWAY_URL}/v1/route", params={
                    "modalities": ",".join(request.get("modalities", ["text"])),
                    "requires_local": request.get("requires_local", False),
                    "requires_tools": request.get("requires_tools", False),
                    "jurisdiction": request.get("jurisdiction", "NZ"),
                }, headers={"Authorization": f"Bearer {GATEWAY_KEY}"})
                r.raise_for_status(); data = r.json()
                pick = next((m for m in REGISTRY.get("models", []) if m["repo_id"].endswith(data["model"])), None)
                if pick:
                    return pick, data.get("score", 0.0), f"gateway top-of-{1 + len(data.get('alternatives', []))}"
        except Exception:
            pass  # fall back to local
    cands = []
    mods = request.get("modalities", ["text"])
    requires_local = request.get("requires_local", False)
    jur = request.get("jurisdiction", "NZ")
    hint = request.get("task_hint", "")
    for m in REGISTRY.get("models", []):
        if any(mod not in (m.get("modalities") or []) for mod in mods): continue
        if requires_local and (m.get("sovereignty_tier") or 5) < 3: continue
        if jur not in (m.get("residency") or []): 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 hint and (hint.lower() in " ".join(m.get("tags") or []).lower() + " " + m["model_id"].lower()):
            score += 3
        cands.append((score, m))
    cands.sort(key=lambda x: -x[0])
    if not cands: return None, 0.0, "no candidate"
    score, pick = cands[0]
    return pick, score, f"local top-of-{len(cands)}"


def _hash(req): return hashlib.sha256(json.dumps(req, sort_keys=True, separators=(",", ":")).encode()).hexdigest()


def render(capsule_text, purpose, jurisdiction, remote_model, for_training,
           export_to, approver, modalities, task_hint, requires_local):
    try:
        capsule = json.loads(capsule_text)
    except Exception as e:
        return f'<div class="lx-result-deny"><b>Invalid capsule JSON:</b> {e}</div>'
    req = {"purpose": purpose, "jurisdiction": jurisdiction,
           "remote_model_required": remote_model, "for_training": for_training,
           "export_to": export_to or None, "approver": approver or None,
           "modalities": modalities, "task_hint": task_hint, "requires_local": requires_local,
           "ts": datetime.now(timezone.utc).isoformat()}
    d = evaluate(capsule, req)
    audit = {"request_sha256": _hash(req), "decision": "ALLOW" if d.allow else "DENY", "gate": d.gate, "obligations": d.obligations}
    if not d.allow:
        return f'''
<div class="lx-result-deny">
  <h2 style="margin:0 0 6px 0;">🚫 DENIED at gate: <code>{d.gate}</code></h2>
  <p>{d.reason}</p>
  <p style="font-size:0.85em; color:var(--lx-muted);"><em>{GATE_EXPLAIN.get(d.gate, '')}</em></p>
  <div class="lx-audit">{json.dumps(audit, indent=2)}</div>
</div>'''
    p = d.route_pick
    if not p:
        return f'<div class="lx-result-deny"><h2>⚠️ Allowed but no model survived routing</h2><p>{d.route_reason}</p></div>'
    tp = p.get("total_params_b"); ap = p.get("active_params_b")
    return f'''
<div class="lx-result-allow">
  <h2 style="margin:0 0 6px 0;">✅ ALLOWED — passed all 6 gates</h2>
  <p>{d.reason}</p>
  <ul>{"".join(f"<li>{o}</li>" for o in d.obligations)}</ul>
</div>
<div class="lx-routing">
  <b>MaramaRoute pick:</b>
  <h3>🎯 {p["title"]}</h3>
  <p><a href="https://huggingface.co/{p['repo_id']}" target="_blank"><code>{p['repo_id']}</code></a> · score <code>{d.route_score:.2f}</code></p>
  <div style="display:grid; grid-template-columns:1fr 1fr 1fr; gap:6px; font-size:0.9em;">
    <div><b>Params:</b> {tp}B{f' / {ap}Ba' if ap else ''}</div>
    <div><b>Ctx:</b> {p.get('context_tokens','—')}</div>
    <div><b>Runtime:</b> {p.get('runtime','—')}</div>
    <div><b>Sovereignty:</b> tier {p.get('sovereignty_tier')}</div>
    <div><b>Modalities:</b> {", ".join(p.get('modalities') or [])}</div>
    <div><b>License:</b> {p.get('license_id','—')}</div>
  </div>
  <p style="font-size:0.85em; color:var(--lx-muted);">{d.route_reason}</p>
</div>
<details><summary>📓 Audit ledger entry</summary><div class="lx-audit">{json.dumps(audit, indent=2)}</div></details>
'''


DEFAULT_CAPSULE = json.dumps({
    "capsule_id": "nz-personal-sovereignty-v1",
    "data_residency": ["NZ"],
    "allow_remote_models": False,
    "allow_training": False,
    "allowed_purposes": ["code-edit","summarize","translate","embed","describe-image"],
    "allowed_exports": ["local-ide"],
    "require_human_approval": False,
}, indent=2)


with gr.Blocks(css=BRAND_CSS, title="LumynaX · SovereignCode + 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():
            capsule_text = gr.Code(value=DEFAULT_CAPSULE, language="json", label="Data Capsule (JSON)", lines=12)
            purpose = gr.Dropdown(["code-edit","summarize","translate","embed","describe-image","train-model","export-data"], value="code-edit", label="Purpose")
            jurisdiction = gr.Dropdown(["NZ","AU","global","US"], value="NZ", label="Jurisdiction")
            modalities = gr.CheckboxGroup(["text","vision","audio"], value=["text"], label="Modalities")
            task_hint = gr.Textbox(label="Task hint", placeholder="coder, reasoning, vision...")
            with gr.Row():
                remote_model = gr.Checkbox(False, label="Remote frontier")
                for_training = gr.Checkbox(False, label="For training")
                requires_local = gr.Checkbox(True, label="Require local")
            export_to = gr.Textbox(label="Export destination", placeholder="local-ide")
            approver = gr.Textbox(label="Approver", placeholder="data-officer@org")
            go = gr.Button("Evaluate & Route", variant="primary")
        with gr.Column():
            out = gr.HTML()
    gr.Examples(
        examples=[
            [DEFAULT_CAPSULE, "code-edit", "NZ", False, False, "local-ide", "", ["text"], "coder", True],
            [DEFAULT_CAPSULE, "train-model", "NZ", False, True, "", "", ["text"], "", True],
            [DEFAULT_CAPSULE, "code-edit", "US", False, False, "", "", ["text"], "coder", True],
            [DEFAULT_CAPSULE, "code-edit", "NZ", True, False, "", "", ["text"], "coder", False],
        ],
        inputs=[capsule_text, purpose, jurisdiction, remote_model, for_training, export_to, approver, modalities, task_hint, requires_local],
    )
    go.click(render, [capsule_text, purpose, jurisdiction, remote_model, for_training, export_to, approver, modalities, task_hint, requires_local], 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)*")


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