"""S2ST Subjective Test (MOS).

Anonymized listening test for speech-to-speech translation. Each sample shows
nine model outputs as `Model A` ... `Model I` in a per-sample randomized order
(rater-specific). Raters score three axes (1-5): translation quality, audio
naturalness, speaker similarity, plus an optional note.

Submissions are appended to mos_results.csv. NOTE: HF Spaces filesystem is
ephemeral by default - download results periodically from the Admin panel.
"""

from __future__ import annotations

import csv
import json
import os
import random
import string
from datetime import datetime
from pathlib import Path

# HF Spaces sets GRADIO_SSR_MODE=True via the runtime; SSR + the HF proxy currently
# generate broken file URLs like /gradio_a/gradio_api/file=... that 404. Force SSR
# off before gradio is imported.
os.environ["GRADIO_SSR_MODE"] = "False"

# Workaround for a long-standing gradio_client bug where the JSON-schema -> Python-type
# walker crashes on schemas where `additionalProperties: True` (boolean) instead of a
# nested object schema. The frontend hits /info on page load, so unpatched the UI fails
# to mount and audio never loads. Patch BEFORE importing gradio.
import gradio_client.utils as _gcu  # noqa: E402

_orig_json_schema = _gcu._json_schema_to_python_type
_orig_get_type = _gcu.get_type


def _safe_get_type(schema):
    if isinstance(schema, bool):
        return "Any"
    return _orig_get_type(schema)


def _safe_json_schema(schema, defs=None):
    if isinstance(schema, bool):
        return "Any"
    return _orig_json_schema(schema, defs)


_gcu.get_type = _safe_get_type
_gcu._json_schema_to_python_type = _safe_json_schema

import gradio as gr  # noqa: E402
import gradio.processing_utils as _gpu  # noqa: E402

# HF Spaces' reverse proxy makes Gradio compute a root_url like
# "https://<space>.hf.space/gradio_a", which then gets prepended to relative file
# URLs as "/gradio_a/gradio_api/file=...". Strip that bogus suffix.
_orig_add_root_url = _gpu.add_root_url


def _fixed_add_root_url(data, root_url, previous_root_url):
    if isinstance(root_url, str) and root_url.endswith("/gradio_a"):
        print(f"[patch] stripping /gradio_a suffix from root_url={root_url!r}", flush=True)
        root_url = root_url[: -len("/gradio_a")]
    return _orig_add_root_url(data, root_url, previous_root_url)


_gpu.add_root_url = _fixed_add_root_url

ROOT = Path(__file__).parent
SAMPLES = json.loads((ROOT / "samples.json").read_text(encoding="utf-8"))
RESULTS_FILE = ROOT / "mos_results.csv"
ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "")  # set in Space Secrets

# When deployed on HF Spaces the platform's reverse proxy corrupts the URLs Gradio
# auto-constructs for cached files (they come out as /gradio_a/gradio_api/file=...).
# Build a complete public URL ourselves; Gradio treats HTTP URLs as opaque and
# passes them straight through to the frontend.
SPACE_HOST = os.environ.get("SPACE_HOST", "").strip()


def audio_url(filename: str) -> str:
    local = ROOT / "audio" / filename
    if SPACE_HOST:
        return f"https://{SPACE_HOST}/gradio_api/file={local}"
    return str(local)

DIRECTIONS = ["en2zh", "zh2en"]
DIRECTION_LABELS = {"en2zh": "English -> Chinese", "zh2en": "Chinese -> English"}
SRC_LANG_LABEL = {"en2zh": "English", "zh2en": "Chinese"}
TGT_LANG_LABEL = {"en2zh": "Chinese", "zh2en": "English"}
NUM_MODELS = 8
LETTERS = list(string.ascii_uppercase[:NUM_MODELS])  # A..H

# CSV header
CSV_HEADER = [
    "timestamp", "user_id", "direction", "sample_idx", "sample_id",
    # real model per letter (so we can de-anonymize later)
    *[f"real_{L}" for L in LETTERS],
    # ratings per letter (each axis 1-5)
    *[f"{L}_translation" for L in LETTERS],
    *[f"{L}_naturalness" for L in LETTERS],
    *[f"{L}_spksim" for L in LETTERS],
    *[f"{L}_note" for L in LETTERS],
]


