"""HF Dataset I/O for the retrieval-tuning loop.

Three files are exchanged between the GitHub Actions tuner and the
chat-logs HF Dataset:

  - turns.jsonl              (chat writes; tuner annotates with judges)
  - retrieval_tuning.json    (tuner writes; chat reads at boot)
  - tuning_log.jsonl         (tuner-only audit trail)

fetch_state() downloads all three to derived/student-chat/, materializing
defaults for any file missing from the dataset (cold start).
push_state() uploads the same three local paths back.

Lazy imports of huggingface_hub via re-exported wrappers so tests can patch
this module's symbols and so this module imports without huggingface_hub
installed (local dev).
"""
from __future__ import annotations
import json
import logging
import shutil
from pathlib import Path

from scripts.lib.chat_eval.tuning import empty_state
from scripts.lib.turns_merge import is_not_found_error


logger = logging.getLogger(__name__)


_DATASET_FILES = ("turns.jsonl", "retrieval_tuning.json", "tuning_log.jsonl")


def hf_hub_download(*args, **kwargs):
    """Re-export so tests can patch this module's reference."""
    from huggingface_hub import hf_hub_download as _real
    return _real(*args, **kwargs)


def HfApi(*args, **kwargs):
    """Re-export so tests can patch this module's reference."""
    from huggingface_hub import HfApi as _real
    return _real(*args, **kwargs)


def fetch_state(*, repo_id: str, token: str, derived_dir: Path) -> None:
    """Download the three dataset files into derived_dir.

    Missing files get sensible defaults so the rest of the loop has real
    files to read:
      - turns.jsonl missing -> empty file (no turns to judge)
      - retrieval_tuning.json missing -> empty_state()
      - tuning_log.jsonl missing -> empty file
    """
    derived_dir.mkdir(parents=True, exist_ok=True)
    for fname in _DATASET_FILES:
        local = derived_dir / fname
        try:
            tmp_path = hf_hub_download(
                repo_id=repo_id,
                filename=fname,
                repo_type="dataset",
                token=token,
            )
            shutil.copyfile(tmp_path, local)
            logger.info("state_io: fetched %s", fname)
        except Exception as e:
            if not is_not_found_error(e):
                raise
            if fname == "retrieval_tuning.json":
                local.write_text(json.dumps(empty_state()), encoding="utf-8")
            else:
                local.write_text("", encoding="utf-8")
            logger.info("state_io: %s missing on remote; wrote default.", fname)


def push_state(*, repo_id: str, token: str, derived_dir: Path,
               commit_message: str) -> None:
    """Upload the three files from derived_dir to the dataset.

    For turns.jsonl: refetch remote and merge by turn_id (local wins on
    conflict; remote-only turns preserved). retrieval_tuning.json and
    tuning_log.jsonl are tuner-only writes and overwrite remote.

    Raises FileNotFoundError if any file is missing locally — we never
    push an empty over a populated remote by accident.
    """
    api = HfApi()
    for fname in _DATASET_FILES:
        local = derived_dir / fname
        if not local.is_file():
            raise FileNotFoundError(
                f"state_io.push_state: missing local file {local}; "
                f"refusing to push."
            )
        staging: Path | None = None
        if fname == "turns.jsonl":
            merged_bytes = _merge_turns_with_remote(
                local, repo_id=repo_id, token=token,
            )
            staging = local.with_suffix(local.suffix + ".merged")
            staging.write_bytes(merged_bytes)
            upload_path = str(staging)
        else:
            upload_path = str(local)
        try:
            api.upload_file(
                path_or_fileobj=upload_path,
                path_in_repo=fname,
                repo_id=repo_id,
                repo_type="dataset",
                token=token,
                commit_message=commit_message,
            )
            logger.info("state_io: pushed %s", fname)
        finally:
            if staging is not None:
                try:
                    staging.unlink()
                except OSError:
                    pass


def _merge_turns_with_remote(local_path: Path, *, repo_id: str,
                             token: str) -> bytes:
    """Tuner-side adapter: read local file, delegate to turns_merge."""
    from scripts.lib.turns_merge import merge_turns_with_remote
    return merge_turns_with_remote(local_path.read_bytes(),
                                   repo_id=repo_id, token=token)