"""Domain-specific loaders for open tabular and catalog sources.

Each loader produces the same ``RealGraph`` contract used by the KG corpus
builder.  Text or numeric fields stay in node attributes; topology carries the
relationships that the graph reasoner can traverse.
"""

from __future__ import annotations

import csv
import gzip
import io
import re
import zipfile
from collections.abc import Iterable, Mapping
from pathlib import Path

import pyarrow.parquet as pq

from graph.graph import EdgeRecord, GraphStore
from graph.real.corpus import RealGraph, SourceMetadata, relation_phrase


def _make_graph(
    source: SourceMetadata,
    nodes: Mapping[str, str],
    edges: Iterable[EdgeRecord],
    *,
    labels: Mapping[str, str],
    node_attributes: Mapping[str, Mapping[str, str]],
    relation_ids: Iterable[str],
) -> RealGraph:
    relation_ids = tuple(sorted(set(relation_ids)))
    relation_attributes = {
        relation: {
            "functional": "false",
            "description": relation_phrase(relation),
        }
        for relation in relation_ids
    }
    edge_list = tuple(edges)
    graph = GraphStore(
        nodes,
        edge_list,
        graph_version=f"{source.source_id}@{source.version}",
        node_attributes=node_attributes,
        relation_attributes=relation_attributes,
    )
    return RealGraph(
        graph=graph,
        node_labels=labels,
        relation_phrases={
            relation: relation_phrase(relation) for relation in relation_ids
        },
        source=source,
    )


def _slug(value: str, *, limit: int = 100) -> str:
    value = re.sub(r"[^A-Za-z0-9]+", "_", value.strip()).strip("_")
    return value[:limit] or "unknown"


def load_math_parquet(
    path: str | Path, source: SourceMetadata, *, max_rows: int | None = None
) -> RealGraph:
    """Load competition mathematics problems as topic/level/solution graphs."""
    rows = pq.read_table(path).to_pylist()
    if max_rows is not None:
        rows = rows[:max_rows]
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relations = {
        "has_topic",
        "has_level",
        "has_solution",
        "topic_at_level",
        "solves_problem",
    }
    for index, row in enumerate(rows):
        problem_id = f"math:problem:{index}"
        solution_id = f"math:solution:{index}"
        topic_id = f"math:topic:{_slug(str(row.get('type', 'unknown')))}"
        level_id = f"math:level:{_slug(str(row.get('level', 'unknown')))}"
        nodes.update(
            {
                problem_id: "math_problem",
                solution_id: "math_solution",
                topic_id: "math_topic",
                level_id: "math_level",
            }
        )
        labels[problem_id] = f"problem {index}"
        labels[solution_id] = f"solution {index}"
        labels[topic_id] = str(row.get("type", "unknown"))
        labels[level_id] = str(row.get("level", "unknown"))
        attributes[problem_id] = {
            "problem": str(row.get("problem", "")),
            "topic": str(row.get("type", "")),
            "level": str(row.get("level", "")),
        }
        attributes[solution_id] = {"solution": str(row.get("solution", ""))}
        edges.extend(
            (
                EdgeRecord(
                    problem_id,
                    "has_topic",
                    topic_id,
                    "math_problem",
                    "math_topic",
                    relation_phrase("has_topic"),
                    {},
                ),
                EdgeRecord(
                    problem_id,
                    "has_level",
                    level_id,
                    "math_problem",
                    "math_level",
                    relation_phrase("has_level"),
                    {},
                ),
                EdgeRecord(
                    problem_id,
                    "has_solution",
                    solution_id,
                    "math_problem",
                    "math_solution",
                    relation_phrase("has_solution"),
                    {},
                ),
                EdgeRecord(
                    solution_id,
                    "solves_problem",
                    problem_id,
                    "math_solution",
                    "math_problem",
                    relation_phrase("solves_problem"),
                    {},
                ),
                EdgeRecord(
                    topic_id,
                    "topic_at_level",
                    level_id,
                    "math_topic",
                    "math_level",
                    relation_phrase("topic_at_level"),
                    {},
                ),
            )
        )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relations,
    )


