"""Read-only, versioned table-lineage projection for the DAX schema browser."""
from __future__ import annotations

import json
from collections import defaultdict
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, Literal

from app.domains.data_access.lineage_catalog import CATALOG_SCHEMA_VERSION

MAX_LINEAGE_NODES = 300
LINEAGE_CATALOG_DIR = Path(__file__).with_name("lineage_catalogs")
LINEAGE_CATALOGS_BY_PROJECT = {
    "gcp-digital-bi": "gcp-digital-bi-dbt.json",
    "gcp-digital-bi-dev": "gcp-digital-bi-dbt.json",
}


class DataAccessLineageCatalogError(RuntimeError):
    """The deployable lineage catalog is malformed or cannot be loaded."""


@dataclass(frozen=True, slots=True)
class DataAccessLineageNode:
    node_id: str
    project_id: str
    dataset_id: str
    table_id: str
    node_type: str
    materialization: str | None
    description: str | None
    level: int
    is_root: bool
    is_target: bool

    @property
    def layer(self) -> str:
        if self.node_type == "source":
            return "source"
        if self.is_target:
            return "serving"
        normalized_dataset = self.dataset_id.casefold()
        if normalized_dataset.startswith("stg_"):
            return "staging"
        if normalized_dataset.startswith(("dwd_", "dwm_", "dwh_", "model_")):
            return "transformation"
        return "intermediate"


@dataclass(frozen=True, slots=True)
class DataAccessLineageEdge:
    upstream_node_id: str
    downstream_node_id: str
    relationship: str


@dataclass(frozen=True, slots=True)
class DataAccessTableLineage:
    status: str
    direction: Literal["upstream", "downstream"]
    catalog_id: str | None
    source_version: str | None
    generated_at: str | None
    project_id: str
    dataset_id: str
    table_id: str
    nodes: tuple[DataAccessLineageNode, ...]
    edges: tuple[DataAccessLineageEdge, ...]
    truncated: bool


@dataclass(frozen=True, slots=True)
class DataAccessLineageCatalogTable:
    table_id: str
    table_type: str
    description: str | None = None
    downstream_reference_count: int = 0
    lineage_direction: Literal["upstream", "downstream"] = "upstream"


@dataclass(frozen=True, slots=True)
class DataAccessLineageCatalogSearchMatch:
    project_id: str
    dataset_id: str
    table: DataAccessLineageCatalogTable


def list_lineage_catalog_tables(
    *,
    project_id: str,
    dataset_id: str,
) -> tuple[DataAccessLineageCatalogTable, ...]:
    """List tracked relations for one dataset from the deployed catalog.

    This read model is intentionally provider-independent.  Callers must keep
    these rows read-only because the catalog proves lineage, not current
    BigQuery existence or schema metadata. Referenced ``source`` relations are
    included so source datasets can use the same request-authoring drawer as
    dbt-produced datasets.
    """
    catalog_filename = LINEAGE_CATALOGS_BY_PROJECT.get(project_id.casefold())
    if catalog_filename is None:
        return ()
    catalog = _load_catalog(LINEAGE_CATALOG_DIR / catalog_filename)
    tables = (
        DataAccessLineageCatalogTable(
            table_id=node.table_id,
            table_type=_catalog_table_type(node.materialization),
            description=node.description,
            downstream_reference_count=(
                len({
                    edge.downstream_node_id
                    for edge in catalog.children_by_node_id.get(node.node_id, ())
                })
                if node.node_type == "source"
                else 0
            ),
            lineage_direction=(
                "downstream"
                if node.node_type == "source"
                else "upstream"
            ),
        )
        for node in catalog.nodes.values()
        if node.project_id == project_id
        and node.dataset_id == dataset_id
    )
    return tuple(sorted(tables, key=lambda table: table.table_id.casefold()))


