"""
Annotation tool for Base_UI vs A11yn / feeda11y comparison.

Deployable to Hugging Face Spaces (Docker).

Persistence:
    - Annotations are written locally to results/annotations_test{1,2}.csv.
    - If HF_TOKEN and HF_RESULTS_DATASET env vars are set, results are also
      synced to a Hugging Face Dataset repo on every submit, and pulled from
      there on startup (so the data survives Space restarts).

Environment variables (all optional):
    PORT                Web port (default 7860, HF Spaces standard).
    HF_TOKEN            HF Hub write token (set as a Space secret).
    HF_RESULTS_DATASET  e.g. "your-username/a11yn-annotation-results"
                        Must be created in advance (private or public).
"""

import csv
import hashlib
import json
import os
import random
import threading
from datetime import datetime
from pathlib import Path

from flask import (
    Flask,
    abort,
    jsonify,
    redirect,
    render_template,
    request,
    send_from_directory,
    session,
    url_for,
)

ROOT = Path(__file__).resolve().parent
JSON_PATH = ROOT / "realuirequest300.json"

FOLDERS = {
    "base":     ROOT / "images" / "base",
    "a11yn":    ROOT / "images" / "a11yn",
    "feeda11y": ROOT / "images" / "feeda11y",
}

TESTS = {
    "1": {"name": "Base_UI vs A11yn",    "left_model": "base", "right_model": "a11yn"},
    "2": {"name": "Base_UI vs feeda11y", "left_model": "base", "right_model": "feeda11y"},
}

CHOICES = ["strong_left", "weak_left", "similar", "weak_right", "strong_right"]

OUT_DIR = ROOT / "results"
OUT_DIR.mkdir(exist_ok=True)
CSV_HEADERS = [
    "timestamp", "annotator", "test_id", "test_name",
    "index", "request_id", "left_model", "right_model", "choice",
]

with open(JSON_PATH) as f:
    REQUESTS = json.load(f)

# ---------- Optional Hugging Face Hub persistence ----------
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_RESULTS_DATASET = os.environ.get("HF_RESULTS_DATASET")
_hf_lock = threading.Lock()
_hf_api = None

def _hf_enabled() -> bool:
    return bool(HF_TOKEN and HF_RESULTS_DATASET and _hf_api is not None)


def _init_hf():
    """Try to import huggingface_hub and pull existing CSVs into ./results."""
    global _hf_api
    if not (HF_TOKEN and HF_RESULTS_DATASET):
        print("[hf] HF_TOKEN or HF_RESULTS_DATASET not set — local CSV only.")
        return
    try:
        from huggingface_hub import HfApi
    except ImportError:
        print("[hf] huggingface_hub not installed — skipping hub sync.")
        return
    _hf_api = HfApi(token=HF_TOKEN)
    # Make sure the dataset repo exists (no-op if it already does).
    try:
        _hf_api.create_repo(
            repo_id=HF_RESULTS_DATASET, repo_type="dataset", exist_ok=True, private=False,
        )
    except Exception as e:
        print(f"[hf] create_repo warning: {e}")
    # Pull existing CSVs (if any) so we resume with previously saved data.
    for tid in TESTS:
        fn = f"annotations_test{tid}.csv"
        try:
            local = _hf_api.hf_hub_download(
                repo_id=HF_RESULTS_DATASET,
                filename=fn,
                repo_type="dataset",
            )
            target = OUT_DIR / fn
            Path(local).replace(target) if False else _copy(local, target)
            print(f"[hf] pulled existing {fn}")
        except Exception:
            pass  # not present yet — fine


def _copy(src, dst):
    with open(src, "rb") as fr, open(dst, "wb") as fw:
        fw.write(fr.read())


def _push_csv_to_hub(test_id: str):
    """Upload the current CSV for `test_id` to the HF dataset repo."""
    if not _hf_enabled():
        return
    fn = f"annotations_test{test_id}.csv"
    local = OUT_DIR / fn
    if not local.exists():
        return
    try:
        _hf_api.upload_file(
            path_or_fileobj=str(local),
            path_in_repo=fn,
            repo_id=HF_RESULTS_DATASET,
            repo_type="dataset",
            commit_message=f"update {fn}",
        )
    except Exception as e:
        print(f"[hf] upload failed for {fn}: {e}")


# ---------- App ----------
app = Flask(__name__)
app.secret_key = os.environ.get("SECRET_KEY") or os.urandom(24).hex()

ACCESS_CODE = os.environ.get("ACCESS_CODE")  # if set, landing page requires it


def _has_access() -> bool:
    """True if no gate configured, or session has been unlocked."""
    return (not ACCESS_CODE) or session.get("unlocked") is True


def csv_path(test_id: str) -> Path:
    return OUT_DIR / f"annotations_test{test_id}.csv"


def ensure_csv(test_id: str) -> None:
    p = csv_path(test_id)
    if not p.exists():
        with open(p, "w", newline="") as f:
            csv.writer(f).writerow(CSV_HEADERS)