def load_world_bank_wdi(
    archive: str | Path,
    source: SourceMetadata,
    *,
    max_rows: int = 50_000,
    years: tuple[str, ...] = ("1990", "2000", "2010", "2020", "2024"),
) -> RealGraph:
    """Load a bounded, multi-year WDI slice as country-observation graphs."""
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relations = {"has_observation", "measures", "in_year", "has_topic"}
    kept = 0
    indicator_topics: dict[str, str] = {}
    with zipfile.ZipFile(archive) as handle:
        with handle.open("WDISeries.csv") as binary:
            text = io.TextIOWrapper(binary, encoding="utf-8-sig", newline="")
            for row in csv.DictReader(text):
                code = str(row.get("Series Code", "")).strip()
                topic = str(row.get("Topic", "")).strip()
                if code and topic:
                    indicator_topics[code] = topic

        with handle.open("WDICSV.csv") as binary:
            text = io.TextIOWrapper(binary, encoding="utf-8-sig", newline="")
            reader = csv.DictReader(text)
            for row in reader:
                country_code = str(row.get("Country Code", "")).strip()
                indicator_code = str(row.get("Indicator Code", "")).strip()
                if not country_code or not indicator_code:
                    continue
                selected = [
                    (year, str(row.get(year, "")).strip())
                    for year in years
                    if str(row.get(year, "")).strip()
                ]
                if not selected:
                    continue
                country_id = f"wdi:country:{country_code}"
                indicator_id = f"wdi:indicator:{indicator_code}"
                nodes[country_id] = "country"
                nodes[indicator_id] = "indicator"
                labels[country_id] = str(row.get("Country Name", country_code))
                labels[indicator_id] = str(row.get("Indicator Name", indicator_code))
                metadata_topic = indicator_topics.get(indicator_code)
                if metadata_topic:
                    topic_id = f"wdi:topic:{_slug(metadata_topic)}"
                    nodes[topic_id] = "topic"
                    labels[topic_id] = metadata_topic
                    edges.append(
                        EdgeRecord(
                            indicator_id,
                            "has_topic",
                            topic_id,
                            "indicator",
                            "topic",
                            relation_phrase("has_topic"),
                            {},
                        )
                    )
                for year, value in selected:
                    observation_id = (
                        f"wdi:observation:{country_code}:{indicator_code}:{year}"
                    )
                    year_id = f"wdi:year:{year}"
                    nodes[observation_id] = "observation"
                    nodes[year_id] = "year"
                    labels[observation_id] = f"{labels[country_id]} {year}"
                    labels[year_id] = year
                    attributes[observation_id] = {
                        "value": value,
                        "indicator": labels[indicator_id],
                    }
                    edges.extend(
                        (
                            EdgeRecord(
                                country_id,
                                "has_observation",
                                observation_id,
                                "country",
                                "observation",
                                relation_phrase("has_observation"),
                                {},
                            ),
                            EdgeRecord(
                                observation_id,
                                "measures",
                                indicator_id,
                                "observation",
                                "indicator",
                                relation_phrase("measures"),
                                {},
                            ),
                            EdgeRecord(
                                observation_id,
                                "in_year",
                                year_id,
                                "observation",
                                "year",
                                relation_phrase("in_year"),
                                {},
                            ),
                        )
                    )
                kept += 1
                if kept >= max_rows:
                    break
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relations,
    )


def _dbf_records(payload: bytes) -> list[dict[str, str]]:
    """Read the dBASE attribute table bundled with a Natural Earth shapefile."""
    if payload[0] not in range(0x03, 0x8C):
        raise ValueError("unsupported DBF version")
    record_count = int.from_bytes(payload[4:8], "little")
    header_length = int.from_bytes(payload[8:10], "little")
    record_length = int.from_bytes(payload[10:12], "little")
    fields: list[tuple[str, int]] = []
    position = 32
    while payload[position] != 0x0D:
        descriptor = payload[position : position + 32]
        name = descriptor[:11].split(b"\0", 1)[0].decode("ascii", "ignore")
        fields.append((name, descriptor[16]))
        position += 32
    records: list[dict[str, str]] = []
    for index in range(record_count):
        start = header_length + index * record_length
        record = payload[start : start + record_length]
        if not record or record[0] == 0x2A:
            continue
        offset = 1
        values: dict[str, str] = {}
        for name, width in fields:
            values[name] = (
                record[offset : offset + width].decode("latin-1", "ignore").strip()
            )
            offset += width
        records.append(values)
    return records


