"""Persistence for the arena: battle logs + Elo state, stored in a
Hugging Face Storage Bucket.

Everything is written under `hf://buckets/<BUCKET_ID>/`:

    state/elo.json                          current ratings (rewritten per vote)
    generations/<arena>/<date>/<id>.json    text, settings and both models
    votes/<arena>/<date>/<id>.json          the vote for that battle + Elo deltas

The bucket holds the logs and, under models/, the weights of every contender; every prompt and vote is readable by anyone who can read the bucket.
Writes go through a single background worker so the UI never blocks on network
I/O and so state/elo.json is never written by two threads at once.

The HF token is read from the HF_TOKEN environment variable (a Space secret).
It is only ever handed to huggingface_hub and is never logged or rendered.
"""

import json
import os
import queue
import tempfile
import threading
import traceback
from datetime import datetime, timezone

BUCKET_ID = "Banaxi-Tech/tts-arena"
STATE_PATH = "state/elo.json"
_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

BUCKET_URL = f"https://huggingface.co/buckets/{BUCKET_ID}"

_LOCAL_MIRROR = ""  # dev fallback: set to a folder to also write logs locally


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds")


def _api():
    from huggingface_hub import HfApi

    return HfApi(token=_TOKEN)


# ---------------------------------------------------------------- writing
def _write_bytes(path: str, data: bytes) -> None:
    """Upload one object to the bucket. Raises on failure."""
    api = _api()
    if hasattr(api, "batch_bucket_files"):
        api.batch_bucket_files(BUCKET_ID, add=[(data, path)])
        return
    # Older huggingface_hub: fall back to the fsspec interface.
    from huggingface_hub import HfFileSystem

    fs = HfFileSystem(token=_TOKEN)
    with fs.open(f"hf://buckets/{BUCKET_ID}/{path}", "wb") as f:
        f.write(data)


def _read_bytes(path: str) -> bytes:
    """Download one object from the bucket. Raises if it can't be read."""
    from huggingface_hub import HfFileSystem

    fs = HfFileSystem(token=_TOKEN)
    uri = f"hf://buckets/{BUCKET_ID}/{path}"
    if not fs.exists(uri):
        raise FileNotFoundError(uri)
    with fs.open(uri, "rb") as f:
        return f.read()


def list_paths(prefix: str) -> list[str]:
    """Every file under `prefix` in the bucket."""
    api = _api()
    return [
        item.path
        for item in api.list_bucket_tree(BUCKET_ID, prefix=prefix, recursive=True)
        if item.type == "file"
    ]


def read_json_many(paths: list[str]) -> list[dict]:
    """Download and parse many JSON objects in one batched call."""
    if not paths:
        return []
    api = _api()
    out: list[dict] = []
    with tempfile.TemporaryDirectory() as tmp:
        targets = [(p, os.path.join(tmp, p.replace("/", "__"))) for p in paths]
        api.download_bucket_files(BUCKET_ID, files=targets)
        for _, local in targets:
            try:
                with open(local, "rb") as f:
                    out.append(json.loads(f.read()))
            except Exception:  # one corrupt record must not sink the replay
                traceback.print_exc()
    return out


def load_vote_log() -> list[dict]:
    """Every vote ever cast, oldest first. This is the authoritative history."""
    records = read_json_many(list_paths("votes/"))
    records.sort(key=lambda r: r.get("created_at", ""))
    return records


def _mirror(path: str, data: bytes) -> None:
    if not _LOCAL_MIRROR:
        return
    full = os.path.join(_LOCAL_MIRROR, path)
    os.makedirs(os.path.dirname(full), exist_ok=True)
    with open(full, "wb") as f:
        f.write(data)


# ---------------------------------------------------------------- worker
_q: "queue.Queue[tuple[str, bytes] | None]" = queue.Queue()
_worker_started = False
_worker_lock = threading.Lock()
last_error: str | None = None


def _worker() -> None:
    global last_error
    while True:
        item = _q.get()
        if item is None:
            return
        path, data = item
        try:
            _mirror(path, data)
            if _TOKEN:
                _write_bytes(path, data)
            else:
                last_error = "HF_TOKEN is not set — logs are not being uploaded."
        except Exception as exc:  # keep the arena usable if the bucket is down
            last_error = f"{type(exc).__name__}: {exc}"
            traceback.print_exc()
        finally:
            _q.task_done()


def _ensure_worker() -> None:
    global _worker_started
    with _worker_lock:
        if not _worker_started:
            threading.Thread(target=_worker, name="arena-uploader", daemon=True).start()
            _worker_started = True


def enqueue_json(path: str, payload: dict) -> None:
    """Queue one JSON document for upload. Never raises."""
    _ensure_worker()
    data = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
    _q.put((path, data))


def flush(timeout: float = 30.0) -> None:
    """Block until queued uploads are done (used by tests / shutdown)."""
    _ensure_worker()
    done = threading.Event()
    threading.Thread(target=lambda: (_q.join(), done.set()), daemon=True).start()
    done.wait(timeout)


# ---------------------------------------------------------------- records
def _day() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def generation_path(arena: str, battle_id: str) -> str:
    return f"generations/{arena}/{_day()}/{battle_id}.json"


def vote_path(arena: str, battle_id: str) -> str:
    return f"votes/{arena}/{_day()}/{battle_id}.json"


def log_generation(record: dict) -> None:
    enqueue_json(generation_path(record["arena"], record["battle_id"]), record)


def log_vote(record: dict) -> None:
    enqueue_json(vote_path(record["arena"], record["battle_id"]), record)


def save_state(state: dict) -> None:
    state = dict(state, updated_at=utc_now())
    enqueue_json(STATE_PATH, state)


def load_state() -> dict | None:
    """The cached ratings blob. Returns None only if it genuinely isn't there."""
    try:
        return json.loads(_read_bytes(STATE_PATH))
    except FileNotFoundError:
        return None


def token_configured() -> bool:
    return bool(_TOKEN)