"""Incremental index state persistence for vector KB builds."""

from __future__ import annotations

import hashlib
import json
import time
from pathlib import Path
from typing import Any, Dict, Iterable, Optional


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def sha256_file(path: Path) -> str:
    hasher = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            hasher.update(chunk)
    return hasher.hexdigest()


class IndexStateStore:
    """Read/write a compact JSON state file for incremental builds."""

    def __init__(self, path: str | Path, *, schema_version: str = "1") -> None:
        self.path = Path(path)
        self.schema_version = schema_version
        self._state: Dict[str, Any] = self._default_state()
        self.load()

    def _default_state(self) -> Dict[str, Any]:
        return {
            "schema_version": self.schema_version,
            "updated_at": "",
            "docs": {},
        }

    @property
    def docs(self) -> Dict[str, Dict[str, Any]]:
        docs = self._state.get("docs")
        if not isinstance(docs, dict):
            docs = {}
            self._state["docs"] = docs
        return docs

    def load(self) -> None:
        if not self.path.exists():
            return
        try:
            data = json.loads(self.path.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                self._state = data
                if "docs" not in self._state or not isinstance(self._state["docs"], dict):
                    self._state["docs"] = {}
        except Exception:
            # keep default on corrupted state file
            self._state = self._default_state()

    def save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._state["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%S")
        self.path.write_text(json.dumps(self._state, ensure_ascii=False, indent=2), encoding="utf-8")

    def get_doc(self, doc_id: str) -> Optional[Dict[str, Any]]:
        value = self.docs.get(doc_id)
        return dict(value) if isinstance(value, dict) else None

    def should_skip_doc(self, doc_id: str, file_hash: str) -> bool:
        doc_state = self.get_doc(doc_id)
        if not doc_state:
            return False
        return doc_state.get("file_hash") == file_hash

    def get_chunk_count(self, doc_id: str) -> int:
        doc_state = self.get_doc(doc_id)
        if not doc_state:
            return 0
        chunk_ids = doc_state.get("chunk_ids") or []
        if isinstance(chunk_ids, list):
            return len(chunk_ids)
        return 0

    def get_chunk_ids(self, doc_id: str) -> list[str]:
        doc_state = self.get_doc(doc_id)
        if not doc_state:
            return []
        chunk_ids = doc_state.get("chunk_ids") or []
        if isinstance(chunk_ids, list):
            return [str(item) for item in chunk_ids]
        return []

    def update_doc(
        self,
        doc_id: str,
        *,
        source: str,
        file_hash: str,
        chunk_ids: Iterable[str],
        chunk_hashes: Iterable[str],
        model_version: str,
    ) -> None:
        self.docs[doc_id] = {
            "source": source,
            "file_hash": file_hash,
            "chunk_ids": list(chunk_ids),
            "chunk_hashes": list(chunk_hashes),
            "model_version": model_version,
            "updated_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
        }

    def remove_docs_not_in(self, valid_doc_ids: set[str]) -> Dict[str, Dict[str, Any]]:
        removed: Dict[str, Dict[str, Any]] = {}
        for doc_id in list(self.docs.keys()):
            if doc_id not in valid_doc_ids:
                removed[doc_id] = self.docs.pop(doc_id)
        return removed


__all__ = ["IndexStateStore", "sha256_file", "sha256_text"]
