from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import streamlit as st
from huggingface_hub import hf_hub_download

APP_DIR = Path(__file__).resolve().parent
ROOT = APP_DIR.parent
DATASET_DIR = ROOT / "data" / "dataset"
SOURCES_DIR = ROOT / "data" / "sources"
MODEL_METADATA_PATH = ROOT / "data" / "model_metadata" / "openrouter_models.json"
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "hoololi/work-translation")
USE_HF_DATASET = os.environ.get("USE_HF_DATASET", "auto").lower()

METRICS = {
    "Semantic similarity": "semantic_similarity_to_original",
    "Vocabulary overlap": "word_overlap_to_original",
    "Surface/edit similarity": "edit_similarity_to_original",
    "Length ratio": "word_count_ratio_to_original",
}

MODEL_COLORS = {
    "openai/gpt-4o-mini": "#4E79A7",
    "poolside/laguna-s-2.1": "#F28E2B",
    "thinkingmachines/inkling-small": "#59A14F",
    "poolside/laguna-s-2.1:free": "#E15759",
}
PIVOT_COLORS = {
    "English": "#8E63B6",
    "Italian": "#D65F5F",
    "Slovak": "#6A9F58",
    "German": "#C49A2C",
    "Spanish": "#3BA3A3",
}
FALLBACK_COLORS = px.colors.qualitative.Set2

METRIC_DEFINITIONS = {
    "Semantic similarity": "Cosine similarity between embeddings of the original text and the translated-back text. Higher means semantically closer.",
    "Vocabulary overlap": "Jaccard similarity of unique lowercased words: shared words divided by all words appearing in either text. Higher means more original vocabulary is preserved.",
    "Surface/edit similarity": "Normalized word-level edit similarity. Higher means fewer insertions, deletions, substitutions, or word-order changes are needed to transform one text into the other.",
    "Length ratio": "Word count of the translated-back text divided by word count of the original. 1.00 means same length; below 1.00 means shorter; above 1.00 means longer.",
}


