"""Q-Prime demo Space (Gradio).

Calls the QGI managed API; does not run any model weights locally.
Set QGI_API_KEY and QGI_API_URL as Space Secrets before launch.
"""

from __future__ import annotations

import os
import time
from dataclasses import dataclass

import gradio as gr
import requests

API_URL = os.environ.get("QGI_API_URL", "https://api.qgi.dev/v1/qprime/embed")
API_KEY = os.environ.get("QGI_API_KEY", "")
MODEL_URL = "https://huggingface.co/QGI-dev/q-prime"
PAPER_URL = "https://huggingface.co/QGI-dev/q-prime#accompanying-papers-qag-series"
GATED_URL = "https://huggingface.co/QGI-dev/q-prime"  # has Request access button

DEFAULT_LABELS = (
    "mandatory, prohibitive, conditional, default, exception, null"
)

DEFAULT_A = "A regulated entity must report any incident within 72 hours."
DEFAULT_B = (
    "A regulated entity must not report an incident while a law-enforcement "
    "investigation is active."
)

RATE_WINDOW_SECS = 60
RATE_MAX_CALLS = 30
_session_log: dict[str, list[float]] = {}


@dataclass
class Signals:
    relevance: list[float]
    overlap: float
    conflict: float
    predicates: list[str]
    born_rule: dict[str, float]


def _rate_limit(session_id: str) -> bool:
    now = time.time()
    hits = _session_log.setdefault(session_id, [])
    hits[:] = [t for t in hits if now - t < RATE_WINDOW_SECS]
    if len(hits) >= RATE_MAX_CALLS:
        return False
    hits.append(now)
    return True


def _call_api(clause_a: str, clause_b: str, labels: list[str]) -> Signals:
    if not API_KEY:
        raise gr.Error(
            "This demo is not yet connected to the Q-Prime API. "
            "Join the waitlist at https://qgi.dev or request evaluation access "
            f"at {GATED_URL}."
        )
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json",
        "User-Agent": "qprime-demo-space/1.0",
    }
    payload = {
        "inputs": [clause_a, clause_b],
        "tasks": ["relevance", "conflict", "overlap", "predicate", "classify"],
        "classify_labels": labels,
    }
    try:
        r = requests.post(API_URL, headers=headers, json=payload, timeout=30)
    except requests.RequestException as exc:
        raise gr.Error(f"Q-Prime API unreachable: {exc}") from exc

    if r.status_code == 401:
        raise gr.Error(
            "Space's QGI_API_KEY was rejected. QGI staff: rotate in Space secrets."
        )
    if r.status_code == 429:
        raise gr.Error("Q-Prime API rate limit hit. Try again in a moment.")
    if r.status_code >= 400:
        raise gr.Error(f"Q-Prime API error {r.status_code}: {r.text[:200]}")

    data = r.json()
    return Signals(
        relevance=data.get("relevance", [0.0, 0.0]),
        overlap=float(data.get("overlap", [[0.0]])[0][1]),
        conflict=float(data.get("conflict", [[0.0]])[0][1]),
        predicates=data.get("predicate", [[], []])[0],
        born_rule=data.get("classify", [{}])[0],
    )


def _format_signed_meter(name: str, value: float) -> str:
    v = max(-1.0, min(1.0, value))
    pos_pct = max(0.0, v) * 100.0
    neg_pct = max(0.0, -v) * 100.0
    colour = "#2E4A8F" if v >= 0 else "#B03A2E"
    if v > 0.05:
        label = "reinforce"
    elif v < -0.05:
        label = "contradict"
    else:
        label = "neutral"
    track_style = (
        "flex:1;display:flex;height:18px;background:#f4f4f6;"
        "border-radius:4px;overflow:hidden"
    )
    left_half_style = "width:50%;display:flex;justify-content:flex-end"
    right_half_style = "width:50%;display:flex;justify-content:flex-start"
    left_fill = (
        f"<div style='height:100%;width:{neg_pct:.1f}%;background:{colour}'></div>"
    )
    right_fill = (
        f"<div style='height:100%;width:{pos_pct:.1f}%;background:{colour}'></div>"
    )
    return (
        "<div style='display:flex;align-items:center;gap:8px;margin:4px 0'>"
        f"<div style='width:150px;font-weight:600'>{name}</div>"
        f"<div style='{track_style}'>"
        f"<div style='{left_half_style}'>{left_fill}</div>"
        f"<div style='{right_half_style}'>{right_fill}</div>"
        "</div>"
        "<div style='width:150px;text-align:right;"
        "font-variant-numeric:tabular-nums'>"
        f"{v:+.3f} <em style='color:#666'>({label})</em>"
        "</div>"
        "</div>"
    )


def _format_probs(probs: dict[str, float]) -> str:
    if not probs:
        return "<em>(no classification returned)</em>"
    rows = sorted(probs.items(), key=lambda kv: -kv[1])
    out = ["<table style='border-collapse:collapse;width:100%'>"]
    for label, p in rows:
        p = max(0.0, min(1.0, float(p)))
        out.append(
            "<tr>"
            f"<td style='padding:2px 8px;font-weight:600'>{label}</td>"
            "<td style='padding:2px 8px;width:70%'>"
            f"<div style='height:12px;background:#f4f4f6;border-radius:3px'>"
            f"<div style='height:100%;width:{p*100:.1f}%;background:#2E4A8F;border-radius:3px'></div>"
            "</div></td>"
            f"<td style='padding:2px 8px;font-variant-numeric:tabular-nums'>{p:.3f}</td>"
            "</tr>"
        )
    out.append("</table>")
    return "".join(out)


