"""Resumable evaluation: checkpoint plugin.

Run-level checkpointing: every completed sample is appended to a jsonl
checkpoint file; on restart (crash/OOM/service down), completed samples are
restored and ONLY the missing ones are generated. The report is then
assembled from restored + fresh predictions.

    from evalharness.eval.checkpoint import CheckpointStore
    store = CheckpointStore(path='run.jsonl', model=model_spec, dataset=name)
    done = store.load()                    # {sample_key: prediction-dict}
    ... generate only missing ...
    store.append(sample_key, pred)         # after EACH sample (crash-safe)
"""

import json
import os
import time
from pathlib import Path
from typing import Any, Dict, Optional


class CheckpointStore:
    """Append-only jsonl checkpoint; keyed by stable sample identity."""

    def __init__(self, path: str, model: str = '', dataset: str = ''):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.model = model
        self.dataset = dataset
        self._entries: Dict[str, Dict[str, Any]] = {}
        self._fh = None

    @staticmethod
    def key_for(sample, idx: int) -> str:
        """Stable per-sample key: prefer dataset-native ids, fall back to
        a hash of the input text (survives re-orderings)."""
        native = (sample.metadata or {}).get('task_id') \
            or (sample.metadata or {}).get('id') \
            or (sample.metadata or {}).get('instance_id')
        if native:
            return str(native)
        import hashlib

        h = hashlib.md5((sample.input_text or '').encode('utf-8')).hexdigest()[:16]
        return f'{idx}:{h}'

    def load(self) -> Dict[str, Dict[str, Any]]:
        """Read all checkpointed predictions (idempotent)."""
        self._entries = {}
        if not self.path.exists():
            return self._entries
        with open(self.path, encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    rec = json.loads(line)
                    self._entries[rec['key']] = rec.get('pred', {})
                except (ValueError, KeyError):
                    continue  # torn tail line from a crash -- safe to skip
        return self._entries

    def append(self, key: str, pred: Dict[str, Any]) -> None:
        """Persist one prediction immediately (fsync-free append is fine:
        worst case loses the last in-flight sample on crash)."""
        rec = {'key': key, 'ts': time.time(), 'pred': pred}
        with open(self.path, 'a', encoding='utf-8') as f:
            f.write(json.dumps(rec, ensure_ascii=False) + '\n')
        self._entries[key] = pred

    def __len__(self) -> int:
        return len(self._entries)

    def summary(self) -> Dict[str, Any]:
        return {'path': str(self.path), 'restored': len(self._entries),
                'size_kb': self.path.stat().st_size // 1024 if self.path.exists() else 0}


def checkpoint_path(root: str, dataset: str, model: str, tag: str = '') -> str:
    """Deterministic path per (dataset, model[, tag]) so re-runs resume."""
    import hashlib

    h = hashlib.md5(f'{dataset}|{model}|{tag}'.encode()).hexdigest()[:10]
    return os.path.join(root, 'ckpt', f'{dataset}-{h}.jsonl')