"""Pipeline state management for resumability."""

import atexit
import json
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

from .config import STATE_JSON

logger = logging.getLogger(__name__)

_FLUSH_INTERVAL = 25  # Write to disk every N mutations


class PipelineState:
    """Track per-document, per-step progress for resumable pipeline runs."""

    def __init__(self, path: Path = STATE_JSON):
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._data: Dict[str, Any] = {"documents": {}}
        self._dirty = 0
        if self.path.exists() and self.path.stat().st_size > 0:
            with open(self.path) as f:
                self._data = json.load(f)
            logger.info("Loaded pipeline state with %d documents", len(self._data.get("documents", {})))
        # Flush any pending writes on interpreter exit
        atexit.register(self.flush)

    def _save(self):
        """Coalesced save — only writes to disk every _FLUSH_INTERVAL mutations."""
        self._dirty += 1
        if self._dirty >= _FLUSH_INTERVAL:
            self.flush()

    def flush(self):
        """Force write state to disk (atomic via temp file)."""
        if self._dirty == 0:
            return
        tmp = self.path.with_suffix(".tmp")
        with open(tmp, "w") as f:
            json.dump(self._data, f, indent=2)
        tmp.replace(self.path)
        self._dirty = 0

    def _doc(self, doc_id: str) -> Dict[str, Any]:
        doc_id = str(doc_id)
        if doc_id not in self._data["documents"]:
            self._data["documents"][doc_id] = {}
        return self._data["documents"][doc_id]

    def get(self, doc_id: str, key: str, default: Any = None) -> Any:
        return self._doc(doc_id).get(key, default)

    def set(self, doc_id: str, key: str, value: Any):
        self._doc(doc_id)[key] = value
        self._save()

    def set_batch(self, doc_id: str, updates: Dict[str, Any]):
        doc = self._doc(doc_id)
        doc.update(updates)
        self._save()

    def is_done(self, doc_id: str, step: str) -> bool:
        return self._doc(doc_id).get(step, False) is True

    def mark_done(self, doc_id: str, step: str):
        self.set(doc_id, step, True)

    def all_doc_ids(self) -> List[str]:
        return list(self._data.get("documents", {}).keys())

    def docs_with_step(self, step: str) -> List[str]:
        """Return doc IDs where the given step is True."""
        return [
            doc_id
            for doc_id, info in self._data.get("documents", {}).items()
            if info.get(step) is True
        ]

    def docs_needing_step(self, step: str, prerequisite: Optional[str] = None) -> List[str]:
        """Return doc IDs that need the given step and have completed the prerequisite."""
        result = []
        for doc_id, info in self._data.get("documents", {}).items():
            if info.get(step) is True:
                continue
            if prerequisite and not info.get(prerequisite):
                continue
            result.append(doc_id)
        return result

    @property
    def summary(self) -> Dict[str, int]:
        """Count documents at each stage."""
        counts: Dict[str, int] = {}
        for info in self._data.get("documents", {}).values():
            for key, val in info.items():
                if val is True:
                    counts[key] = counts.get(key, 0) + 1
        return counts
