"""Assemble reachable dbt source-dataset metadata into the DAX inventory.

The inventory is the authority for which active dbt result datasets Helix
offers for requests.  Starting only from those requestable result datasets,
the lineage catalog is traversed upstream until a dbt ``source`` relation is
reached.  Source nodes are terminal: a dataset that is itself only a source is
never used as a new traversal origin.

This keeps unrelated dbt source declarations out of Helix while preserving
manually governed inventory rows and their owners.  Rows previously derived by
this module are removed from the projection when they are no longer reachable;
the authoritative catalog feed then suspends their runtime rows.
"""

from __future__ import annotations

from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Final

DEFAULT_LINEAGE_ROOT_OWNER_EMAIL: Final = "xiaopeng.zhu@pingcap.com"
LINEAGE_ROOT_CATALOGED_BY: Final = "dbt_lineage"


@dataclass(frozen=True, slots=True)
class LineageRootDatasetAddition:
    project_id: str
    dataset_id: str
    owner_email: str


def sync_lineage_root_datasets(
    inventory: dict[str, Any],
    catalog: dict[str, Any],
    *,
    default_owner_email: str = DEFAULT_LINEAGE_ROOT_OWNER_EMAIL,
) -> tuple[dict[str, Any], tuple[LineageRootDatasetAddition, ...]]:
    """Make reachable dbt ``source`` datasets requestable in the inventory.

    Existing owners are never rewritten. A missing ``(project_id, dataset_id)``
    locator is assembled with the default owner; an existing active source row
    is promoted to requestable while its owner and lifecycle stay authoritative.
    Provenance labels on still-reachable rows previously assembled by this
    function are normalized. Stale derived rows are omitted, but manually
    cataloged rows are never removed.
    """
    owner_email = default_owner_email.strip().casefold()
    if not owner_email or "@" not in owner_email:
        raise ValueError("default_owner_email must be a non-empty email")

    datasets = inventory.get("datasets")
    nodes = catalog.get("nodes")
    edges = catalog.get("edges")
    if not isinstance(datasets, list):
        raise ValueError("dataset inventory must contain a datasets array")
    if not isinstance(nodes, list):
        raise ValueError("lineage catalog must contain a nodes array")
    if not isinstance(edges, list):
        raise ValueError("lineage catalog must contain an edges array")

    existing: set[tuple[str, str]] = set()
    for index, row in enumerate(datasets):
        if not isinstance(row, dict):
            raise ValueError(f"dataset inventory row #{index} must be an object")
        locator = (_required_text(row, "project_id"), _required_text(row, "dataset_id"))
        if locator in existing:
            raise ValueError(f"duplicate dataset inventory locator: {locator[0]}.{locator[1]}")
        existing.add(locator)

    target_locators = {
        (_required_text(row, "project_id"), _required_text(row, "dataset_id"))
        for row in datasets
        if (
            row.get("requestable") is True
            and str(row.get("status") or "active") == "active"
            and not _is_derived_lineage_root(row)
        )
    }
    source_locators, target_relation_count = _reachable_source_locators(
        nodes=nodes,
        edges=edges,
        target_locators=target_locators,
    )
    source_locator_set = set(source_locators)
    additions = tuple(
        LineageRootDatasetAddition(
            project_id=project_id,
            dataset_id=dataset_id,
            owner_email=owner_email,
        )
        for project_id, dataset_id in source_locators
        if (project_id, dataset_id) not in existing
    )

    projected = deepcopy(inventory)
    projected_datasets = [
        row
        for row in projected["datasets"]
        if (
            not _is_derived_lineage_root(row)
            or (
                _required_text(row, "project_id"),
                _required_text(row, "dataset_id"),
            )
            in source_locator_set
        )
    ]
    projected["datasets"] = projected_datasets
    for row in projected_datasets:
        locator = (
            _required_text(row, "project_id"),
            _required_text(row, "dataset_id"),
        )
        if locator in source_locator_set and str(row.get("status") or "active") == "active":
            row["requestable"] = True
        labels = row.get("labels")
        if not isinstance(labels, dict) or labels.get("lineage_root") != "true":
            continue
        labels.pop("managed_by", None)
        labels["cataloged_by"] = LINEAGE_ROOT_CATALOGED_BY
        labels["dbt_generated"] = "false"
        row["description"] = _source_requestable_description()
    projected_datasets.extend(_inventory_row(addition, catalog=catalog) for addition in additions)
    projected_datasets.sort(
        key=lambda row: (
            str(row["project_id"]).casefold(),
            str(row["dataset_id"]).casefold(),
        ),
    )

    metadata = projected.setdefault("metadata", {})
    if not isinstance(metadata, dict):
        raise ValueError("dataset inventory metadata must be an object")
    metadata["lineage_root_catalog_id"] = _required_text(catalog, "catalog_id")
    metadata["lineage_root_source_version"] = _required_text(catalog, "source_version")
    metadata["lineage_root_target_dataset_count"] = len(target_locators)
    metadata["lineage_root_target_relation_count"] = target_relation_count
    metadata["lineage_root_dataset_count"] = len(source_locators)
    metadata["lineage_root_default_owner_email"] = owner_email
    metadata["requestable_count"] = sum(
        1
        for row in projected_datasets
        if row.get("requestable") is True and str(row.get("status") or "active") == "active"
    )
    return projected, additions


