import os, json, html, traceback
import gradio as gr
import pandas as pd
import spaces
from google import genai
from google.genai import types

import judge_prompts as J
import trust_analysis as ta

# ------------------------------------------------------------------ config
GEMINI_KEY   = os.environ.get("GEMINI_API_KEY")
APP_PASSWORD = os.environ.get("APP_PASSWORD", "")
MODEL        = os.environ.get("JUDGE_MODEL", "gemini-3.1-flash-lite")

MAX_SOURCE_CHARS  = 50000
MAX_SUMMARY_CHARS = 6000
SESSION_LIMIT     = 15

client = genai.Client(api_key=GEMINI_KEY) if GEMINI_KEY else None

DIMENSIONS = {
    "Faithfulness":              ("faithfulness",   "faithfulness_score",        "unsupported_claims", ("span","type","explanation")),
    "Overclaiming":              ("overclaiming",   "overclaiming_score",        "overclaims",         ("span","type","explanation")),
    "Uncertainty preservation":  ("uncertainty",    "uncertainty_score",         "dropped_hedges",     ("rendered_as","source_hedge",None)),
    "Methods & limitations":     ("methods_limits", "methods_limitations_score", None,                 None),
}
ALL_DIMS = list(DIMENSIONS)


# ------------------------------------------------------------------ judge
@spaces.GPU
def _zeroGPU_dummy():
    return True

def parse_json(text):
    obj, _ = json.JSONDecoder().raw_decode(text.strip())
    return obj

def judge(dimension, task, source, output):
    prompt = J.build_prompt(dimension, task, source, output)
    resp = client.models.generate_content(
        model=MODEL, contents=prompt,
        config=types.GenerateContentConfig(
            system_instruction=J.SYSTEM_PROMPT, temperature=0,
            response_mime_type="application/json"))
    return parse_json(resp.text)

def screen(output):
    df = ta.screen_validity(pd.DataFrame([{"generated_output": output,
                                           "has_output": bool(str(output).strip())}]))
    row = df.iloc[0]
    return bool(row["is_valid"]), str(row["degeneracy_type"])

# ------------------------------------------------------------------ rendering
def esc(s): return html.escape(str(s))

def score_colour(s):
    try: s = float(s)
    except: return "#888"
    return "#b5523f" if s <= 2 else "#c9a227" if s < 4 else "#2e8f6a"

def highlight(summary, spans):
    out = esc(summary)
    for sp in sorted({s for s in spans if s and len(str(s)) > 4}, key=len, reverse=True):
        out = out.replace(esc(sp), f"<mark style='background:#ffe0b3;color:#000000'>{esc(sp)}</mark>")
    return out.replace("\n", "<br>")

TEXT = "#000000"
MUTED = "#1a1a1a"

def render_scores(results, valid, degen):
    """Middle column: validity banner + one box per scored dimension."""
    parts = []
    if not valid:
        parts.append(f"<div style='padding:10px;border-left:4px solid #b5523f;background:#fbeae6;"
                     f"font-family:system-ui;color:#000000'>"
                     f"<b>Validity screen: DEGENERATE</b> — flagged as <code>{esc(degen)}</code>. "
                     f"Fault-absence dimensions (overclaiming, uncertainty) are not scored for "
                     f"non-functional output; faithfulness is reported as minimal.</div>")
    else:
        parts.append(f"<div style='padding:8px;border-left:4px solid #2e8f6a;background:#eaf6f0;"
                     f"font-family:system-ui;color:#000000'><b>Validity screen: functional output</b></div>")

    all_spans = []
    for label, res in results.items():
        key, sfield, ifield, ikeys = DIMENSIONS[label]
        if "error" in res:
            parts.append(f"<div style='margin-top:12px;color:#000000'><b>{esc(label)}</b>: "
                         f"<span style='color:#b5523f'>could not score ({esc(res['error'])})</span></div>")
            continue
        score = res.get(sfield, "–")
        col = score_colour(score)
        head = (f"<div style='margin-top:14px;padding:8px;border:1px solid #e3e7ee;border-radius:8px;"
                f"font-family:system-ui;background:#ffffff;color:#000000'>"
                f"<span style='font-family:system-ui;font-size:15px;font-weight:700'>{esc(label)}</span> "
                f"<span style='float:right;font-weight:700;color:{col}'>{esc(score)} / 5</span>")
        body = ""
        if label == "Methods & limitations":
            flags = [("Methods", res.get("methods_present")),
                     ("Limitations", res.get("limitations_present")),
                     ("Population", res.get("population_present"))]
            body += f"<div style='margin-top:6px;font-size:13px;color:#000000'>" + " · ".join(
                f"{n}: {'✓' if v else '✗'}" for n, v in flags) + "</div>"
        elif ifield:
            issues = res.get(ifield) or []
            if issues:
                body += f"<ul style='margin:6px 0 0 0;font-size:13px;color:#000000'>"
                for it in issues:
                    if not isinstance(it, dict): continue
                    span = it.get(ikeys[0], "")
                    all_spans.append(span)
                    typ = it.get(ikeys[1], "") if ikeys[1] else ""
                    exp = it.get(ikeys[2], "") if ikeys[2] else ""
                    line = f"<li><i>“{esc(span)}”</i>"
                    if typ: line += f" — <code>{esc(typ)}</code>"
                    if exp: line += f": {esc(exp)}"
                    body += line + "</li>"
                body += "</ul>"
            else:
                body += f"<div style='margin-top:6px;font-size:13px;color:#2e8f6a'>No issues flagged.</div>"
        parts.append(head + body + "</div>")

    html_out = "<div style='font-family:system-ui;color:" + TEXT + "'>" + "".join(parts) + "</div>"
    return html_out, all_spans


