"""Recon types — the dataset's knowledge graph as seen by the rest of the api.

The dataset-model types (Column, Table, Metric, SchemaContext) used to live
in `api/tools/types.py`. They're authored here now because they describe the
dataset, not the tools that consume it.

`Recon` is a SchemaContext extended with derived indexes that let Analyses
ask things like "which table owns column A2?" or "how do I join loan to
district?" without re-deriving from raw YAML each time.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


# ── Schema model ──

@dataclass
class Column:
    name: str
    sql_type: str
    nullable: bool
    description: str | None = None
    # Domain-level type that the composer dispatches on for filter rendering.
    # Examples: "date_yymmdd" (BIRD's int-encoded YYMMDD dates), "date" (real
    # postgres date). None = treat literally with the column's sql_type.
    semantic_type: str | None = None

    @classmethod
    def from_inspector(cls, info: dict[str, Any], description: str | None = None) -> "Column":
        """Build from a SQLAlchemy Inspector column dict."""
        return cls(
            name=info["name"],
            sql_type=str(info["type"]).upper(),
            nullable=bool(info["nullable"]),
            description=description,
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "sql_type": self.sql_type,
            "nullable": self.nullable,
            "description": self.description,
            "semantic_type": self.semantic_type,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "Column":
        return cls(
            name=d["name"], sql_type=d["sql_type"],
            nullable=d["nullable"], description=d.get("description"),
            semantic_type=d.get("semantic_type"),
        )


@dataclass
class Table:
    name: str
    description: str | None
    columns: list[Column]

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "columns": [c.to_dict() for c in self.columns],
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "Table":
        return cls(
            name=d["name"],
            description=d.get("description"),
            columns=[Column.from_dict(c) for c in d.get("columns", [])],
        )


@dataclass
class Metric:
    name: str
    description: str
    sql: str
    from_table: str
    filter: str | None
    unit: str | None

    @classmethod
    def from_spec(cls, name: str, spec: dict[str, Any]) -> "Metric":
        return cls(
            name=name,
            description=spec.get("description", ""),
            sql=spec["sql"],
            from_table=spec["from_table"],
            filter=spec.get("filter"),
            unit=spec.get("unit"),
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "sql": self.sql,
            "from_table": self.from_table,
            "filter": self.filter,
            "unit": self.unit,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "Metric":
        return cls(
            name=d["name"], description=d.get("description", ""),
            sql=d["sql"], from_table=d["from_table"],
            filter=d.get("filter"), unit=d.get("unit"),
        )


# ── Relationships ──

@dataclass
class Relationship:
    from_table: str
    from_column: str
    to_table: str
    to_column: str

    @classmethod
    def parse(cls, raw: dict[str, str]) -> "Relationship":
        """Accept either `{from: 'a.b', to: 'c.d'}` or `{from_table, from_column, ...}`."""
        if "from_table" in raw:
            return cls(raw["from_table"], raw["from_column"], raw["to_table"], raw["to_column"])
        ft, fc = raw["from"].split(".", 1)
        tt, tc = raw["to"].split(".", 1)
        return cls(ft, fc, tt, tc)

    def to_dict(self) -> dict[str, str]:
        return {
            "from_table": self.from_table, "from_column": self.from_column,
            "to_table": self.to_table, "to_column": self.to_column,
        }


# ── Recon — top-level dataset knowledge bundle ──

@dataclass
class Recon:
    schema: str
    tables: dict[str, Table]
    metrics: dict[str, Metric]
    column_to_tables: dict[str, list[str]] = field(default_factory=dict)
    relationships: list[Relationship] = field(default_factory=list)

    # ── Read-side helpers used by Analyses + the composer ──

    def table_names(self) -> list[str]:
        return sorted(self.tables)

    def metric_names(self) -> list[str]:
        return sorted(self.metrics)

    def owning_tables(self, column: str) -> list[str]:
        """Tables that have a column with this name (case-sensitive)."""
        return self.column_to_tables.get(column, [])

    def resolve_column(self, ref: str) -> tuple[str, "Column"]:
        """Resolve a `column` or `table.column` reference to its (table_name,
        Column) pair. Raises ValueError if the column doesn't exist or is
        ambiguous (lives on multiple tables) without an explicit `table.` prefix.

        Used by the composer to bind every Pick.group_by entry to a single
        owning table before SQL composition. Disambiguation is on the caller:
        if you mean `account.district_id`, say so."""
        if "." in ref:
            table, col = ref.split(".", 1)
            if table not in self.tables:
                raise ValueError(f"unknown table {table!r} in column ref {ref!r}")
            for c in self.tables[table].columns:
                if c.name == col:
                    return table, c
            raise ValueError(f"column {col!r} not found on table {table!r}")
        owners = self.owning_tables(ref)
        if not owners:
            raise ValueError(f"no table has a column named {ref!r}")
        if len(owners) > 1:
            raise ValueError(
                f"column {ref!r} is ambiguous (lives in: {', '.join(owners)}); "
                f"qualify it as 'table.{ref}'"
            )
        table = owners[0]
        for c in self.tables[table].columns:
            if c.name == ref:
                return table, c
        # Index says it's there but the table doesn't — shouldn't happen.
        raise ValueError(f"column {ref!r} indexed under {table!r} but not found")

    def join_path(self, src: str, dst: str) -> list[str] | None:
        """Shortest sequence of tables from src to dst via declared relationships.
        Returns None if no path exists. The path includes both endpoints."""
        if src == dst:
            return [src]
        if src not in self.tables or dst not in self.tables:
            return None
        # Build undirected adjacency once.
        adj: dict[str, set[str]] = {t: set() for t in self.tables}
        for r in self.relationships:
            adj.setdefault(r.from_table, set()).add(r.to_table)
            adj.setdefault(r.to_table, set()).add(r.from_table)
        # BFS.
        from collections import deque
        prev: dict[str, str | None] = {src: None}
        q: deque[str] = deque([src])
        while q:
            cur = q.popleft()
            if cur == dst:
                # reconstruct
                path: list[str] = []
                node: str | None = cur
                while node is not None:
                    path.append(node)
                    node = prev[node]
                return list(reversed(path))
            for nb in adj.get(cur, ()):
                if nb not in prev:
                    prev[nb] = cur
                    q.append(nb)
        return None

    # ── Prompt-rendering helpers ──

    def render_tables(self, names: list[str] | None = None) -> str:
        """CREATE TABLE-like rendering for prompt context."""
        sel = [self.tables[n] for n in (names or self.table_names()) if n in self.tables]
        out: list[str] = []
        for t in sel:
            header = f"-- {t.description}\n" if t.description else ""
            cols: list[str] = []
            for c in t.columns:
                line = f'  "{c.name}" {c.sql_type}'
                if not c.nullable:
                    line += " NOT NULL"
                if c.description:
                    line += f"  -- {c.description}"
                cols.append(line)
            out.append(header + f'CREATE TABLE "{t.name}" (\n' + ",\n".join(cols) + "\n);")
        return "\n\n".join(out)

    def render_metrics(self) -> str:
        if not self.metrics:
            return "(no metrics defined)"
        lines: list[str] = []
        for m in self.metrics.values():
            lines.append(
                f"- {m.name} ({m.unit or 'unitless'}): {m.description}\n"
                f"    sql: {m.sql}\n"
                f"    from: {m.from_table}"
                + (f"\n    filter: {m.filter}" if m.filter else "")
            )
        return "\n".join(lines)

    def render_metrics_brief(self) -> str:
        """One line per metric — name, unit, description. No SQL, no table.

        For the planner: it picks WHICH analyses to run; it doesn't need
        implementation detail. Saves prompt tokens and cuts the noise the
        model has to filter through."""
        if not self.metrics:
            return "(no metrics defined)"
        return "\n".join(
            f"- {m.name} [{m.unit or 'unitless'}]: {m.description}"
            for m in self.metrics.values()
        )

    def render_tables_brief(self) -> str:
        """One line per table — name + description. No columns, no types."""
        if not self.tables:
            return "(no tables)"
        lines: list[str] = []
        for t in self.tables.values():
            desc = t.description or "(no description)"
            lines.append(f"- {t.name}: {desc}")
        return "\n".join(lines)

    def render_relationships(self) -> str:
        if not self.relationships:
            return "(no relationships declared)"
        return "\n".join(
            f"  {r.from_table}.{r.from_column} → {r.to_table}.{r.to_column}"
            for r in self.relationships
        )

    # ── Cache serialisation ──

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema": self.schema,
            "tables": {n: t.to_dict() for n, t in self.tables.items()},
            "metrics": {n: m.to_dict() for n, m in self.metrics.items()},
            "column_to_tables": self.column_to_tables,
            "relationships": [r.to_dict() for r in self.relationships],
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "Recon":
        return cls(
            schema=d["schema"],
            tables={n: Table.from_dict(t) for n, t in d.get("tables", {}).items()},
            metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()},
            column_to_tables=d.get("column_to_tables", {}),
            relationships=[Relationship(**r) for r in d.get("relationships", [])],
        )