def get_pair_mapping(annotator: str, test_id: str, index: int) -> dict:
    test = TESTS[test_id]
    seed = hashlib.md5(f"{annotator}|{test_id}|{index}".encode()).hexdigest()
    rng = random.Random(seed)
    models = [test["left_model"], test["right_model"]]
    rng.shuffle(models)
    return {"left_model": models[0], "right_model": models[1]}


def load_existing(annotator: str, test_id: str) -> dict:
    p = csv_path(test_id)
    if not p.exists():
        return {}
    existing = {}
    with open(p, newline="") as f:
        for row in csv.DictReader(f):
            if row["annotator"] == annotator:
                existing[int(row["index"])] = row["choice"]
    return existing


@app.route("/")
def index():
    return render_template(
        "index.html",
        tests=TESTS,
        needs_code=bool(ACCESS_CODE),
        unlocked=_has_access(),
        error=request.args.get("error"),
    )


@app.route("/start", methods=["POST"])
def start():
    if ACCESS_CODE and not session.get("unlocked"):
        code = (request.form.get("access_code") or "").strip()
        if code != ACCESS_CODE:
            return redirect(url_for("index", error="wrong_code"))
        session["unlocked"] = True
    annotator = (request.form.get("annotator") or "").strip()
    test_id = request.form.get("test_id", "1")
    if not annotator or test_id not in TESTS:
        return redirect(url_for("index"))
    existing = load_existing(annotator, test_id)
    next_idx = next(
        (i for i in range(1, len(REQUESTS) + 1) if i not in existing), len(REQUESTS)
    )
    return redirect(
        url_for("annotate", test_id=test_id, index=next_idx, annotator=annotator)
    )


@app.route("/annotate/<test_id>/<int:index>")
def annotate(test_id, index):
    if not _has_access():
        return redirect(url_for("index"))
    annotator = (request.args.get("annotator") or "").strip()
    if not annotator or test_id not in TESTS:
        return redirect(url_for("index"))
    if index < 1 or index > len(REQUESTS):
        return redirect(url_for("index"))

    item = REQUESTS[index - 1]
    mapping = get_pair_mapping(annotator, test_id, index)
    filename = f"{index:03d}_{item['request_id']}.png"

    existing = load_existing(annotator, test_id)
    current_choice = existing.get(index)
    done_count = len(existing)

    return render_template(
        "annotate.html",
        test=TESTS[test_id],
        test_id=test_id,
        index=index,
        total=len(REQUESTS),
        annotator=annotator,
        item=item,
        mapping=mapping,
        filename=filename,
        choices=CHOICES,
        current_choice=current_choice,
        done_count=done_count,
        existing_indices=sorted(existing.keys()),
    )


@app.route("/submit", methods=["POST"])
def submit():
    if not _has_access():
        return jsonify({"ok": False, "error": "unauthorized"}), 401
    data = request.get_json(force=True)
    annotator = (data.get("annotator") or "").strip()
    test_id = str(data.get("test_id"))
    index = int(data.get("index"))
    choice = data.get("choice")

    if (
        not annotator
        or test_id not in TESTS
        or choice not in CHOICES
        or not (1 <= index <= len(REQUESTS))
    ):
        return jsonify({"ok": False, "error": "invalid input"}), 400

    item = REQUESTS[index - 1]
    mapping = get_pair_mapping(annotator, test_id, index)

    # Serialize CSV mutations to avoid concurrent-write corruption on the Space.
    with _hf_lock:
        ensure_csv(test_id)
        p = csv_path(test_id)
        with open(p, newline="") as f:
            rows = list(csv.DictReader(f))
        rows = [r for r in rows
                if not (r["annotator"] == annotator and int(r["index"]) == index)]
        rows.append({
            "timestamp": datetime.utcnow().isoformat(timespec="seconds") + "Z",
            "annotator": annotator,
            "test_id": test_id,
            "test_name": TESTS[test_id]["name"],
            "index": index,
            "request_id": item["request_id"],
            "left_model": mapping["left_model"],
            "right_model": mapping["right_model"],
            "choice": choice,
        })
        with open(p, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=CSV_HEADERS)
            w.writeheader()
            w.writerows(rows)
        _push_csv_to_hub(test_id)

    return jsonify({"ok": True})


@app.route("/image/<model>/<path:filename>")
def image(model, filename):
    if not _has_access():
        abort(401)
    if model not in FOLDERS:
        abort(404)
    if "/" in filename or ".." in filename:
        abort(400)
    folder = FOLDERS[model]
    full = folder / filename
    if not full.exists():
        abort(404)
    # Cache aggressively — image files are immutable per filename.
    resp = send_from_directory(folder, filename)
    resp.headers["Cache-Control"] = "public, max-age=86400"
    return resp


@app.route("/results/<test_id>.csv")
def download_results(test_id):
    """Download the current CSV (gated by ACCESS_CODE if set)."""
    if not _has_access():
        abort(401)
    if test_id not in TESTS:
        abort(404)
    p = csv_path(test_id)
    if not p.exists():
        # return empty CSV with headers
        return ",".join(CSV_HEADERS) + "\n", 200, {"Content-Type": "text/csv"}
    return send_from_directory(OUT_DIR, p.name, as_attachment=True)


_init_hf()

if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    app.run(host="0.0.0.0", port=port)