from __future__ import annotations

import json
from pathlib import Path

from .models import ParentDocument, RetrievalDocument, SourceRef


def _read_jsonl(path: Path) -> list[dict]:
    with path.open(encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]


def _source_ref(raw: dict) -> SourceRef:
    bbox = raw.get("bbox")
    page_range = raw.get("page_range")
    printed_page_range = raw.get("printed_page_range")
    return SourceRef(
        physical_page=int(raw["physical_page"]),
        precision=raw["precision"],
        block_id=raw.get("block_id"),
        bbox=tuple(bbox) if bbox else None,
        source_crop=raw.get("source_crop"),
        page_range=tuple(page_range) if page_range else None,
        printed_page=raw.get("printed_page"),
        printed_page_range=(
            tuple(printed_page_range) if printed_page_range else None
        ),
    )


def load_documents(path: Path) -> list[RetrievalDocument]:
    documents = []
    for raw in _read_jsonl(path):
        documents.append(RetrievalDocument(
            doc_id=raw["doc_id"],
            drug_id=raw["drug_id"],
            kind=raw["kind"],
            text=raw["text"],
            section_key=raw["section_key"],
            source_refs=tuple(_source_ref(item) for item in raw.get("source_refs", [])),
            parent_id=raw.get("parent_id"),
            requires_visual_check=raw.get("requires_visual_check", False),
            drug_name=raw.get("drug_name"),
        ))
    return documents


def load_parents(path: Path) -> list[ParentDocument]:
    parents = []
    for raw in _read_jsonl(path):
        parents.append(ParentDocument(
            parent_id=raw["logical_table_id"],
            kind=raw["kind"],
            text=raw["markdown"],
            source_refs=tuple(_source_ref(item) for item in raw.get("source_refs", [])),
            requires_visual_check=raw.get("requires_visual_check", False),
        ))
    return parents


def load_aliases(path: Path | None) -> dict[str, set[str]]:
    if path is None:
        return {}
    raw = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(raw, dict) and "entities" in raw:
        return {
            entity["drug_id"]: set(entity["aliases"])
            for entity in raw["entities"]
        }
    return {drug_id: set(aliases) for drug_id, aliases in raw.items()}


def build_drug_catalog(
    documents: list[RetrievalDocument],
    extra_aliases: dict[str, set[str]] | None = None,
) -> dict[str, set[str]]:
    catalog: dict[str, set[str]] = {}
    for document in documents:
        aliases = catalog.setdefault(document.drug_id, set())
        aliases.add(document.drug_id.replace("_", " "))
        if document.drug_name:
            aliases.add(document.drug_name)
    for drug_id, aliases in (extra_aliases or {}).items():
        catalog.setdefault(drug_id, set()).update(aliases)
    return catalog