"""Recon build — two-stage:

1. **DDL extraction** (auto, from Postgres): reads every table's columns +
   types + nullability via SQLAlchemy Inspector and writes
   `api/datasets/<name>/extracted_schema.json`. This is the structural
   source of truth. Re-run whenever the warehouse changes.

2. **Augmentation merge** (human YAML): reads `schema_docs.yaml` (sparse —
   table descriptions, column descriptions, relationships) and `metrics.yaml`,
   merges them onto the extracted schema, and writes
   `api/datasets/<name>/recon.json` — the artefact the runtime reads.

Run via `make recon` or `python -m api.recon.build`. Auto-triggered by
load_recon() on first call if recon.json is missing.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any

import yaml
from sqlalchemy import inspect

from api.datasets import DATASETS_DIR
from api.recon.types import Column, Metric, Recon, Relationship, Table
from api.tools.db import get_engine

logger = logging.getLogger("nvi.recon.build")


# ── Stage 1: DDL extraction ──

def extract_ddl(dataset: str) -> dict[str, Any]:
    """Introspect Postgres for `<dataset>` schema → structural dict.

    Shape:
      {
        "schema": "<name>",
        "tables": {
          "<table>": {
            "columns": [
              {"name": ..., "sql_type": ..., "nullable": ...},
              ...
            ]
          },
          ...
        }
      }
    """
    insp = inspect(get_engine())
    tables: dict[str, dict[str, Any]] = {}
    for name in insp.get_table_names(schema=dataset):
        cols = []
        for c in insp.get_columns(name, schema=dataset):
            cols.append({
                "name": c["name"],
                "sql_type": str(c["type"]).upper(),
                "nullable": bool(c["nullable"]),
            })
        tables[name] = {"columns": cols}
    return {"schema": dataset, "tables": tables}


def write_extracted_schema(dataset: str, extracted: dict[str, Any]) -> Path:
    out = DATASETS_DIR / dataset / "extracted_schema.json"
    out.write_text(json.dumps(extracted, indent=2) + "\n", encoding="utf-8")
    return out


# ── Stage 2: augmentation merge ──

def _read_yaml(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    return yaml.safe_load(path.read_text()) or {}


def _parse_column_descs(raw: dict[str, Any]) -> dict[str, str]:
    """Accept either the flat form `{table.col: desc}` or the nested form
    `{table: {col: desc}}`. Returns flat form."""
    out: dict[str, str] = {}
    for k, v in raw.items():
        if isinstance(v, dict):
            for col, desc in v.items():
                out[f"{k}.{col}"] = desc
        else:
            out[k] = v
    return out


def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon:
    """Combine extracted DDL with human-authored augmentation YAML."""
    schema_name = extracted.get("schema", dataset)

    aug = _read_yaml(DATASETS_DIR / dataset / "schema_docs.yaml")
    metrics_yaml = _read_yaml(DATASETS_DIR / dataset / "metrics.yaml")

    table_descs: dict[str, str] = aug.get("tables", {}) or {}
    col_descs = _parse_column_descs(aug.get("columns", {}) or {})

    tables: dict[str, Table] = {}
    for tname, t_data in extracted["tables"].items():
        cols = [
            Column(
                name=c["name"],
                sql_type=c["sql_type"],
                nullable=c["nullable"],
                description=col_descs.get(f"{tname}.{c['name']}"),
            )
            for c in t_data["columns"]
        ]
        tables[tname] = Table(
            name=tname,
            description=table_descs.get(tname),
            columns=cols,
        )

    metrics: dict[str, Metric] = {
        name: Metric.from_spec(name, spec)
        for name, spec in (metrics_yaml.get("metrics", {}) or {}).items()
    }

    relationships = [Relationship.parse(r) for r in (aug.get("relationships", []) or [])]

    column_to_tables: dict[str, list[str]] = {}
    for t in tables.values():
        for c in t.columns:
            column_to_tables.setdefault(c.name, []).append(t.name)
    for k in column_to_tables:
        column_to_tables[k] = sorted(column_to_tables[k])

    return Recon(
        schema=schema_name,
        tables=tables,
        metrics=metrics,
        column_to_tables=column_to_tables,
        relationships=relationships,
    )


def write_recon(dataset: str, recon: Recon) -> Path:
    out = DATASETS_DIR / dataset / "recon.json"
    out.write_text(json.dumps(recon.to_dict(), indent=2) + "\n", encoding="utf-8")
    return out


# ── Public entry: full build ──

def build_recon(dataset: str) -> Recon:
    """End-to-end: extract DDL, write the extracted artefact, merge with
    YAML, write the recon artefact. Returns the in-memory Recon."""
    extracted = extract_ddl(dataset)
    write_extracted_schema(dataset, extracted)
    recon = merge_into_recon(dataset, extracted)
    write_recon(dataset, recon)
    return recon


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
    ap = argparse.ArgumentParser(description="Build the recon artefacts for one or all datasets.")
    ap.add_argument("--dataset", help="Dataset name (default: all subdirs of api/datasets/).")
    args = ap.parse_args()

    if args.dataset:
        targets = [args.dataset]
    else:
        targets = [
            p.name for p in DATASETS_DIR.iterdir()
            if p.is_dir() and not p.name.startswith("_") and not p.name.startswith(".")
        ]

    if not targets:
        print("no datasets to build", file=sys.stderr)
        sys.exit(1)

    for name in targets:
        recon = build_recon(name)
        print(
            f"recon[{name}]: {len(recon.tables)} tables, "
            f"{len(recon.metrics)} metrics, {len(recon.relationships)} relationships "
            f"→ api/datasets/{name}/{{extracted_schema,recon}}.json"
        )


if __name__ == "__main__":
    main()