app/schedule_docs.py

"""Assemble fee-code documents from the real OHIP FSM + Schedule of Benefits.

Kept free of OpenSearch / embedding / torch imports so the HF Spaces pilot
image can build or refresh a SQLite FTS index without GPU stacks.
"""
from __future__ import annotations

import hashlib
import logging

from .datasource import fetch_fsm, fetch_sob_pdfs
from .fsm_parser import FeeRecord, parse_fsm
from .search_aliases import apply_search_aliases
from .sob_descriptions import (
    extract_differentiators,
    extract_merged,
    extract_present_codes,
)

logger = logging.getLogger(__name__)

PREFIX_CATEGORY = {
    "A": "Consultations and Visits",
    "C": "Hospital In-Patient Services",
    "E": "Additional / Bundled Procedures",
    "F": "Family Practice and Obstetrics",
    "G": "Laboratory Medicine and Diagnostics",
    "H": "Emergency and Hospital Care",
    "J": "Diagnostic and Therapeutic Procedures",
    "K": "Counselling and Special Services",
    "M": "Reproductive and Genito-Urinary",
    "P": "Obstetrics",
    "Q": "Special Payments and Premiums",
    "R": "Surgical Procedures",
    "X": "Diagnostic Radiology",
    "Z": "Surgical and Diagnostic Procedures",
}


def category_for_code(code: str) -> str:
    return PREFIX_CATEGORY.get(code[:1], "General Listings")


def embedding_input(doc: dict) -> str:
    return (
        f"Section: {doc['parent_section']}. "
        f"OHIP code {doc['billing_code']}: {doc['description_text']}. "
        f"Constraints: {doc['rules_and_constraints']}"
    )


def doc_hash(doc: dict) -> str:
    key = f"{doc['description_text']}|{doc['parent_section']}"
    return hashlib.sha256(key.encode()).hexdigest()


def build_documents() -> list[dict]:
    """Fetch + parse the real sources and assemble child documents."""
    fsm_file = fetch_fsm()
    records: list[FeeRecord] = parse_fsm(fsm_file.path)

    sob_files = fetch_sob_pdfs()
    descriptions: dict[str, dict] = extract_merged(sob_files) if sob_files else {}
    differentiators: dict[str, dict] = (
        extract_differentiators(sob_files) if sob_files else {}
    )
    present: set[str] = extract_present_codes(sob_files) if sob_files else set()

    docs: list[dict] = []
    for rec in records:
        code = rec.billing_code
        enrich = descriptions.get(code)
        diff = differentiators.get(code)
        in_current = (not present) or (code in present)

        if enrich:
            description = enrich["description"]
            parent_section = enrich["section"]
            has_desc = True
        elif diff and diff.get("kind") == "matrix":
            parent_section = category_for_code(code)
            description = diff["description"]
            has_desc = True
        elif not in_current:
            parent_section = category_for_code(code)
            description = (
                f"Legacy OHIP code {code} — billable in the Fee Schedule Master "
                f"but not listed in the current Schedule of Benefits."
            )
            has_desc = False
        else:
            parent_section = category_for_code(code)
            description = f"{parent_section} service (OHIP code {code})."
            has_desc = False

        reference = diff.get("reference") if diff else None
        diff_obj = None
        if diff:
            diff_obj = {
                k: diff[k]
                for k in ("cap", "person_seen", "time_band", "reference")
                if diff.get(k)
            } or None

        rules = f"Refer to the Schedule of Benefits section '{parent_section}'."
        if reference:
            rules += f" Schedule reference: {reference}."

        docs.append(
            apply_search_aliases(
                {
                    "billing_code": code,
                    "description_text": description,
                    "rules_and_constraints": rules,
                    "base_fee_cad": round(rec.primary_fee, 2),
                    "fee_components": rec.all_fees,
                    "effective_date": rec.effective_date,
                    "termination_date": rec.termination_date,
                    "parent_section": parent_section,
                    "parent_id": parent_section.lower().replace(" ", "_"),
                    "has_description": has_desc,
                    "reference": reference,
                    "differentiators": diff_obj,
                    "in_current_schedule": in_current,
                }
            )
        )
    logger.info(
        "Built %d docs (%d described, %d matrix-differentiated, %d referenced, "
        "%d legacy/not-in-current-schedule)",
        len(docs),
        sum(1 for d in docs if d["has_description"]),
        sum(1 for c, d in differentiators.items() if d.get("kind") == "matrix"),
        sum(1 for d in docs if d["reference"]),
        sum(1 for d in docs if not d["in_current_schedule"]),
    )
    return docs