import json

import gradio as gr
import pandas as pd
import spaces

from inference import (
    CHAMPION,
    MODELS,
    REFERENCE,
    generate_text,
    hardware_summary,
    start_stream,
)
from portugality import compare_table, looks_like_correction, portugality

SYSTEM_PROMPT = open("system_prompt.txt", encoding="utf-8").read()

# Extra instruction when the user wants to see what changed (JSON output).
EXPLAIN_SUFFIX = (
    "\n\n## Modo explicado\nEm vez de texto simples, devolve um JSON válido com "
    'as chaves: "corrigido" (string) e "alteracoes" (lista de objetos com '
    '"original", "correcao", "motivo"). Não escrevas nada fora do JSON.'
)

REGISTER_MAP = {
    "Automático": "",
    "Informal (tu)": "\n\nUsa tratamento informal por 'tu'.",
    "Formal (o senhor/a senhora)": "\n\nUsa tratamento formal por 'o senhor/a senhora'.",
}

CATEGORY_LABELS = {
    "lexicon": "Léxico pt-BR",
    "spelling": "Ortografia pt-BR",
    "gerund": "Gerúndio",
    "address": "Tratamento (você / a gente)",
    "proclisis": "Próclise inicial",
}


def _wrap_user_text(text):
    """Delimit the input to reduce the risk of prompt injection."""
    return f"<<<TEXTO_A_CORRIGIR>>>\n{text}\n<<<FIM>>>"


def _ipt_badge(before_text, after_text):
    """Markdown line comparing the IPT of the input and the output."""
    before, _ = portugality(before_text)
    after, _ = portugality(after_text)
    arrow = "📈" if after > before else ("📉" if after < before else "➡️")
    return f"**IPT** (portugalidade): `{before}` → `{after}` {arrow}"


# --------------------------------------------------------------------------
# Tab 1 - single correction with streaming
# --------------------------------------------------------------------------
@spaces.GPU(duration=120)
def correct(text, model_name, register, explain):
    if not text or not text.strip():
        yield "Escreve ou cola texto para corrigir.", ""
        return

    system = SYSTEM_PROMPT + REGISTER_MAP.get(register, "")
    if explain:
        system += EXPLAIN_SUFFIX

    output = ""
    try:
        for chunk in start_stream(model_name, system, _wrap_user_text(text)):
            output += chunk
            yield output, gr.update()
    except Exception as exc:
        yield f"⚠️ **Erro ao gerar:** {exc}", ""
        return

    corrected = output
    # If an explanation was requested, try to pretty-print the JSON at the end.
    if explain:
        try:
            data = json.loads(output)
            corrected = data["corrigido"]
            lines = [corrected, "\n---\n**Alterações:**"]
            for change in data.get("alteracoes", []):
                lines.append(
                    f"- `{change['original']}` → **{change['correcao']}** "
                    f"— {change['motivo']}"
                )
            output = "\n".join(lines)
        except Exception:
            pass  # not valid JSON -> show it raw

    badge = _ipt_badge(text, corrected)
    if not looks_like_correction(text, corrected):
        badge = (
            "⚠️ A resposta não parece uma correção válida — experimenta outro "
            "modelo.\n\n" + badge
        )
    yield output, badge


# --------------------------------------------------------------------------
# Tab 2 - multi-model benchmark with the Portugality Index (IPT)
# --------------------------------------------------------------------------
def _rows_to_frames(rows, reference):
    """Split ranked rows into a display dataframe and a bar-plot dataframe."""
    table = pd.DataFrame(
        [
            {
                "Modelo": row["model"],
                "IPT": row["IPT"],
                "Marcadores": row["weighted_markers"],
                f"Δ vs {reference}": row["delta_vs_ref"],
                "Texto corrigido": row["corrected"],
            }
            for row in rows
        ]
    )
    plot = pd.DataFrame([{"Modelo": r["model"], "IPT": r["IPT"]} for r in rows])
    return table, plot


@spaces.GPU(duration=300)
def benchmark(text, selected, reference):
    if not text or not text.strip():
        yield "Escreve ou cola texto para comparar.", None, None
        return
    if not selected:
        yield "Seleciona pelo menos um modelo.", None, None
        return

    # Make sure the reference model is always evaluated so deltas exist.
    order = list(dict.fromkeys(selected + ([reference] if reference else [])))
    user_text = _wrap_user_text(text)

    results = {}
    log_lines = []
    for name in order:
        log_lines.append(f"⏳ A processar **{name}**…")
        yield "\n\n".join(log_lines), gr.update(), gr.update()

        try:
            corrected = generate_text(name, SYSTEM_PROMPT, user_text)
        except Exception as exc:  # gated model, OOM, etc. - keep going
            log_lines[-1] = f"⚠️ **{name}** falhou: {exc}"
            yield "\n\n".join(log_lines), gr.update(), gr.update()
            continue

        if not looks_like_correction(text, corrected):
            log_lines[-1] = (
                f"⚠️ **{name}** excluído: a resposta não parece uma correção "
                f"(início: `{corrected[:80]}`)"
            )
            yield "\n\n".join(log_lines), gr.update(), gr.update()
            continue

        results[name] = corrected
        rows = compare_table(results, reference)
        table, plot = _rows_to_frames(rows, reference)
        log_lines[-1] = f"✅ **{name}** — IPT={rows[0]['IPT']}"
        yield "\n\n".join(log_lines), table, plot

    if results:
        rows = compare_table(results, reference)
        table, plot = _rows_to_frames(rows, reference)
        best = rows[0]
        summary = [f"🏁 **Terminado.** Melhor: **{best['model']}** (IPT={best['IPT']})"]
        if reference in results:
            for row in rows:
                if row["model"] != reference:
                    summary.append(
                        f"- {row['model']}: {row['delta_vs_ref']:+.1f} IPT vs {reference}"
                    )
        yield "\n".join(summary), table, plot