def search_lineage_catalog_tables(
    search: str,
) -> tuple[DataAccessLineageCatalogSearchMatch, ...]:
    """Search deployed, curated relations without scanning BigQuery providers."""
    normalized_search = search.strip().casefold()
    if not normalized_search:
        return ()

    matches: dict[
        tuple[str, str, str],
        DataAccessLineageCatalogSearchMatch,
    ] = {}
    for catalog_filename in sorted(set(LINEAGE_CATALOGS_BY_PROJECT.values())):
        catalog = _load_catalog(LINEAGE_CATALOG_DIR / catalog_filename)
        for node in catalog.nodes.values():
            if (
                normalized_search not in node.table_id.casefold()
                and normalized_search not in str(node.description or "").casefold()
            ):
                continue
            locator = (node.project_id, node.dataset_id, node.table_id)
            matches[locator] = DataAccessLineageCatalogSearchMatch(
                project_id=node.project_id,
                dataset_id=node.dataset_id,
                table=DataAccessLineageCatalogTable(
                    table_id=node.table_id,
                    table_type=_catalog_table_type(node.materialization),
                    description=node.description,
                    downstream_reference_count=(
                        len({
                            edge.downstream_node_id
                            for edge in catalog.children_by_node_id.get(node.node_id, ())
                        })
                        if node.node_type == "source"
                        else 0
                    ),
                    lineage_direction=(
                        "downstream" if node.node_type == "source" else "upstream"
                    ),
                ),
            )
    return tuple(matches[key] for key in sorted(matches))


def _catalog_table_type(materialization: str | None) -> str:
    normalized = str(materialization or "").strip().casefold()
    if normalized == "view":
        return "VIEW"
    if normalized == "materialized_view":
        return "MATERIALIZED_VIEW"
    return "TABLE"


def get_table_lineage(
    *,
    project_id: str,
    dataset_id: str,
    table_id: str,
) -> DataAccessTableLineage:
    """Return ancestors for one exact table locator from the dbt read model.

    An untracked relation is intentionally a successful ``not_tracked`` result:
    it may be a direct raw table, a manually maintained object, or simply not
    be in the deployed dbt manifest.  Helix must not invent an upstream path.
    """
    catalog_filename = LINEAGE_CATALOGS_BY_PROJECT.get(project_id.casefold())
    if catalog_filename is None:
        return _not_tracked(project_id, dataset_id, table_id)
    catalog = _load_catalog(LINEAGE_CATALOG_DIR / catalog_filename)
    target_ids = catalog.locator_index.get((project_id, dataset_id, table_id), ())
    if not target_ids:
        return _not_tracked(
            project_id,
            dataset_id,
            table_id,
            catalog_id=catalog.catalog_id,
            source_version=catalog.source_version,
            generated_at=catalog.generated_at,
        )
    if len(target_ids) != 1:
        raise DataAccessLineageCatalogError(
            "Lineage catalog has duplicate nodes for one BigQuery relation.",
        )

    target_id = target_ids[0]
    direction: Literal["upstream", "downstream"] = (
        "downstream"
        if catalog.nodes[target_id].node_type == "source"
        else "upstream"
    )
    if direction == "downstream":
        accepted_ids, accepted_edges, truncated = _collect_descendants(
            catalog,
            target_id,
        )
    else:
        accepted_ids, accepted_edges, truncated = _collect_ancestors(
            catalog,
            target_id,
        )
    level_by_node_id = _levels_from_sources(catalog, accepted_ids)
    root_ids = {
        node_id
        for node_id in accepted_ids
        if not any(
            edge.upstream_node_id in accepted_ids
            for edge in catalog.parents_by_node_id.get(node_id, ())
        )
    }
    nodes = tuple(
        DataAccessLineageNode(
            node_id=node_id,
            project_id=catalog.nodes[node_id].project_id,
            dataset_id=catalog.nodes[node_id].dataset_id,
            table_id=catalog.nodes[node_id].table_id,
            node_type=catalog.nodes[node_id].node_type,
            materialization=catalog.nodes[node_id].materialization,
            description=catalog.nodes[node_id].description,
            level=level_by_node_id[node_id],
            is_root=node_id in root_ids,
            is_target=node_id == target_id,
        )
        for node_id in sorted(
            accepted_ids,
            key=lambda value: (
                level_by_node_id[value],
                catalog.nodes[value].dataset_id.casefold(),
                catalog.nodes[value].table_id.casefold(),
                value,
            ),
        )
    )
    return DataAccessTableLineage(
        status="available",
        direction=direction,
        catalog_id=catalog.catalog_id,
        source_version=catalog.source_version,
        generated_at=catalog.generated_at,
        project_id=project_id,
        dataset_id=dataset_id,
        table_id=table_id,
        nodes=nodes,
        edges=tuple(sorted(
            accepted_edges,
            key=lambda edge: (
                edge.upstream_node_id,
                edge.downstream_node_id,
                edge.relationship,
            ),
        )),
        truncated=truncated,
    )


