"""Exact statutory-text retrieval from a private local Chroma snapshot.

The section crosswalk decides which provisions correspond.  This store only
returns the exact Act + section record requested by that mapping; it never runs
similarity search and never substitutes a neighbouring provision.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from statute_crosswalk import normalise_act, normalise_section


COLLECTION_NAME = "indian_statutes"
STORED_ACT_CODES = {
    "IPC": ["IPC"],
    "BNS": ["BNS"],
    "CRPC": ["CRPC", "CrPC"],
    "BNSS": ["BNSS"],
    "IEA": ["IEA"],
    "BSA": ["BSA"],
}


def _chroma_root(value: str | os.PathLike[str] | None) -> Path | None:
    """Accept either the Chroma root or its UUID segment directory."""
    if not value:
        return None
    candidate = Path(value).expanduser().resolve()
    if (candidate / "chroma.sqlite3").is_file():
        return candidate
    if candidate.is_dir() and (candidate.parent / "chroma.sqlite3").is_file():
        return candidate.parent
    return candidate


class ExactStatuteLibrary:
    """Provide exact bare-act records with a JSON fallback for older releases."""

    def __init__(
        self,
        chroma_path: str | os.PathLike[str] | None = None,
        *,
        fallback_path: str | os.PathLike[str] | None = None,
        collection_name: str = COLLECTION_NAME,
    ) -> None:
        self._collection = None
        self._provider = "unavailable"
        self._count = 0
        self._fallback: dict[tuple[str, str], dict[str, Any]] = {}
        self._error: str | None = None
        self._root = _chroma_root(chroma_path)
        self._collection_name = collection_name

        if self._root and (self._root / "chroma.sqlite3").is_file():
            try:
                import chromadb

                client = chromadb.PersistentClient(path=str(self._root))
                self._collection = client.get_collection(self._collection_name)
                self._count = int(self._collection.count())
                self._provider = "private_chroma"
                return
            except Exception as exc:  # keep the API alive if the optional store is damaged
                self._error = type(exc).__name__

        self._load_fallback(fallback_path)

    @classmethod
    def from_env(
        cls,
        *,
        fallback_path: str | os.PathLike[str] | None = None,
    ) -> "ExactStatuteLibrary":
        return cls(
            os.environ.get("THEMIS_STATUTE_CHROMA", "").strip() or None,
            fallback_path=fallback_path,
            collection_name=os.environ.get(
                "THEMIS_STATUTE_COLLECTION", COLLECTION_NAME
            ).strip()
            or COLLECTION_NAME,
        )

    def _load_fallback(self, path: str | os.PathLike[str] | None) -> None:
        candidate = Path(path).expanduser().resolve() if path else None
        if not candidate or not candidate.is_file():
            return
        try:
            payload = json.loads(candidate.read_text(encoding="utf-8"))
            for item in payload if isinstance(payload, list) else []:
                metadata = item.get("metadata") if isinstance(item, dict) else None
                if not isinstance(metadata, dict):
                    continue
                act = normalise_act(metadata.get("act_short"))
                section = normalise_section(metadata.get("section_number"))
                if not act or not section:
                    continue
                self._fallback[(act, section)] = {
                    "act": act,
                    "act_name": metadata.get("act_name"),
                    "section": section,
                    "title": metadata.get("title"),
                    "text": str(item.get("retrieval_text") or "").replace("\x00", "").strip(),
                }
            if self._fallback:
                self._provider = "release_json"
                self._count = len(self._fallback)
        except Exception as exc:
            self._error = self._error or type(exc).__name__

    def lookup(self, act: object, section: object) -> dict[str, Any] | None:
        code = normalise_act(act)
        number = normalise_section(section)
        if not code or not number:
            return None
        if self._collection is None:
            value = self._fallback.get((code, number))
            return dict(value) if value else None

        try:
            result = self._collection.get(
                where={
                    "$and": [
                        {"act_short": {"$in": STORED_ACT_CODES.get(code, [code])}},
                        {"section_number": {"$eq": number}},
                    ]
                },
                include=["documents", "metadatas"],
            )
        except Exception:
            return None
        if not result.get("ids"):
            return None
        metadata = (result.get("metadatas") or [{}])[0] or {}
        documents = result.get("documents") or [""]
        return {
            "act": code,
            "act_name": metadata.get("act_name"),
            "section": number,
            "title": metadata.get("title"),
            "text": str(documents[0] or "").replace("\x00", "").strip(),
        }

    def status(self) -> dict[str, Any]:
        return {
            "configured": self._root is not None,
            "ready": self._provider != "unavailable",
            "provider": self._provider,
            "collection": self._collection_name,
            "provisions": self._count,
            "lookup": "exact_act_and_section_only",
            "semantic_conversion": False,
            "judgment_embeddings_used": False,
            "error": self._error,
        }