"""
code_refs.py - Pluggable code/standard reference retriever for Stage 7.

Loads a local corpus of code/standard text and returns the excerpts most
relevant to the active review paths and the drawing's subject matter. Kept
dependency-free (keyword scoring, no embeddings) so it ships today; the
retrieve() interface is the seam to swap in a vector store later.

Corpus layout: drop plain-text files in backend/code_corpus/. The standard is
inferred from the filename prefix (ada*.txt -> ADA, tas*.txt -> TAS,
ibc*.txt -> IBC, ifc*.txt -> IFC, iebc*.txt -> IEBC, energy*.txt -> ENERGY).
Each file is chunked on blank lines; a leading "<number> ..." token on a chunk
is captured as its section id. An empty corpus is fine -- Stage 7 then runs
reasoning-only and cites nothing.
"""

import os
import re
from typing import Dict, List

_CORPUS_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                           "code_corpus")

# review_paths flag -> standard tag the corpus uses.
PATH_TO_STANDARD = {
    "ada_review": "ADA",
    "tas_tdlr_review": "TAS",
    "ibc_review": "IBC",
    "iebc_review": "IEBC",
    "fire_code_review": "IFC",
    "energy_review": "ENERGY",
}

_FILE_PREFIX_TO_STANDARD = [
    ("ada", "ADA"), ("tas", "TAS"), ("iebc", "IEBC"),
    ("ibc", "IBC"), ("ifc", "IFC"), ("energy", "ENERGY"),
]

_SECTION_RE = re.compile(r"^\s*(\d+[\d.]*\b)")
_WORD_RE = re.compile(r"[a-z]{4,}")

_corpus_cache: List[Dict] = None  # type: ignore


def _standard_for_file(name: str) -> str:
    stem = name.lower()
    for prefix, std in _FILE_PREFIX_TO_STANDARD:
        if stem.startswith(prefix):
            return std
    return stem.split("_")[0].split(".")[0].upper()


def load_corpus() -> List[Dict]:
    """Load and chunk the corpus once. Each chunk: {standard, section, text}."""
    global _corpus_cache
    if _corpus_cache is not None:
        return _corpus_cache
    chunks: List[Dict] = []
    if os.path.isdir(_CORPUS_DIR):
        for fname in sorted(os.listdir(_CORPUS_DIR)):
            if not fname.lower().endswith(".txt"):
                continue
            std = _standard_for_file(fname)
            with open(os.path.join(_CORPUS_DIR, fname), encoding="utf-8") as f:
                raw = f.read()
            for para in re.split(r"\n\s*\n", raw):
                text = para.strip()
                if len(text) < 20:
                    continue
                m = _SECTION_RE.match(text)
                chunks.append({
                    "standard": std,
                    "section": m.group(1) if m else None,
                    "text": text,
                })
    _corpus_cache = chunks
    return chunks


def _query_terms(assertions: List[Dict]) -> set:
    blob = " ".join(
        f"{a.get('attribute','')} {a.get('value','')} {a.get('source_text','')}"
        for a in assertions
    ).lower()
    return set(_WORD_RE.findall(blob))


def retrieve(review_paths: Dict, assertions: List[Dict], k: int = 24) -> List[Dict]:
    """Top-k corpus chunks for the active standards, scored by term overlap."""
    corpus = load_corpus()
    if not corpus:
        return []
    active = {std for flag, std in PATH_TO_STANDARD.items() if (review_paths or {}).get(flag)}
    # If the profile is empty/unknown, fall back to accessibility standards.
    pool = [c for c in corpus if not active or c["standard"] in active]
    if not pool:
        pool = corpus

    terms = _query_terms(assertions)
    scored = []
    for c in pool:
        ctext = c["text"].lower()
        score = sum(1 for t in terms if t in ctext)
        if score:
            scored.append((score, c))
    scored.sort(key=lambda sc: sc[0], reverse=True)
    return [c for _, c in scored[:k]]


def format_excerpts(excerpts: List[Dict]) -> str:
    """Render retrieved excerpts for the prompt."""
    if not excerpts:
        return "(no code excerpts available; do not cite section numbers)"
    out = []
    for e in excerpts:
        sec = f" {e['section']}" if e.get("section") else ""
        out.append(f"[{e['standard']}{sec}] {e['text']}")
    return "\n\n".join(out)