@dataclass(frozen=True, slots=True)
class _CatalogNode:
    node_id: str
    project_id: str
    dataset_id: str
    table_id: str
    node_type: str
    materialization: str | None
    description: str | None


@dataclass(frozen=True, slots=True)
class _LineageCatalog:
    catalog_id: str
    source_version: str
    generated_at: str
    nodes: dict[str, _CatalogNode]
    parents_by_node_id: dict[str, tuple[DataAccessLineageEdge, ...]]
    children_by_node_id: dict[str, tuple[DataAccessLineageEdge, ...]]
    locator_index: dict[tuple[str, str, str], tuple[str, ...]]


def _not_tracked(
    project_id: str,
    dataset_id: str,
    table_id: str,
    *,
    catalog_id: str | None = None,
    source_version: str | None = None,
    generated_at: str | None = None,
) -> DataAccessTableLineage:
    return DataAccessTableLineage(
        status="not_tracked",
        direction="upstream",
        catalog_id=catalog_id,
        source_version=source_version,
        generated_at=generated_at,
        project_id=project_id,
        dataset_id=dataset_id,
        table_id=table_id,
        nodes=(),
        edges=(),
        truncated=False,
    )


def _collect_ancestors(
    catalog: _LineageCatalog,
    target_id: str,
) -> tuple[set[str], set[DataAccessLineageEdge], bool]:
    accepted_ids: set[str] = {target_id}
    accepted_edges: set[DataAccessLineageEdge] = set()
    pending = [target_id]
    truncated = False
    while pending:
        downstream_id = pending.pop()
        for edge in catalog.parents_by_node_id.get(downstream_id, ()):
            if edge.upstream_node_id not in accepted_ids:
                if len(accepted_ids) >= MAX_LINEAGE_NODES:
                    truncated = True
                    continue
                accepted_ids.add(edge.upstream_node_id)
                pending.append(edge.upstream_node_id)
            if edge.upstream_node_id in accepted_ids:
                accepted_edges.add(edge)
    return accepted_ids, accepted_edges, truncated


def _collect_descendants(
    catalog: _LineageCatalog,
    source_id: str,
) -> tuple[set[str], set[DataAccessLineageEdge], bool]:
    accepted_ids: set[str] = {source_id}
    accepted_edges: set[DataAccessLineageEdge] = set()
    pending = [source_id]
    truncated = False
    while pending:
        upstream_id = pending.pop()
        for edge in catalog.children_by_node_id.get(upstream_id, ()):
            if edge.downstream_node_id not in accepted_ids:
                if len(accepted_ids) >= MAX_LINEAGE_NODES:
                    truncated = True
                    continue
                accepted_ids.add(edge.downstream_node_id)
                pending.append(edge.downstream_node_id)
            if edge.downstream_node_id in accepted_ids:
                accepted_edges.add(edge)
    return accepted_ids, accepted_edges, truncated


def _levels_from_sources(
    catalog: _LineageCatalog,
    accepted_ids: set[str],
) -> dict[str, int]:
    cache: dict[str, int] = {}
    active: set[str] = set()

    def level(node_id: str) -> int:
        if node_id in cache:
            return cache[node_id]
        if node_id in active:
            raise DataAccessLineageCatalogError("Lineage catalog contains a cycle.")
        active.add(node_id)
        parent_levels = [
            level(edge.upstream_node_id)
            for edge in catalog.parents_by_node_id.get(node_id, ())
            if edge.upstream_node_id in accepted_ids
        ]
        active.remove(node_id)
        cache[node_id] = max(parent_levels, default=-1) + 1
        return cache[node_id]

    for node_id in accepted_ids:
        level(node_id)
    return cache