def _reachable_source_locators(
    *,
    nodes: list[Any],
    edges: list[Any],
    target_locators: set[tuple[str, str]],
) -> tuple[list[tuple[str, str]], int]:
    nodes_by_id: dict[str, dict[str, Any]] = {}
    for index, node in enumerate(nodes):
        if not isinstance(node, dict):
            raise ValueError(f"lineage catalog node #{index} must be an object")
        node_id = _required_text(node, "node_id")
        if node_id in nodes_by_id:
            raise ValueError(f"duplicate lineage catalog node_id: {node_id}")
        nodes_by_id[node_id] = node

    parents_by_node_id: dict[str, set[str]] = {}
    for index, edge in enumerate(edges):
        if not isinstance(edge, dict):
            raise ValueError(f"lineage catalog edge #{index} must be an object")
        upstream_id = _required_text(edge, "upstream_node_id")
        downstream_id = _required_text(edge, "downstream_node_id")
        if upstream_id not in nodes_by_id or downstream_id not in nodes_by_id:
            raise ValueError(
                f"lineage catalog edge references an unknown node: {upstream_id}->{downstream_id}",
            )
        parents_by_node_id.setdefault(downstream_id, set()).add(upstream_id)

    target_node_ids = {
        node_id
        for node_id, node in nodes_by_id.items()
        if (
            str(node.get("node_type") or "") != "source"
            and (
                _required_text(node, "project_id"),
                _required_text(node, "dataset_id"),
            )
            in target_locators
        )
    }
    source_locators: set[tuple[str, str]] = set()
    visited: set[str] = set()
    stack = list(target_node_ids)
    while stack:
        node_id = stack.pop()
        if node_id in visited:
            continue
        visited.add(node_id)
        for upstream_id in parents_by_node_id.get(node_id, ()):
            upstream = nodes_by_id[upstream_id]
            if str(upstream.get("node_type") or "") == "source":
                source_locators.add(
                    (
                        _required_text(upstream, "project_id"),
                        _required_text(upstream, "dataset_id"),
                    )
                )
                continue
            stack.append(upstream_id)
    return sorted(source_locators), len(target_node_ids)


def _is_derived_lineage_root(row: dict[str, Any]) -> bool:
    labels = row.get("labels")
    return bool(
        isinstance(labels, dict)
        and labels.get("cataloged_by") == LINEAGE_ROOT_CATALOGED_BY
        and labels.get("lineage_root") == "true"
    )


def _inventory_row(
    addition: LineageRootDatasetAddition,
    *,
    catalog: dict[str, Any],
) -> dict[str, Any]:
    catalog_id = _required_text(catalog, "catalog_id")
    return {
        "project_id": addition.project_id,
        "dataset_id": addition.dataset_id,
        "owner_email": addition.owner_email,
        "display_name": _display_name(addition.dataset_id),
        "description": _source_requestable_description(),
        "labels": {
            "inventory": catalog_id,
            "cataloged_by": LINEAGE_ROOT_CATALOGED_BY,
            "dbt_generated": "false",
            "layer": "source",
            "lineage_root": "true",
        },
        "requestable": True,
        "status": "active",
    }


def _source_requestable_description() -> str:
    return "dbt 血缘引用的上游源数据集；已纳入 Helix，可按具体数据表申请访问。"


def _display_name(dataset_id: str) -> str:
    return " ".join(part.capitalize() for part in dataset_id.split("_") if part)


def _required_text(payload: dict[str, Any], field: str) -> str:
    value = str(payload.get(field) or "").strip()
    if not value:
        raise ValueError(f"{field} is required")
    return value