"""Loads the knowledge base (knowledge_base/) into memory for retrieval and citation lookup.

Only the content documented in knowledge_base/README.md as genuinely esure-sourced is here —
see that file before adding anything new to this directory.
"""
import json
from pathlib import Path

KB_DIR = Path(__file__).parent / "knowledge_base"


class EvidenceStore:
    """In-memory index over the procedures and corpus documents.

    A procedure's own `branches` can reference a rule_id or procedure_id that isn't in this
    store (see the Known gap section of knowledge_base/README.md) — callers must treat a
    missing reference as "no real content for this path", never synthesize one.
    """

    def __init__(self, kb_dir=KB_DIR):
        self.kb_dir = Path(kb_dir)
        procedures_path = self.kb_dir / "esure-procedures.json"
        data = json.loads(procedures_path.read_text())
        self.procedures = {p["procedure_id"]: p for p in data.get("procedures", [])}
        self.corpus = {c["article_id"]: c for c in data.get("corpus", [])}

        vuln_doc_path = self.kb_dir / "vulnerable-customer-process-source.txt"
        self._vuln_doc_text = vuln_doc_path.read_text()

    def procedure(self, procedure_id):
        return self.procedures.get(procedure_id)

    def document_text(self, article_id):
        """Full verbatim text for a citation, by article id. None if not held here."""
        if article_id == "VULNERABLE-CUSTOMER-PROCESS":
            return self._vuln_doc_text
        corpus_entry = self.corpus.get(article_id)
        return corpus_entry.get("passage") if corpus_entry else None

    def known_procedure_ids(self):
        return set(self.procedures.keys())

    def documents(self):
        """All citable documents, in the {article_id: {...}} shape callers can list from."""
        out = {}
        for article_id, entry in self.corpus.items():
            out[article_id] = {
                "article_id": article_id,
                "title": entry.get("article_title"),
                "text": entry.get("passage"),
            }
        out["VULNERABLE-CUSTOMER-PROCESS"] = {
            "article_id": "VULNERABLE-CUSTOMER-PROCESS",
            "title": "Vulnerable Customer Process",
            "text": self._vuln_doc_text,
        }
        return out