@lru_cache(maxsize=8)
def _load_catalog(path: Path) -> _LineageCatalog:
    try:
        raw = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise DataAccessLineageCatalogError(
            "The deployed DAX lineage catalog is unavailable.",
        ) from exc
    if not isinstance(raw, dict):
        raise DataAccessLineageCatalogError("Lineage catalog must be a JSON object.")
    if raw.get("schema_version") != CATALOG_SCHEMA_VERSION:
        raise DataAccessLineageCatalogError("Lineage catalog schema version is unsupported.")
    catalog_id = _required_text(raw, "catalog_id")
    source_version = _required_text(raw, "source_version")
    generated_at = _required_text(raw, "generated_at")
    nodes_raw = raw.get("nodes")
    edges_raw = raw.get("edges")
    if not isinstance(nodes_raw, list) or not isinstance(edges_raw, list):
        raise DataAccessLineageCatalogError("Lineage catalog nodes and edges must be lists.")

    nodes: dict[str, _CatalogNode] = {}
    locator_index: dict[tuple[str, str, str], list[str]] = defaultdict(list)
    for item in nodes_raw:
        node = _parse_node(item)
        if node.node_id in nodes:
            raise DataAccessLineageCatalogError("Lineage catalog contains duplicate node ids.")
        nodes[node.node_id] = node
        locator_index[(node.project_id, node.dataset_id, node.table_id)].append(node.node_id)

    parents_by_node_id: dict[str, list[DataAccessLineageEdge]] = defaultdict(list)
    children_by_node_id: dict[str, list[DataAccessLineageEdge]] = defaultdict(list)
    for item in edges_raw:
        edge = _parse_edge(item)
        if edge.upstream_node_id not in nodes or edge.downstream_node_id not in nodes:
            raise DataAccessLineageCatalogError(
                "Lineage catalog edge refers to an unknown node.",
            )
        parents_by_node_id[edge.downstream_node_id].append(edge)
        children_by_node_id[edge.upstream_node_id].append(edge)
    return _LineageCatalog(
        catalog_id=catalog_id,
        source_version=source_version,
        generated_at=generated_at,
        nodes=nodes,
        parents_by_node_id={
            node_id: tuple(sorted(edges, key=lambda edge: edge.upstream_node_id))
            for node_id, edges in parents_by_node_id.items()
        },
        children_by_node_id={
            node_id: tuple(sorted(edges, key=lambda edge: edge.downstream_node_id))
            for node_id, edges in children_by_node_id.items()
        },
        locator_index={
            locator: tuple(sorted(node_ids))
            for locator, node_ids in locator_index.items()
        },
    )


def clear_lineage_catalog_cache() -> None:
    """Test-only seam for a replacement catalog path."""
    _load_catalog.cache_clear()


def _parse_node(value: object) -> _CatalogNode:
    if not isinstance(value, dict):
        raise DataAccessLineageCatalogError("Lineage catalog node must be an object.")
    return _CatalogNode(
        node_id=_required_text(value, "node_id"),
        project_id=_required_text(value, "project_id"),
        dataset_id=_required_text(value, "dataset_id"),
        table_id=_required_text(value, "table_id"),
        node_type=_required_text(value, "node_type"),
        materialization=_optional_text(value.get("materialization")),
        description=_optional_text(value.get("description")),
    )


def _parse_edge(value: object) -> DataAccessLineageEdge:
    if not isinstance(value, dict):
        raise DataAccessLineageCatalogError("Lineage catalog edge must be an object.")
    return DataAccessLineageEdge(
        upstream_node_id=_required_text(value, "upstream_node_id"),
        downstream_node_id=_required_text(value, "downstream_node_id"),
        relationship=_required_text(value, "relationship"),
    )


def _required_text(value: dict[str, Any], key: str) -> str:
    text = _optional_text(value.get(key))
    if text is None:
        raise DataAccessLineageCatalogError(f"Lineage catalog {key} is required.")
    return text


def _optional_text(value: object) -> str | None:
    text = str(value or "").strip()
    return text or None