def ensure_csv():
    if not RESULTS_FILE.exists():
        with open(RESULTS_FILE, "w", encoding="utf-8", newline="") as f:
            csv.writer(f).writerow(CSV_HEADER)


def sample_at(direction: str, idx: int):
    return SAMPLES[direction]["samples"][idx]


def total_samples(direction: str) -> int:
    return len(SAMPLES[direction]["samples"])


def build_letter_assignment(rater_id: str, direction: str, sample_idx_zero_based: int):
    """Return a list mapping displayed letter index (0=A..8=I) -> real model key.
    Deterministic per (rater, direction, sample_idx) so navigation is stable
    within a session.
    """
    real_models = SAMPLES[direction]["model_order"]
    seed = f"{rater_id or 'anon'}::{direction}::{sample_idx_zero_based}"
    rng = random.Random(seed)
    order = list(real_models)
    rng.shuffle(order)
    return order


def render_sample(rater_id: str, direction: str, idx_zero: int):
    """Build all UI-bound values for the requested sample."""
    s = sample_at(direction, idx_zero)
    real_per_letter = build_letter_assignment(rater_id, direction, idx_zero)
    src_audio_path = audio_url(s["source_audio"])

    audio_paths = []
    for letter_idx in range(NUM_MODELS):
        real = real_per_letter[letter_idx]
        rel = s["models"][real]["audio"]
        audio_paths.append(audio_url(rel))

    progress = f"Sample {idx_zero + 1} / {total_samples(direction)} ({DIRECTION_LABELS[direction]})"
    src_header = f"Source ({SRC_LANG_LABEL[direction]})"
    return src_audio_path, s["source_text"], src_header, progress, real_per_letter, audio_paths


def reset_widget_outputs():
    """Return default values for the 9*(audio+3 sliders+1 note) widget set."""
    out = []
    for _ in range(NUM_MODELS):
        out.append(None)         # audio
        out.append(3)            # translation
        out.append(3)            # naturalness
        out.append(3)            # spksim
        out.append("")           # note
    return out