@st.cache_data
def dataset_file(filename: str) -> Path:
    """Return a local dataset file, downloading from HF if needed/configured."""
    local_path = DATASET_DIR / filename
    if USE_HF_DATASET != "1" and local_path.exists():
        return local_path
    return Path(hf_hub_download(repo_id=HF_DATASET_REPO, repo_type="dataset", filename=filename))


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    with path.open("r", encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def parse_source_file(path: Path) -> dict[str, Any]:
    raw = path.read_text(encoding="utf-8")
    lines = raw.splitlines()
    meta: dict[str, str] = {}
    body = raw
    if lines and lines[0].strip() == "---":
        end = next((i for i, line in enumerate(lines[1:], start=1) if line.strip() == "---"), None)
        if end is not None:
            for line in lines[1:end]:
                if line.strip() and ":" in line:
                    key, value = line.split(":", 1)
                    meta[key.strip()] = value.strip().strip('"').strip("'")
            body = "\n".join(lines[end + 1 :]).strip("\n")
    meta["source_file"] = path.name
    meta["text"] = body
    return meta


@st.cache_data
def load_sources() -> pd.DataFrame:
    # Prefer the dataset-level source metadata; it contains the source text and
    # avoids downloading individual Markdown files from the dataset repository.
    try:
        sources_jsonl = dataset_file("sources.jsonl")
        rows = read_jsonl(sources_jsonl)
        if rows:
            return pd.DataFrame(rows)
    except Exception:
        pass
    return pd.DataFrame([parse_source_file(p) for p in sorted(SOURCES_DIR.glob("*.md"))])


@st.cache_data
def load_translations() -> pd.DataFrame:
    rows = read_jsonl(dataset_file("translations.jsonl"))
    if not rows:
        return pd.DataFrame()
    df = pd.DataFrame(rows)
    for col in ["experiment_id", "timestamp", "translation_model", "provider", "pivot_language", "language_pair", "source_id", "text"]:
        if col not in df.columns:
            df[col] = None
    df["step"] = pd.to_numeric(df["step"], errors="coerce").astype("Int64")
    return df


@st.cache_data
def load_metrics() -> pd.DataFrame:
    # Prefer the clean pivot-to-original table, which contains the main app metrics.
    try:
        path = dataset_file("pivot_to_original_metrics.csv")
    except Exception:
        return pd.DataFrame()
    df = pd.read_csv(path)
    for col in ["source_id", "pivot_language", "translation_model", "embedding_model"]:
        if col not in df.columns:
            df[col] = None
    df = df.dropna(subset=["source_id", "pivot_language", "translation_model"])
    df["step"] = pd.to_numeric(df["step"], errors="coerce").astype("Int64")
    df = df[df["step"].notna() & (df["step"].astype(int) % 2 == 0)]
    for col in METRICS.values():
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    keys = ["source_id", "pivot_language", "translation_model", "embedding_model", "step"]
    df = df.drop_duplicates(subset=keys, keep="last")
    return df


@st.cache_data
def load_model_metadata() -> dict[str, Any]:
    try:
        path = dataset_file("openrouter_models.json")
    except Exception:
        path = MODEL_METADATA_PATH
    if not path.exists():
        return {}
    payload = json.loads(path.read_text(encoding="utf-8"))
    return {m.get("id"): m for m in payload.get("models", []) if isinstance(m, dict) and m.get("id")}


def model_short(model: str) -> str:
    for prefix in ["openai/", "poolside/", "thinkingmachines/", "voyageai/"]:
        model = model.replace(prefix, "")
    return model


def source_display_options(sources: pd.DataFrame, metric_source_ids: set[str]) -> dict[str, str]:
    options: dict[str, str] = {}
    for _, row in sources.iterrows():
        source_id = row.get("source_id")
        if source_id not in metric_source_ids:
            continue
        slug = row.get("source_slug") or source_id
        category = row.get("category") or ""
        title = row.get("title") or ""
        options[f"{slug} — {title} ({category})"] = source_id
    return options


def original_text_for(sources: pd.DataFrame, source_id: str) -> str:
    row = sources[sources["source_id"] == source_id]
    return "" if row.empty else str(row.iloc[0].get("text", ""))


def looks_like_failed_output(text: str, original_text: str) -> bool:
    normalized = text.strip().casefold()
    if normalized in {"fin du texte.", "fin du texte", "koniec textu.", "koniec textu"}:
        return True
    if original_text and len(text) < max(40, len(original_text) * 0.15):
        return True
    return False


def output_text(translations: pd.DataFrame, source_id: str, model: str, pivot: str, step: int) -> str:
    df = translations[
        (translations["source_id"] == source_id)
        & (translations["translation_model"] == model)
        & (translations["pivot_language"] == pivot)
        & (translations["step"] == step)
    ]
    if df.empty:
        return "No text found for this source/model/pivot/step."
    # If repeated runs exist with same factors, take the latest timestamp.
    df = df.sort_values("timestamp", na_position="first")
    return str(df.iloc[-1]["text"])


def selected_metric_row(metrics: pd.DataFrame, source_id: str, model: str, pivot: str, embedding_model: str, step: int) -> pd.Series | None:
    df = metrics[
        (metrics["source_id"] == source_id)
        & (metrics["translation_model"] == model)
        & (metrics["pivot_language"] == pivot)
        & (metrics["embedding_model"] == embedding_model)
        & (metrics["step"] == step)
    ]
    if df.empty:
        return None
    return df.iloc[-1]


def color_map(values: list[str] | pd.Series, fixed: dict[str, str]) -> dict[str, str]:
    all_values = sorted(str(x) for x in pd.Series(values).dropna().unique())
    return {value: fixed.get(value, FALLBACK_COLORS[i % len(FALLBACK_COLORS)]) for i, value in enumerate(all_values)}


def model_color_map(models: list[str] | pd.Series) -> dict[str, str]:
    return color_map(models, MODEL_COLORS)


def series_name(value: str, series_col: str) -> str:
    return model_short(value) if series_col == "translation_model" else value


def trajectory_chart(df: pd.DataFrame, metric_labels: list[str], series_values: list[str], series_col: str, series_label: str, full_y_axis: bool):
    available = [m for m in METRICS.keys() if m in metric_labels]
    if not available:
        return px.line(title="No metric selected")

    fig = make_subplots(
        rows=2,
        cols=2,
        subplot_titles=available + [""] * (4 - len(available)),
        horizontal_spacing=0.07,
        vertical_spacing=0.14,
    )
    fixed = MODEL_COLORS if series_col == "translation_model" else PIVOT_COLORS
    colors = color_map(df[series_col], fixed)
    steps = sorted(int(x) for x in df["step"].dropna().unique())
    tickvals = [s for s in steps if s == 0 or s % 4 == 0 or s == max(steps)]

    any_trace = False
    for idx, label in enumerate(available):
        row = idx // 2 + 1
        col = idx % 2 + 1
        metric = METRICS[label]
        for value in series_values:
            tmp = df[df[series_col] == value].dropna(subset=[metric]).copy()
            if tmp.empty:
                continue
            any_trace = True
            fig.add_trace(
                go.Scatter(
                    x=tmp["step"].astype(int),
                    y=tmp[metric] * 100,
                    mode="lines+markers",
                    name=series_name(value, series_col),
                    legendgroup=value,
                    showlegend=idx == 0,
                    line=dict(width=1.8, color=colors[value]),
                    marker=dict(size=4, color=colors[value]),
                    hovertemplate=(
                        f"{label}<br>{series_label}: {series_name(value, series_col)}<br>"
                        "step: %{x}<br>value: %{y:.1f}%<extra></extra>"
                    ),
                ),
                row=row,
                col=col,
            )

    if not any_trace:
        return px.line(title="No data for selected series")

    fig.update_annotations(font_size=11, yshift=8)
    fig.update_xaxes(tickmode="array", tickvals=tickvals, title_text="", showline=True, mirror=True, linewidth=1, linecolor="rgba(180,180,180,0.45)", gridcolor="rgba(180,180,180,0.18)")
    fig.update_xaxes(showticklabels=False, row=1)
    fig.update_xaxes(title_text="step", row=2)
    if full_y_axis:
        fig.update_yaxes(title_text="%", range=[0, 105], showline=True, mirror=True, linewidth=1, linecolor="rgba(180,180,180,0.45)", gridcolor="rgba(180,180,180,0.18)")
    else:
        fig.update_yaxes(title_text="%", showline=True, mirror=True, linewidth=1, linecolor="rgba(180,180,180,0.45)", gridcolor="rgba(180,180,180,0.18)")
    fig.update_layout(height=560, margin=dict(l=8, r=8, t=55, b=8), legend=dict(orientation="h", y=1.08, font=dict(size=10), title_text=series_label), legend_itemclick=False, legend_itemdoubleclick=False, font=dict(size=11))
    return fig


def selected_step_bar_chart(df: pd.DataFrame, metric_labels: list[str], series_values: list[str], series_col: str, series_label: str, step: int, full_y_axis: bool):
    rows = []
    for label in metric_labels:
        metric = METRICS[label]
        tmp = df[(df[series_col].isin(series_values)) & (df["step"] == step)].dropna(subset=[metric]).copy()
        for _, r in tmp.iterrows():
            value = str(r[series_col])
            rows.append({
                "metric": label,
                "value": float(r[metric]) * 100,
                series_label: series_name(value, series_col),
                "raw_series": value,
            })
    plot_df = pd.DataFrame(rows)
    if plot_df.empty:
        return px.bar(title="No selected-step metric data")
    fixed = MODEL_COLORS if series_col == "translation_model" else PIVOT_COLORS
    cmap = {series_name(k, series_col): v for k, v in color_map(df[series_col], fixed).items()}
    fig = px.bar(
        plot_df,
        x="metric",
        y="value",
        color=series_label,
        barmode="group",
        text_auto=".1f",
        hover_data=["raw_series"],
        labels={"metric": "metric", "value": f"value at step {step} (%)"},
        color_discrete_map=cmap,
    )
    ymax = float(plot_df["value"].max()) if not plot_df.empty else 1.0
    fig.update_layout(
        height=340,
        margin=dict(l=8, r=8, t=58, b=8),
        legend=dict(orientation="h", y=1.24, x=0, font=dict(size=10), title_text=series_label, bgcolor="rgba(0,0,0,0)"),
        legend_itemclick=False,
        legend_itemdoubleclick=False,
        font=dict(size=11),
    )
    if full_y_axis:
        fig.update_yaxes(range=[0, 105])
    else:
        fig.update_yaxes(range=[0, min(130, max(5, ymax * 1.18))])
    return fig


def aggregate_chart(metrics: pd.DataFrame, embedding_model: str, metric_label: str, full_y_axis: bool):
    metric = METRICS[metric_label]
    df = metrics[metrics["embedding_model"] == embedding_model].dropna(subset=[metric]).copy()
    df["LLM"] = df["translation_model"].map(model_short)
    df["value_percent"] = df[metric] * 100
    fig = px.box(
        df,
        x="step",
        y="value_percent",
        color="LLM",
        points="all",
        hover_data=["source_id", "pivot_language", "translation_model"],
        labels={"step": "step", "value_percent": f"{metric_label} (%)"},
        color_discrete_map={model_short(k): v for k, v in model_color_map(metrics["translation_model"]).items()},
    )
    fig.update_layout(height=280, margin=dict(l=8, r=8, t=18, b=8), legend=dict(orientation="h", y=1.15, font=dict(size=10)))
    if full_y_axis:
        fig.update_yaxes(range=[0, 105])
    return fig


def metric_cards(row: pd.Series | None) -> None:
    cols = st.columns(4)
    for col, (label, metric) in zip(cols, METRICS.items()):
        if row is None or pd.isna(row.get(metric)):
            value = "—"
        else:
            value = f"{float(row[metric]) * 100:.1f}%"
        col.metric(label, value)


def usd_per_million_tokens(value: Any) -> str:
    try:
        return f"${float(value) * 1_000_000:.3f}"
    except (TypeError, ValueError):
        return "—"


def model_page_url(model: str) -> str:
    return f"https://openrouter.ai/{model}"


def show_model_metadata_table(models: list[str], metadata: dict[str, Any]) -> None:
    rows = []
    for model in models:
        info = metadata.get(model, {})
        pricing = info.get("pricing") or {}
        top_provider = info.get("top_provider") or {}
        rows.append({
            "Model": model_short(model),
            "Context tokens": info.get("context_length", "—"),
            "Max output tokens": top_provider.get("max_completion_tokens", "—"),
            "Prompt USD / 1M tokens": usd_per_million_tokens(pricing.get("prompt")),
            "Completion USD / 1M tokens": usd_per_million_tokens(pricing.get("completion")),
            "Model page": model_page_url(model),
        })

    st.dataframe(
        pd.DataFrame(rows),
        use_container_width=True,
        hide_index=True,
        column_config={"Model page": st.column_config.LinkColumn("Model page")},
    )
    st.caption("OpenRouter prices are shown as USD per 1 million input tokens (prompt) or output tokens (completion), not per experiment or per 50-translation trajectory.")


def main() -> None:
    st.set_page_config(page_title="Repeated translation explorer", layout="wide")
    st.title("Repeated translation explorer")

    st.markdown(
        """
<style>
.intro-box {
  background: rgba(180, 180, 180, 0.13);
  border: 1px solid rgba(180, 180, 180, 0.32);
  border-radius: 0.6rem;
  padding: 1.1rem 1.25rem;
  margin: 0.75rem 0 1.25rem 0;
}
.intro-box h3 { margin-top: 0; }
</style>
<div class="intro-box">
<p><strong>What happens to a text when it is translated back and forth repeatedly by an LLM?</strong></p>
<p>This experiment starts from short French texts and sends them through different pivot languages, one translation at a time. Each step is a new, stateless API call.</p>
<p>The resulting translation trajectories, together with several comparison metrics, form a dataset.</p>
<p>This app lets you explore that dataset: compare models and pivot languages, follow how a text changes over successive translations, and inspect the actual outputs behind the metrics.</p>
<p>Metrics are computed on returned-to-French steps only: 0, 2, 4, …</p>
</div>
        """,
        unsafe_allow_html=True,
    )

    sources = load_sources()
    translations = load_translations()
    metrics = load_metrics()
    model_metadata = load_model_metadata()

    if translations.empty or metrics.empty:
        st.error("Missing dataset export. Run `python combine_jsonl.py` and `python export_dataset.py` first.")
        return

    metric_source_ids = set(metrics["source_id"].dropna().unique())
    source_options = source_display_options(sources, metric_source_ids)
    if not source_options:
        st.error("No translated sources found in metrics.")
        return

    with st.sidebar:
        st.header("Choose an experiment")
        st.markdown("**1 · Source**")
        source_label = st.selectbox("Source", list(source_options.keys()), label_visibility="collapsed")
        source_id = source_options[source_label]

        source_metrics_all = metrics[metrics["source_id"] == source_id].copy()

        st.markdown("**2 · Embedding model**")
        embedding_models = sorted(source_metrics_all["embedding_model"].dropna().unique())
        embedding_model = st.selectbox("Embedding model", embedding_models, label_visibility="collapsed")
        source_metrics_all = source_metrics_all[source_metrics_all["embedding_model"] == embedding_model]

        st.markdown("**3 · Compare**")
        comparison_mode = st.radio("Comparison mode", ["LLM models", "Pivot languages"], label_visibility="collapsed")

        if comparison_mode == "LLM models":
            st.markdown("**4 · Pivot language**")
            pivots = sorted(source_metrics_all["pivot_language"].dropna().unique())
            pivot = st.selectbox("Pivot language", pivots, label_visibility="collapsed")
            source_metrics = source_metrics_all[source_metrics_all["pivot_language"] == pivot].copy()

            st.markdown("**5 · LLMs**")
            models = sorted(source_metrics["translation_model"].dropna().unique())
            selected_models = []
            for model in models:
                if st.checkbox(model_short(model), value=True, key=f"model_{source_id}_{pivot}_{embedding_model}_{model}"):
                    selected_models.append(model)
            if not selected_models:
                st.warning("Select at least one LLM.")
                st.stop()
            selected_pivots = [pivot]
            chart_series = selected_models
            chart_series_col = "translation_model"
            chart_series_label = "LLM"
        else:
            st.markdown("**4 · LLM model**")
            models = sorted(source_metrics_all["translation_model"].dropna().unique())
            selected_model = st.selectbox("LLM model", models, format_func=model_short, label_visibility="collapsed")
            source_metrics = source_metrics_all[source_metrics_all["translation_model"] == selected_model].copy()

            st.markdown("**5 · Pivot languages**")
            pivots = sorted(source_metrics["pivot_language"].dropna().unique())
            selected_pivots = []
            for pvt in pivots:
                if st.checkbox(pvt, value=True, key=f"pivot_{source_id}_{selected_model}_{embedding_model}_{pvt}"):
                    selected_pivots.append(pvt)
            if not selected_pivots:
                st.warning("Select at least one pivot language.")
                st.stop()
            pivot = selected_pivots[0]
            selected_models = [selected_model]
            chart_series = selected_pivots
            chart_series_col = "pivot_language"
            chart_series_label = "Pivot"

        selected_metric_labels = list(METRICS.keys())

        st.markdown("**6 · Display**")
        full_y_axis = st.checkbox("Use full 0–100% y-axis", value=False)

    steps = sorted(int(x) for x in source_metrics["step"].dropna().unique())

    source_row = sources[sources["source_id"] == source_id]
    source_meta = source_row.iloc[0] if not source_row.empty else pd.Series(dtype=object)

    with st.container(border=True):
        st.markdown(f"#### Selected source: {source_label}")
        meta_cols = st.columns(4)
        meta_cols[0].write(f"**Source id:** {source_id}")
        meta_cols[1].write(f"**Category:** {source_meta.get('category', '')}")
        meta_cols[2].write(f"**Mode:** {comparison_mode}")
        meta_cols[3].write(f"**Embedding:** {embedding_model}")
        with st.expander("Source/provenance details", expanded=False):
            st.write(f"**Title:** {source_meta.get('title', '')}")
            st.write(f"**Source:** {source_meta.get('source', '')}")
            st.write(f"**License:** {source_meta.get('license', '')}")

    with st.expander("What do these metrics mean?", expanded=False):
        st.markdown("""
<style>
.metric-def-row { display: grid; grid-template-columns: 190px 1fr; gap: 1rem; margin: 0.55rem 0; align-items: start; }
.metric-def-name { color: #4EA3FF; font-weight: 700; }
.metric-def-text { color: inherit; }
</style>
        """, unsafe_allow_html=True)
        for label, definition in METRIC_DEFINITIONS.items():
            st.markdown(
                f'<div class="metric-def-row"><div class="metric-def-name">{label}</div><div class="metric-def-text">{definition}</div></div>',
                unsafe_allow_html=True,
            )

    st.markdown("#### Metric trajectories")
    st.plotly_chart(trajectory_chart(source_metrics, selected_metric_labels, chart_series, chart_series_col, chart_series_label, full_y_axis), use_container_width=True)

    st.markdown("#### Texts")
    original_text = original_text_for(sources, source_id)
    with st.expander("Original text", expanded=True):
        st.text_area("Original", original_text, height=220, label_visibility="collapsed")

    step = st.select_slider(
        "Returned-to-original-language step",
        options=steps,
        value=max(steps),
        key=f"step_{source_id}_{comparison_mode}_{embedding_model}",
    )

    if comparison_mode == "LLM models":
        cols = st.columns(len(selected_models))
        for col, model in zip(cols, selected_models):
            with col:
                st.markdown(f"**{model_short(model)} · {pivot} · step {step}**")
                out = output_text(translations, source_id, model, pivot, int(step))
                if looks_like_failed_output(out, original_text):
                    st.warning("Possible failed translation / end marker.")
                st.text_area(model, out, height=360, label_visibility="collapsed")
    else:
        model = selected_models[0]
        cols = st.columns(min(len(selected_pivots), 3))
        for col, pvt in zip(cols, selected_pivots[:3]):
            with col:
                st.markdown(f"**{pvt} pivot · {model_short(model)} · step {step}**")
                out = output_text(translations, source_id, model, pvt, int(step))
                if looks_like_failed_output(out, original_text):
                    st.warning("Possible failed translation / end marker.")
                st.text_area(f"{pvt} pivot", out, height=360, label_visibility="collapsed")

    st.markdown(f"#### Selected-step metrics · step {step}")
    if comparison_mode == "LLM models":
        tabs = st.tabs([model_short(m) for m in selected_models])
        for tab, model in zip(tabs, selected_models):
            with tab:
                metric_cards(selected_metric_row(metrics, source_id, model, pivot, embedding_model, int(step)))
    else:
        tabs = st.tabs(selected_pivots)
        model = selected_models[0]
        for tab, pvt in zip(tabs, selected_pivots):
            with tab:
                metric_cards(selected_metric_row(metrics, source_id, model, pvt, embedding_model, int(step)))

    st.markdown(f"#### Selected-step comparison chart · step {step}")
    st.plotly_chart(selected_step_bar_chart(source_metrics, selected_metric_labels, chart_series, chart_series_col, chart_series_label, int(step), full_y_axis), use_container_width=True)

    with st.expander("Table: all pivots/models for this source and step", expanded=False):
        table = metrics[(metrics["source_id"] == source_id) & (metrics["step"] == int(step)) & (metrics["embedding_model"] == embedding_model)].copy()
        cols = ["translation_model", "pivot_language", "step"] + list(METRICS.values())
        st.dataframe(table[[c for c in cols if c in table.columns]].sort_values(["pivot_language", "translation_model"]), use_container_width=True, hide_index=True)

    with st.expander("Aggregate distribution across experiments", expanded=False):
        agg_metric = st.selectbox("Aggregate metric", list(METRICS.keys()))
        st.plotly_chart(aggregate_chart(metrics, embedding_model, agg_metric, full_y_axis), use_container_width=True)

    with st.expander("Model metadata", expanded=True):
        all_models_for_source = sorted(source_metrics_all["translation_model"].dropna().unique())
        show_model_metadata_table(all_models_for_source, model_metadata)


if __name__ == "__main__":
    main()