"""Thread-safe per-job blackboard for the Agent pipeline."""

import copy
import json
import os
import threading
from typing import Any, Dict, Iterable, Optional


_COLLECTION_KEYS = {"sheets", "clusters", "findings", "decisions", "rfis",
                    "suppressed"}
_MAPPING_KEYS = {"sheet_index", "jurisdiction", "object_graph"}
_MEMORY_KEYS = _COLLECTION_KEYS | _MAPPING_KEYS


class ProjectMemory:
    """Owns intermediate Agent-mode state and optional debug snapshots."""

    def __init__(self, artifact_dir: Optional[str] = None) -> None:
        self.artifact_dir = artifact_dir
        self._lock = threading.RLock()
        self._data: Dict[str, Any] = {
            **{key: [] for key in _COLLECTION_KEYS},
            **{key: {} for key in _MAPPING_KEYS},
        }
        if artifact_dir:
            os.makedirs(artifact_dir, exist_ok=True)

    def replace(self, key: str, value: Any) -> None:
        """Replace one named memory section."""
        self._validate_key(key)
        with self._lock:
            self._data[key] = copy.deepcopy(value)

    def append(self, key: str, value: Dict[str, Any]) -> None:
        """Append one artifact to a list-backed memory section."""
        if key not in _COLLECTION_KEYS:
            raise KeyError(f"{key!r} is not an appendable memory section")
        with self._lock:
            self._data[key].append(copy.deepcopy(value))

    def extend(self, key: str, values: Iterable[Dict[str, Any]]) -> None:
        """Append several artifacts under one lock."""
        if key not in _COLLECTION_KEYS:
            raise KeyError(f"{key!r} is not an appendable memory section")
        with self._lock:
            self._data[key].extend(copy.deepcopy(list(values)))

    def snapshot(self) -> Dict[str, Any]:
        """Return a detached, JSON-serializable view of current state."""
        with self._lock:
            return copy.deepcopy(self._data)

    def dump(self, filename: str = "memory.json") -> Optional[str]:
        """Persist a snapshot when this job has an artifact directory."""
        if not self.artifact_dir:
            return None
        path = os.path.join(self.artifact_dir, filename)
        temp_path = f"{path}.tmp"
        with open(temp_path, "w", encoding="utf-8") as f:
            json.dump(self.snapshot(), f, indent=2)
        os.replace(temp_path, path)
        return path

    @staticmethod
    def _validate_key(key: str) -> None:
        if key not in _MEMORY_KEYS:
            raise KeyError(f"Unknown project memory section: {key!r}")