def load_natural_earth_countries(
    archive: str | Path, source: SourceMetadata
) -> RealGraph:
    """Load Natural Earth country attributes and administrative hierarchy."""
    with zipfile.ZipFile(archive) as handle:
        dbf_name = next(name for name in handle.namelist() if name.endswith(".dbf"))
        records = _dbf_records(handle.read(dbf_name))
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relations = {
        "in_continent",
        "in_region",
        "in_subregion",
        "within_region",
        "within_continent",
    }
    subregion_regions: dict[str, str] = {}
    region_continents: dict[str, str] = {}
    for record in records:
        iso = record.get("ADM0_A3") or record.get("ISO_A3") or record.get("NAME")
        if not iso or iso == "-99":
            iso = record.get("NAME", "unknown")
        country_id = f"earth:country:{_slug(iso)}"
        name = record.get("NAME_LONG") or record.get("NAME") or iso
        nodes[country_id] = "country"
        labels[country_id] = name
        attributes[country_id] = {
            key: value
            for key in ("CONTINENT", "REGION_UN", "SUBREGION", "POP_EST", "GDP_MD")
            if (value := record.get(key, ""))
        }
        for field, relation, node_type in (
            ("CONTINENT", "in_continent", "continent"),
            ("REGION_UN", "in_region", "region"),
            ("SUBREGION", "in_subregion", "subregion"),
        ):
            value = record.get(field, "").strip()
            if not value:
                continue
            target_id = f"earth:{node_type}:{_slug(value)}"
            nodes[target_id] = node_type
            labels[target_id] = value
            if field == "SUBREGION":
                subregion_regions[target_id] = ""
            elif field == "REGION_UN":
                region_continents[target_id] = ""
            edges.append(
                EdgeRecord(
                    country_id,
                    relation,
                    target_id,
                    "country",
                    node_type,
                    relation_phrase(relation),
                    {},
                )
            )
        continent = record.get("CONTINENT", "").strip()
        region = record.get("REGION_UN", "").strip()
        subregion = record.get("SUBREGION", "").strip()
        if region and continent:
            region_continents[f"earth:region:{_slug(region)}"] = (
                f"earth:continent:{_slug(continent)}"
            )
        if subregion and region:
            subregion_regions[f"earth:subregion:{_slug(subregion)}"] = (
                f"earth:region:{_slug(region)}"
            )
    for subregion_id, region_id in subregion_regions.items():
        if region_id:
            edges.append(
                EdgeRecord(
                    subregion_id,
                    "within_region",
                    region_id,
                    "subregion",
                    "region",
                    relation_phrase("within_region"),
                    {},
                )
            )
    for region_id, continent_id in region_continents.items():
        if continent_id:
            edges.append(
                EdgeRecord(
                    region_id,
                    "within_continent",
                    continent_id,
                    "region",
                    "continent",
                    relation_phrase("within_continent"),
                    {},
                )
            )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relations,
    )


def load_gutenberg_catalog(
    path: str | Path,
    source: SourceMetadata,
    *,
    max_rows: int = 50_000,
) -> RealGraph:
    """Load Project Gutenberg catalog metadata, not copyrighted book text."""
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    relations = {
        "written_by",
        "in_language",
        "has_subject",
        "in_bookshelf",
        "author_has_book",
        "language_has_book",
        "subject_has_book",
        "bookshelf_has_book",
    }
    reverse_degree_cap = 32
    reverse_counts: dict[str, int] = {}
    with gzip.open(path, "rt", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle)
        for index, row in enumerate(reader):
            if index >= max_rows:
                break
            book_id = f"gutenberg:book:{row.get('Text#', index)}"
            title = " ".join(str(row.get("Title", "")).split())
            nodes[book_id] = "book"
            labels[book_id] = title[:160] or book_id
            attributes[book_id] = {
                key: str(row.get(key, ""))[:4000]
                for key in ("Title", "Issued", "LoCC")
                if row.get(key)
            }
            for field, relation, reverse_relation, node_type in (
                ("Language", "in_language", "language_has_book", "language"),
                ("Authors", "written_by", "author_has_book", "author"),
                ("Subjects", "has_subject", "subject_has_book", "subject"),
                ("Bookshelves", "in_bookshelf", "bookshelf_has_book", "bookshelf"),
            ):
                values = [
                    item.strip()
                    for item in str(row.get(field, "")).split(";")
                    if item.strip()
                ]
                for value in values[:32]:
                    target_id = f"gutenberg:{node_type}:{_slug(value)}"
                    nodes[target_id] = node_type
                    labels[target_id] = value[:160]
                    edges.append(
                        EdgeRecord(
                            book_id,
                            relation,
                            target_id,
                            "book",
                            node_type,
                            relation_phrase(relation),
                            {},
                        )
                    )
                    reverse_key = f"{target_id}:{reverse_relation}"
                    if reverse_counts.get(reverse_key, 0) < reverse_degree_cap:
                        edges.append(
                            EdgeRecord(
                                target_id,
                                reverse_relation,
                                book_id,
                                node_type,
                                "book",
                                relation_phrase(reverse_relation),
                                {},
                            )
                        )
                        reverse_counts[reverse_key] = (
                            reverse_counts.get(reverse_key, 0) + 1
                        )
    return _make_graph(
        source,
        nodes,
        edges,
        labels=labels,
        node_attributes=attributes,
        relation_ids=relations,
    )