# -*- coding: utf-8 -*-
"""Lightweight task-scoped checkpoint persistence for SmartPaper executors.

This module intentionally depends only on the Python standard library. It can
be used immediately after a Worker claim without importing DeepSeek, RAG,
statistics, document, or Streamlit modules.
"""
from __future__ import annotations

import json
import os
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


def _atomic_json(path: Path, value: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False), encoding="utf-8")
    os.replace(temporary, path)


class TaskCheckpointStore:
    def __init__(self, task_dir: Path, request_hash: str, execution_epoch: int):
        self.task_dir = Path(task_dir)
        self.request_hash = str(request_hash)
        self.execution_epoch = int(execution_epoch)
        self.directory = self.task_dir / "checkpoints"
        self.directory.mkdir(parents=True, exist_ok=True)
        self.sequence = 0

    def load_latest(self) -> dict[str, Any] | None:
        candidates: list[tuple[int, Path, dict[str, Any]]] = []
        for path in self.directory.glob("checkpoint_*.json"):
            try:
                payload = json.loads(path.read_text(encoding="utf-8"))
                if not isinstance(payload, dict):
                    continue
                if payload.get("request_hash") != self.request_hash:
                    continue
                seq = int(payload.get("checkpoint_seq") or 0)
                if seq <= 0:
                    continue
                candidates.append((seq, path, payload))
            except Exception:
                continue
        if not candidates:
            return None
        seq, _, payload = max(candidates, key=lambda item: item[0])
        self.sequence = seq
        return payload

    def save(self, stage: str, state: dict[str, Any]) -> tuple[int, str]:
        self.sequence += 1
        stage_name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(stage or "checkpoint"))[:80]
        payload = dict(state or {})
        payload.update(
            {
                "request_hash": self.request_hash,
                "checkpoint_seq": self.sequence,
                "checkpoint_stage": stage_name,
                "checkpoint_execution_epoch": self.execution_epoch,
                "checkpoint_saved_at": datetime.now(timezone.utc).isoformat(),
            }
        )
        path = self.directory / f"checkpoint_{self.sequence:05d}_{stage_name}.json"
        _atomic_json(path, payload)
        _atomic_json(self.directory / "latest.json", payload)
        return self.sequence, stage_name