# --------------------------------------------------------------------------
# Tab 3 - instant IPT analysis (deterministic, no GPU)
# --------------------------------------------------------------------------
def analyse(text):
    if not text or not text.strip():
        return "Escreve ou cola texto para analisar.", None

    score, breakdown = portugality(text)
    if score >= 90:
        verdict = "🇵🇹 Muito europeu"
    elif score >= 60:
        verdict = "🤝 Misto"
    else:
        verdict = "🇧🇷 Muito brasileiro"

    header = (
        f"## IPT = {score} / 100 — {verdict}\n"
        f"{breakdown['n_words']} palavras · "
        f"{breakdown['weighted_markers']} marcadores ponderados"
    )

    rows = []
    for key, label in CATEGORY_LABELS.items():
        for match in breakdown.get(key, []):
            rows.append({"Categoria": label, "Marcador encontrado": match})
    if not rows:
        return header + "\n\nNenhum marcador brasileiro encontrado. 🎉", None
    return header, pd.DataFrame(rows)


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------
EXAMPLES = [
    ["Você pode pegar o ônibus, tomar um café da manhã e depois ir no banheiro do trem."],
    ["Meu celular tá na geladeira, foi mal, salvei o arquivo errado no seu usuário."],
]

with gr.Blocks(title="Des-abrasileirador · pt-PT") as demo:
    gr.Markdown(
        "# 🇵🇹 Des-abrasileirador\n"
        "Reescreve texto em **português europeu** correto com o modelo "
        "[AMALIA-9B](https://hf.co/amalia-llm/AMALIA-9B-0626-SFT) e mede a "
        "**portugalidade** com um índice determinístico (IPT)."
    )

    with gr.Tab("Corrigir"):
        with gr.Row():
            with gr.Column():
                text_in = gr.Textbox(
                    label="Texto original",
                    lines=10,
                    placeholder="Vou pegar o ônibus e tomar café da manhã…",
                )
                with gr.Row():
                    model_sel = gr.Dropdown(
                        list(MODELS),
                        value=CHAMPION,
                        label="Modelo",
                    )
                    register = gr.Dropdown(
                        list(REGISTER_MAP),
                        value="Automático",
                        label="Registo",
                    )
                    explain = gr.Checkbox(label="Mostrar alterações", value=False)
                with gr.Row():
                    correct_btn = gr.Button("Corrigir", variant="primary")
                    stop_btn = gr.Button("Parar", variant="stop")
            with gr.Column():
                text_out = gr.Markdown(label="Resultado")
                ipt_out = gr.Markdown()

        correct_evt = correct_btn.click(
            correct, [text_in, model_sel, register, explain], [text_out, ipt_out]
        )
        stop_btn.click(None, cancels=[correct_evt])
        gr.Examples(EXAMPLES, inputs=text_in)

    with gr.Tab("Comparar modelos (IPT)"):
        gr.Markdown(
            "Corre o mesmo texto por vários modelos e mede o **Índice de "
            "Portugalidade** (0–100). O delta é calculado contra o modelo de "
            "referência — a baseline justa da mesma dimensão do AMALIA."
        )
        bench_in = gr.Textbox(
            label="Texto original",
            lines=6,
            placeholder="Você pode pegar o ônibus e tomar café da manhã…",
        )
        with gr.Row():
            model_pick = gr.CheckboxGroup(
                list(MODELS),
                value=[CHAMPION, REFERENCE],
                label="Modelos a comparar",
            )
            ref_pick = gr.Dropdown(
                list(MODELS),
                value=REFERENCE,
                label="Referência (para o delta)",
            )
        bench_btn = gr.Button("Comparar", variant="primary")

        bench_status = gr.Markdown()
        bench_plot = gr.BarPlot(
            x="Modelo",
            y="IPT",
            y_lim=[0, 100],
            title="Índice de Portugalidade por modelo",
            tooltip=["Modelo", "IPT"],
        )
        bench_table = gr.Dataframe(
            headers=["Modelo", "IPT", "Marcadores", "Δ", "Texto corrigido"],
            wrap=True,
            label="Detalhe",
        )

        bench_btn.click(
            benchmark,
            [bench_in, model_pick, ref_pick],
            [bench_status, bench_table, bench_plot],
        )
        gr.Examples(EXAMPLES, inputs=bench_in)

    with gr.Tab("Analisar (IPT)"):
        gr.Markdown(
            "Mede a portugalidade de qualquer texto **instantaneamente** — o IPT "
            "é determinístico e não usa GPU nem modelos."
        )
        analyse_in = gr.Textbox(
            label="Texto a analisar",
            lines=8,
            placeholder="Cola aqui um texto para ver os marcadores brasileiros…",
        )
        analyse_btn = gr.Button("Analisar", variant="primary")
        analyse_out = gr.Markdown()
        analyse_table = gr.Dataframe(
            headers=["Categoria", "Marcador encontrado"],
            wrap=True,
            label="Marcadores",
        )
        analyse_btn.click(analyse, analyse_in, [analyse_out, analyse_table])
        gr.Examples(EXAMPLES, inputs=analyse_in)

    gr.Markdown(
        f"⚙️ Inferência: **{hardware_summary()}**\n\n"
        "ℹ️ Alguns modelos (Llama, Gemma) são *gated* no Hugging Face: define o "
        "segredo `HF_TOKEN` no Space e aceita as licenças para os poderes usar."
    )

demo.queue().launch(theme=gr.themes.Soft())