def compare(clause_a: str, clause_b: str, labels_raw: str, request: gr.Request):
    if not clause_a.strip() or not clause_b.strip():
        raise gr.Error("Both clauses must be non-empty.")
    if len(clause_a) > 2000 or len(clause_b) > 2000:
        raise gr.Error("Clauses limited to 2000 chars in the demo.")

    session_id = getattr(request, "session_hash", "default") or "default"
    if not _rate_limit(session_id):
        raise gr.Error(
            f"Demo rate limit hit ({RATE_MAX_CALLS} calls / {RATE_WINDOW_SECS}s). "
            f"Request a full evaluation key at {GATED_URL}."
        )

    labels = [x.strip() for x in labels_raw.split(",") if x.strip()]
    sig = _call_api(clause_a, clause_b, labels)

    header = "<h3 style='margin-bottom:4px'>Intelligence signals</h3>"
    signed = _format_signed_meter("Conflict (signed)", sig.conflict)
    overlap = _format_signed_meter("Overlap", sig.overlap)
    rel_a = _format_signed_meter("Relevance A", sig.relevance[0] if sig.relevance else 0.0)
    rel_b = _format_signed_meter("Relevance B", sig.relevance[1] if len(sig.relevance) > 1 else 0.0)

    predicate_html = (
        "<h4 style='margin-top:14px'>Predicates (clause A)</h4><ul>"
        + "".join(f"<li><code>{p}</code></li>" for p in sig.predicates)
        + "</ul>"
        if sig.predicates
        else "<h4 style='margin-top:14px'>Predicates (clause A)</h4><em>none</em>"
    )

    classify_html = (
        "<h4 style='margin-top:14px'>Born-rule classifier (clause A)</h4>"
        + _format_probs(sig.born_rule)
    )

    return header + signed + overlap + rel_a + rel_b + predicate_html + classify_html


def _disabled_banner() -> str:
    if API_KEY:
        return ""
    return (
        "<div style='padding:10px 14px;background:#fff3cd;border:1px solid #ffeeba;"
        "border-radius:6px;margin-bottom:12px;color:#664d03'>"
        "<strong>Demo is inactive.</strong> This Space ships without a live key. "
        f"Request an evaluation API key via the <a href='{GATED_URL}'>Q-Prime model "
        "card</a> to try the full API, or read the "
        f"<a href='{PAPER_URL}'>accompanying papers</a>."
        "</div>"
    )


EXAMPLES = [
    [
        DEFAULT_A,
        DEFAULT_B,
        DEFAULT_LABELS,
    ],
    [
        "All employees shall complete annual security training.",
        "Employees on medical leave are exempt from annual security training.",
        DEFAULT_LABELS,
    ],
    [
        "Customer data must be encrypted at rest using AES-256.",
        "Customer data may be stored in plaintext in development environments.",
        DEFAULT_LABELS,
    ],
]


with gr.Blocks(
    title="Q-Prime — QAG conflict and intelligence signals",
    theme=gr.themes.Soft(primary_hue="blue"),
) as demo:
    gr.Markdown(
        f"""
# Q-Prime — live QAG signals

Enter two rule-bearing clauses and the demo calls the Q-Prime API to return the
signed conflict signal, the relevance scores, the overlap, predicates, and a
zero-shot Born-rule classification.

Clauses never leave QGI infrastructure; this Space forwards them to the QGI
managed API. The model weights are not distributed.

Model card: [{MODEL_URL}]({MODEL_URL}) · Papers: [QAG series]({PAPER_URL})
"""
    )

    gr.HTML(_disabled_banner())

    with gr.Row():
        with gr.Column():
            in_a = gr.Textbox(
                label="Clause A",
                value=DEFAULT_A,
                lines=3,
                placeholder="A regulated entity must report ...",
            )
            in_b = gr.Textbox(
                label="Clause B",
                value=DEFAULT_B,
                lines=3,
                placeholder="A regulated entity must not ...",
            )
            in_labels = gr.Textbox(
                label="Zero-shot labels (comma-separated)",
                value=DEFAULT_LABELS,
                lines=1,
            )
            btn = gr.Button("Compare", variant="primary")
        with gr.Column():
            out = gr.HTML(
                "<em>Results will appear here.</em>",
                elem_id="qp-signals",
            )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[in_a, in_b, in_labels],
        label="Try one",
    )

    btn.click(fn=compare, inputs=[in_a, in_b, in_labels], outputs=out)

    gr.Markdown(
        f"""
---

Q-Prime is a commercial model licensed under the QGI Commercial Model License v1.0.
[Request evaluation access]({GATED_URL}) · `contact@qgi.dev` · © 2025–2026 Quantum General Intelligence, Inc.
"""
    )


if __name__ == "__main__":
    demo.queue(max_size=32).launch()