def render_summary(summary, all_spans):
    hi = highlight(summary, all_spans)
    return (f"<div style='font-family:system-ui;color:#000000'>"
           f"<div style='padding:10px;background:#fbfaf7;border:1px solid #ece7dc;border-radius:8px;"
           f"font-size:14px;line-height:1.6;color:#000000'>{hi}</div></div>")

# ------------------------------------------------------------------ main callback
def _err(msg, remaining):
    return f"<b style='color:#b5523f'>{esc(msg)}</b>", "", remaining

def evaluate(source, summary, task, dims, remaining):
    if client is None:
        return _err("Server not configured: GEMINI_API_KEY secret is missing.", remaining)
    if remaining <= 0:
        return _err("Session evaluation limit reached. Reload the page to reset.", remaining)
    source, summary = (source or "").strip(), (summary or "").strip()
    if not source or not summary:
        return _err("Please provide both a source text and an AI summary.", remaining)
    if len(source) > MAX_SOURCE_CHARS:
        return _err(f"Source too long ({len(source)} chars; max {MAX_SOURCE_CHARS}). Please trim.", remaining)
    if len(summary) > MAX_SUMMARY_CHARS:
        return _err(f"Summary too long ({len(summary)} chars; max {MAX_SUMMARY_CHARS}). Please trim.", remaining)
    if not dims:
        return _err("Select at least one dimension to run.", remaining)

    task_key = "structured" if task.startswith("Structured") else "lay"
    valid, degen = screen(summary)

    results = {}
    for label in dims:
        key, sfield, *_ = DIMENSIONS[label]
        # gate fault-absence dims on validity
        if not valid and key in ("overclaiming", "uncertainty"):
            continue
        if not valid and key in ("faithfulness", "methods_limits"):
            results[label] = {sfield: 1}
            continue
        try:
            results[label] = judge(key, task_key, source, summary)
        except Exception as e:
            results[label] = {"error": str(e)[:120]}

    scores_html, all_spans = render_scores(
        {k: v for k, v in results.items() if k in DIMENSIONS}, valid, degen)
    summary_html = render_summary(summary, all_spans)
    return scores_html, summary_html, remaining - 1

# ------------------------------------------------------------------ UI
CSS = ".gradio-container{max-width:1450px!important}"
with gr.Blocks(title="Trust Evaluation Toolkit") as demo:
    gr.Markdown("## Biomedical Summary Trustworthiness Toolkit\n"
                "Scores an AI-generated summary against a validated multi-dimensional trust rubric. "
                "Paste the source article and the AI summary, choose the register and dimensions, and run.")
    remaining = gr.State(SESSION_LIMIT)
    with gr.Row():
        with gr.Column(scale=4):
            source = gr.Textbox(label="Source text (the article)", lines=12, max_lines=20)
            summary = gr.Textbox(label="AI-generated summary", lines=8, max_lines=14)
            task = gr.Radio(["Structured summary", "Lay summary"], value="Structured summary",
                            label="Register")
            dims = gr.CheckboxGroup(ALL_DIMS, value=ALL_DIMS, label="Trust dimensions (LLM-judge)")
            btn = gr.Button("Evaluate", variant="primary")
        with gr.Column(scale=3):
            gr.Markdown("**Scores**")
            scores_out = gr.HTML()
        with gr.Column(scale=3):
            gr.Markdown("**Summary with flagged spans**")
            summary_out = gr.HTML()
    btn.click(evaluate, [source, summary, task, dims, remaining], [scores_out, summary_out, remaining])
    gr.Markdown("<small>Scores are produced by an LLM judge validated against human ratings; they are "
                "comparative, not absolute ground truth. Access is rate-limited per session.</small>")

def check_auth(username, password):
    return bool(APP_PASSWORD) and password == APP_PASSWORD

demo.launch(
    auth=check_auth,
    auth_message="Enter the shared access password.",
    css=CSS,
)