def load_existing_ratings(rater_id: str, direction: str, sample_id: str, real_per_letter):
    """If this rater has already rated this (direction, sample) — return previous values."""
    if not RESULTS_FILE.exists():
        return None
    try:
        with open(RESULTS_FILE, "r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            last = None
            for row in reader:
                if (row.get("user_id") == (rater_id or "anonymous")
                        and row.get("direction") == direction
                        and row.get("sample_id") == sample_id):
                    last = row
            if last is None:
                return None
        # Build a real_model -> ratings dict from the saved row
        per_real = {}
        for L in LETTERS:
            real = last.get(f"real_{L}")
            if not real:
                continue
            per_real[real] = {
                "translation": int(last.get(f"{L}_translation") or 3),
                "naturalness": int(last.get(f"{L}_naturalness") or 3),
                "spksim": int(last.get(f"{L}_spksim") or 3),
                "note": last.get(f"{L}_note") or "",
            }
        # Project onto current letter assignment for the rater
        proj = []
        for letter_idx in range(NUM_MODELS):
            real = real_per_letter[letter_idx]
            r = per_real.get(real, {"translation": 3, "naturalness": 3, "spksim": 3, "note": ""})
            proj.extend([r["translation"], r["naturalness"], r["spksim"], r["note"]])
        return proj
    except Exception:
        return None


# ---------------------------- Gradio callbacks -------------------------------

def on_load_sample(rater_id, direction, idx_zero):
    rater_id = (rater_id or "").strip()
    n = total_samples(direction)
    idx_zero = max(0, min(idx_zero, n - 1))
    src_audio, src_text, src_header, progress, real_per_letter, audios = render_sample(
        rater_id, direction, idx_zero
    )
    # Existing ratings (if any) for this rater/direction/sample
    s = sample_at(direction, idx_zero)
    existing = load_existing_ratings(rater_id, direction, s["id"], real_per_letter)

    # Build all output widget values
    widget_vals = []
    for letter_idx in range(NUM_MODELS):
        widget_vals.append(audios[letter_idx])  # audio
        if existing is None:
            widget_vals.extend([3, 3, 3, ""])
        else:
            widget_vals.extend(existing[letter_idx * 4: letter_idx * 4 + 4])

    nav_prev_interactive = idx_zero > 0
    nav_next_label = "Submit & Next ->" if idx_zero < n - 1 else "Submit & Finish"
    msg = ""
    if existing is not None:
        msg = "Loaded your previously submitted ratings for this sample."

    return (
        src_audio, src_text, src_header, progress, idx_zero, real_per_letter, msg,
        gr.update(interactive=nav_prev_interactive),
        gr.update(value=nav_next_label),
        *widget_vals,
    )


def on_submit(rater_id, direction, idx_zero, real_per_letter, *vals):
    """vals: per letter (translation, naturalness, spksim, note) repeated NUM_MODELS times."""
    rater_id = (rater_id or "").strip() or "anonymous"
    ensure_csv()
    s = sample_at(direction, idx_zero)
    row = {
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "user_id": rater_id,
        "direction": direction,
        "sample_idx": s["sample_idx"],
        "sample_id": s["id"],
    }
    for letter_idx, L in enumerate(LETTERS):
        row[f"real_{L}"] = real_per_letter[letter_idx]
    for letter_idx, L in enumerate(LETTERS):
        base = letter_idx * 4
        row[f"{L}_translation"] = vals[base]
        row[f"{L}_naturalness"] = vals[base + 1]
        row[f"{L}_spksim"] = vals[base + 2]
        row[f"{L}_note"] = vals[base + 3]

    # Overwrite if same (rater, direction, sample) submitted before; otherwise append.
    rows = []
    if RESULTS_FILE.exists():
        with open(RESULTS_FILE, "r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            rows = [r for r in reader
                    if not (r.get("user_id") == rater_id
                            and r.get("direction") == direction
                            and r.get("sample_id") == s["id"])]
    rows.append(row)
    with open(RESULTS_FILE, "w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=CSV_HEADER)
        w.writeheader()
        w.writerows(rows)

    # Advance
    n = total_samples(direction)
    next_idx = min(idx_zero + 1, n - 1)
    return next_idx, "Saved. " + (f"Moving to sample {next_idx + 1}." if next_idx > idx_zero
                                  else "This was the last sample for this direction.")


def on_prev(idx_zero):
    return max(0, idx_zero - 1)


def on_direction_change(_direction):
    # Resetting to sample 0 of the new direction
    return 0


def admin_download(password):
    if not RESULTS_FILE.exists():
        return None, "No submissions yet."
    if ADMIN_PASSWORD and password != ADMIN_PASSWORD:
        return None, "Wrong password."
    return str(RESULTS_FILE), f"OK. {RESULTS_FILE.stat().st_size} bytes."


def admin_clear(password):
    if ADMIN_PASSWORD and password != ADMIN_PASSWORD:
        return "Wrong password."
    if RESULTS_FILE.exists():
        RESULTS_FILE.unlink()
    return "mos_results.csv cleared."


# ----------------------------------- UI --------------------------------------

CSS = """
.model-card {border: 1px solid #d0d0d0; border-radius: 8px; padding: 10px;
            margin-bottom: 8px; background: #fbfbfb;}
.model-card h3 {margin: 0 0 6px 0; font-size: 16px;}
.muted {color: #666; font-size: 12px;}
"""

gr.set_static_paths(paths=[str(ROOT / "audio")])

with gr.Blocks(title="S2ST Subjective Test", css=CSS, analytics_enabled=False) as demo:
    gr.Markdown("## Speech-to-Speech Translation - Subjective Listening Test")
    gr.Markdown(
        f"Rate each anonymous model (A through {LETTERS[-1]}) on three axes "
        "(1=worst, 5=best): **Translation Quality**, **Audio Naturalness**, "
        "**Speaker Similarity**. Models are presented in a different random order "
        "on every sample, so 'Model A' on this page is NOT the same model as "
        "'Model A' on the next page."
    )

    with gr.Row():
        rater_id = gr.Textbox(label="Rater ID (your name or email)", scale=2)
        direction = gr.Radio(
            choices=[(DIRECTION_LABELS[d], d) for d in DIRECTIONS],
            value="en2zh",
            label="Translation direction",
            scale=2,
        )

    idx_state = gr.State(0)
    real_per_letter_state = gr.State([])
    progress_md = gr.Markdown()
    status_md = gr.Markdown()

    src_header_md = gr.Markdown()
    with gr.Row():
        src_text = gr.Textbox(label="Source text", interactive=False, lines=2, scale=3)
        src_audio = gr.Audio(label="Source audio", interactive=False, scale=2)

    gr.Markdown("### Model outputs (anonymized)")

    audio_widgets = []
    translation_widgets = []
    naturalness_widgets = []
    spksim_widgets = []
    note_widgets = []
    for i, L in enumerate(LETTERS):
        with gr.Group(elem_classes="model-card"):
            gr.Markdown(f"### Model {L}")
            with gr.Row():
                a = gr.Audio(label=f"Model {L} audio", interactive=False, scale=2)
                with gr.Column(scale=3):
                    t = gr.Slider(1, 5, value=3, step=1, label=f"Translation Quality (Model {L})")
                    nat = gr.Slider(1, 5, value=3, step=1, label=f"Audio Naturalness (Model {L})")
                    spk = gr.Slider(1, 5, value=3, step=1, label=f"Speaker Similarity (Model {L})")
                    note = gr.Textbox(label=f"Note (Model {L}, optional)", lines=1)
            audio_widgets.append(a)
            translation_widgets.append(t)
            naturalness_widgets.append(nat)
            spksim_widgets.append(spk)
            note_widgets.append(note)

    with gr.Row():
        prev_btn = gr.Button("<- Previous", variant="secondary", interactive=False)
        submit_btn = gr.Button("Submit & Next ->", variant="primary")

    with gr.Accordion("Admin", open=False):
        admin_pwd = gr.Textbox(label="Admin password", type="password")
        with gr.Row():
            download_btn = gr.Button("Refresh & download mos_results.csv")
            clear_btn = gr.Button("Clear all results", variant="stop")
        download_file = gr.File(label="Latest mos_results.csv")
        admin_status = gr.Markdown()

    # Flat list of all per-letter widgets in the same order as on_load_sample returns
    per_letter_widgets = []
    for i in range(NUM_MODELS):
        per_letter_widgets.extend([
            audio_widgets[i],
            translation_widgets[i],
            naturalness_widgets[i],
            spksim_widgets[i],
            note_widgets[i],
        ])

    load_outputs = [
        src_audio, src_text, src_header_md, progress_md, idx_state, real_per_letter_state, status_md,
        prev_btn, submit_btn,
        *per_letter_widgets,
    ]

    # Triggers that re-render the current sample
    rater_id.change(on_load_sample, [rater_id, direction, idx_state], load_outputs)
    direction.change(on_direction_change, [direction], [idx_state]).then(
        on_load_sample, [rater_id, direction, idx_state], load_outputs
    )
    idx_state.change(on_load_sample, [rater_id, direction, idx_state], load_outputs)
    demo.load(on_load_sample, [rater_id, direction, idx_state], load_outputs)

    # Submit: collect per-letter ratings, write CSV, advance idx
    rating_inputs = []
    for i in range(NUM_MODELS):
        rating_inputs.extend([
            translation_widgets[i],
            naturalness_widgets[i],
            spksim_widgets[i],
            note_widgets[i],
        ])

    submit_btn.click(
        on_submit,
        [rater_id, direction, idx_state, real_per_letter_state, *rating_inputs],
        [idx_state, status_md],
    )
    prev_btn.click(on_prev, [idx_state], [idx_state])

    download_btn.click(admin_download, [admin_pwd], [download_file, admin_status])
    clear_btn.click(admin_clear, [admin_pwd], [admin_status])


ensure_csv()
demo.queue(default_concurrency_limit=4)

if __name__ == "__main__":
    demo.launch(allowed_paths=[str(ROOT / "audio")], show_api=False, ssr_mode=False)