from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from govoplan_dataflow.backend.expressions import (
    ExpressionError,
    infer_expression_type,
    parse_expression,
)
from govoplan_dataflow.backend.operator_registry import OPERATOR_REGISTRY
from govoplan_dataflow.backend.schemas import (
    DataflowDiagnostic,
    GraphNode,
    PipelineGraph,
)


@dataclass(frozen=True, slots=True)
class SchemaState:
    columns: frozenset[str]
    open: bool = False
    types: dict[str, str] = field(default_factory=dict)

    def knows(self, column: str) -> bool:
        return self.open or column in self.columns

    def type_of(self, column: str) -> str:
        return self.types.get(column, "unknown")


@dataclass(frozen=True, slots=True)
class SchemaPropagationContext:
    node: GraphNode
    input_state: SchemaState
    input_states: tuple[SchemaState, ...]
    inputs_by_port: dict[str, list[str]]
    schemas: dict[str, SchemaState]

    def port_state(self, port: str) -> SchemaState:
        source_ids = self.inputs_by_port.get(port, ())
        if not source_ids:
            return unknown_schema()
        return self.schemas.get(source_ids[0], unknown_schema())


@dataclass(frozen=True, slots=True)
class SchemaPropagationResult:
    state: SchemaState
    diagnostics: tuple[DataflowDiagnostic, ...] = ()


def validate_graph_schemas(
    graph: PipelineGraph,
    *,
    ordered: list[str],
) -> list[DataflowDiagnostic]:
    _, diagnostics = propagate_graph_schemas(graph, ordered=ordered)
    return diagnostics


def propagate_graph_schemas(
    graph: PipelineGraph,
    *,
    ordered: list[str],
) -> tuple[dict[str, SchemaState], list[DataflowDiagnostic]]:
    node_by_id = {node.id: node for node in graph.nodes}
    inputs = _graph_inputs_by_port(graph)
    schemas: dict[str, SchemaState] = {}
    diagnostics: list[DataflowDiagnostic] = []
    for node_id in ordered:
        node = node_by_id[node_id]
        context = _propagation_context(node, inputs, schemas)
        propagator = OPERATOR_REGISTRY.schema_propagator(node.type)
        if propagator is None:
            diagnostics.append(
                _error(
                    "node.schema_propagator_missing",
                    f"Node type {node.type!r} has no schema propagator.",
                    node_id=node.id,
                )
            )
            schemas[node.id] = context.input_state
            continue
        result = propagator(context)
        if not isinstance(result, SchemaPropagationResult):
            raise TypeError(
                f"Schema propagator for {node.type!r} returned "
                f"{type(result).__name__}, not SchemaPropagationResult."
            )
        schemas[node.id] = result.state
        diagnostics.extend(result.diagnostics)
    return schemas, diagnostics


def _graph_inputs_by_port(
    graph: PipelineGraph,
) -> dict[str, dict[str, list[str]]]:
    result: dict[str, dict[str, list[str]]] = {}
    for edge in graph.edges:
        result.setdefault(edge.target, {}).setdefault(
            edge.target_port,
            [],
        ).append(edge.source)
    return result


def _propagation_context(
    node: GraphNode,
    inputs: dict[str, dict[str, list[str]]],
    schemas: dict[str, SchemaState],
) -> SchemaPropagationContext:
    node_inputs = inputs.get(node.id, {})
    input_states = tuple(
        schemas[source_id]
        for port_sources in node_inputs.values()
        for source_id in port_sources
        if source_id in schemas
    )
    return SchemaPropagationContext(
        node=node,
        input_state=input_states[0] if input_states else unknown_schema(),
        input_states=input_states,
        inputs_by_port=node_inputs,
        schemas=schemas,
    )


def _inline_source(
    context: SchemaPropagationContext,
) -> SchemaPropagationResult:
    return SchemaPropagationResult(_inline_schema(context.node.config.get("rows")))


def _reference_source(
    context: SchemaPropagationContext,
) -> SchemaPropagationResult:
    return SchemaPropagationResult(
        _configured_schema(context.node.config.get("source_columns"))
    )


def _union(context: SchemaPropagationContext) -> SchemaPropagationResult:
    diagnostics: list[DataflowDiagnostic] = []
    closed_shapes = {
        state.columns
        for state in context.input_states
        if not state.open
    }
    if len(closed_shapes) > 1:
        diagnostics.append(
            _warning(
                "union.schema_mismatch",
                "Appended inputs use different columns; missing values will be null.",
                node_id=context.node.id,
            )
        )
    return SchemaPropagationResult(
        SchemaState(
            frozenset().union(
                *(state.columns for state in context.input_states)
            ),
            open=any(state.open for state in context.input_states),
            types=_merged_schema_types(context.input_states),
        ),
        tuple(diagnostics),
    )


def _join(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    left_state = context.port_state("left")
    right_state = context.port_state("right")
    diagnostics = [
        *_unknown_columns(
            node,
            left_state,
            node.config.get("left_keys"),
            field_name="left_keys",
        ),
        *_unknown_columns(
            node,
            right_state,
            node.config.get("right_keys"),
            field_name="right_keys",
        ),
    ]
    if node.config.get("join_type", "inner") in {"semi", "anti"}:
        return SchemaPropagationResult(left_state, tuple(diagnostics))
    prefix = str(node.config.get("right_prefix", "right_"))
    prefixed_right = {
        f"{prefix}{column}"
        for column in right_state.columns
    }
    collisions = left_state.columns & prefixed_right
    if collisions:
        diagnostics.append(
            _error(
                "join.output_collision",
                f"Join output columns collide: {', '.join(sorted(collisions))}.",
                node_id=node.id,
                field="right_prefix",
            )
        )
    return SchemaPropagationResult(
        SchemaState(
            frozenset(left_state.columns | prefixed_right),
            open=left_state.open or right_state.open,
            types={
                **left_state.types,
                **{
                    f"{prefix}{column}": right_state.type_of(column)
                    for column in right_state.columns
                },
            },
        ),
        tuple(diagnostics),
    )


def _filter(context: SchemaPropagationContext) -> SchemaPropagationResult:
    return _passthrough_with_columns(
        context,
        [context.node.config.get("column")],
        field_name="column",
    )


def _filter_expression(
    context: SchemaPropagationContext,
) -> SchemaPropagationResult:
    parsed, diagnostics = _node_expression(context.node, "expression")
    if parsed is not None:
        diagnostics.extend(
            _unknown_columns(
                context.node,
                context.input_state,
                list(parsed.columns),
                field_name="expression",
            )
        )
    return SchemaPropagationResult(
        context.input_state,
        tuple(diagnostics),
    )


def _distinct(context: SchemaPropagationContext) -> SchemaPropagationResult:
    return _passthrough_with_columns(
        context,
        context.node.config.get("columns"),
        field_name="columns",
    )


def _select(context: SchemaPropagationContext) -> SchemaPropagationResult:
    selected, output = _selected_columns(context.node.config.get("fields"))
    diagnostics = [
        *_unknown_columns(
            context.node,
            context.input_state,
            selected,
            field_name="fields",
        ),
        *_duplicate_outputs(
            context.node,
            output,
            field_name="fields",
        ),
    ]
    return SchemaPropagationResult(
        SchemaState(
            frozenset(output),
            types={
                target: context.input_state.type_of(source)
                for source, target in zip(selected, output, strict=False)
            },
        ),
        tuple(diagnostics),
    )


def _selected_columns(value: object) -> tuple[list[str], list[str]]:
    selected: list[str] = []
    output: list[str] = []
    if not isinstance(value, list):
        return selected, output
    for item in value:
        if isinstance(item, str):
            selected.append(item)
            output.append(item)
        elif isinstance(item, dict):
            source = item.get("column")
            target = item.get("alias") or source
            if isinstance(source, str):
                selected.append(source)
            if isinstance(target, str):
                output.append(target)
    return selected, output


def _derive(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    source_columns = _text_items(node.config.get("source_columns"))
    diagnostics = _unknown_columns(
        node,
        context.input_state,
        source_columns,
        field_name="source_columns",
    )
    target = node.config.get("target_column")
    if not isinstance(target, str) or not target:
        return SchemaPropagationResult(
            context.input_state,
            tuple(diagnostics),
        )
    if target in context.input_state.columns:
        diagnostics.append(
            _warning(
                "derive.overwrites_column",
                f"Derived column {target!r} replaces an existing value.",
                node_id=node.id,
                field="target_column",
            )
        )
    result_type = _derive_result_type(
        str(node.config.get("operation") or ""),
        [
            context.input_state.type_of(column)
            for column in source_columns
        ],
    )
    return SchemaPropagationResult(
        _with_column(context.input_state, target, result_type),
        tuple(diagnostics),
    )


def _expression(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    parsed, diagnostics = _node_expression(node, "expression")
    inferred = "unknown"
    if parsed is not None:
        diagnostics.extend(
            _unknown_columns(
                node,
                context.input_state,
                list(parsed.columns),
                field_name="expression",
            )
        )
        inferred = infer_expression_type(
            parsed,
            {
                name: context.input_state.type_of(name)  # type: ignore[dict-item]
                for name in context.input_state.columns
            },
        )
    expected = str(node.config.get("result_type") or "unknown")
    if expected != "unknown" and inferred not in {"unknown", "null", expected}:
        diagnostics.append(
            _warning(
                "expression.type_mismatch",
                f"Expression infers {inferred}, not {expected}.",
                node_id=node.id,
                field="result_type",
            )
        )
    target = str(node.config.get("target_column") or "")
    state = (
        _with_column(
            context.input_state,
            target,
            expected if expected != "unknown" else inferred,
        )
        if target
        else context.input_state
    )
    return SchemaPropagationResult(state, tuple(diagnostics))


def _calculate(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    state = context.input_state
    diagnostics: list[DataflowDiagnostic] = []
    for item in _mapping_items(node.config.get("calculations")):
        source = str(item.get("expression") or "")
        target = str(item.get("target_column") or "")
        try:
            parsed = parse_expression(source)
        except ExpressionError as exc:
            diagnostics.append(
                _error(
                    "expression.invalid",
                    str(exc),
                    node_id=node.id,
                    field="calculations",
                )
            )
            continue
        diagnostics.extend(
            _unknown_columns(
                node,
                state,
                list(parsed.columns),
                field_name="calculations",
            )
        )
        inferred = infer_expression_type(
            parsed,
            {
                name: state.type_of(name)  # type: ignore[dict-item]
                for name in state.columns
            },
        )
        expected = str(item.get("result_type") or "unknown")
        if expected != "unknown" and inferred not in {
            "unknown",
            "null",
            expected,
        }:
            diagnostics.append(
                _warning(
                    "calculate.type_mismatch",
                    (
                        f"Calculation for {target!r} infers {inferred}, "
                        f"not {expected}."
                    ),
                    node_id=node.id,
                    field="calculations",
                )
            )
        if target:
            state = _with_column(
                state,
                target,
                expected if expected != "unknown" else inferred,
            )
    return SchemaPropagationResult(state, tuple(diagnostics))


def _convert_or_replace(
    context: SchemaPropagationContext,
) -> SchemaPropagationResult:
    node = context.node
    source = str(node.config.get("source_column") or "")
    target = str(node.config.get("target_column") or "")
    diagnostics = _unknown_columns(
        node,
        context.input_state,
        [source],
        field_name="source_column",
    )
    target_type = (
        str(node.config.get("target_type") or "unknown")
        if node.type == "convert"
        else context.input_state.type_of(source)
    )
    state = (
        _with_column(context.input_state, target, target_type)
        if target
        else context.input_state
    )
    return SchemaPropagationResult(state, tuple(diagnostics))


def _aggregate(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    group_by = _text_items(node.config.get("group_by"))
    aggregates = _mapping_items(node.config.get("aggregates"))
    source_columns = [
        str(item.get("column"))
        for item in aggregates
        if item.get("column") not in (None, "", "*")
    ]
    aliases = [
        str(item.get("alias"))
        for item in aggregates
        if item.get("alias")
    ]
    output = [*group_by, *aliases]
    diagnostics = [
        *_unknown_columns(
            node,
            context.input_state,
            [*group_by, *source_columns],
            field_name="aggregates",
        ),
        *_duplicate_outputs(node, output, field_name="aggregates"),
    ]
    aggregate_types = {
        str(item["alias"]): (
            "integer"
            if item.get("function") == "count"
            else context.input_state.type_of(str(item.get("column") or ""))
        )
        for item in aggregates
        if item.get("alias")
    }
    return SchemaPropagationResult(
        SchemaState(
            frozenset(output),
            types={
                **{
                    column: context.input_state.type_of(column)
                    for column in group_by
                },
                **aggregate_types,
            },
        ),
        tuple(diagnostics),
    )


def _sort(context: SchemaPropagationContext) -> SchemaPropagationResult:
    columns = [
        str(item.get("column"))
        for item in _mapping_items(context.node.config.get("fields"))
        if item.get("column")
    ]
    return _passthrough_with_columns(
        context,
        columns,
        field_name="fields",
    )


def _rank(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    columns = [
        *_text_items(node.config.get("partition_by")),
        *[
            str(item.get("column"))
            for item in _mapping_items(node.config.get("order_by"))
            if item.get("column")
        ],
    ]
    diagnostics = _unknown_columns(
        node,
        context.input_state,
        columns,
        field_name="order_by",
    )
    target = str(node.config.get("target_column") or "")
    state = (
        _with_column(context.input_state, target, "integer")
        if target
        else context.input_state
    )
    return SchemaPropagationResult(state, tuple(diagnostics))


def _quality(context: SchemaPropagationContext) -> SchemaPropagationResult:
    rules = _mapping_items(context.node.config.get("rules"))
    diagnostics = _unknown_columns(
        context.node,
        context.input_state,
        [str(rule.get("column") or "") for rule in rules],
        field_name="rules",
    )
    if context.node.config.get("action", "annotate") != "annotate":
        return SchemaPropagationResult(
            context.input_state,
            tuple(diagnostics),
        )
    state = _with_column(context.input_state, "_quality_valid", "boolean")
    state = _with_column(state, "_quality_errors", "array")
    return SchemaPropagationResult(state, tuple(diagnostics))


def _reconcile(context: SchemaPropagationContext) -> SchemaPropagationResult:
    node = context.node
    left_state = context.port_state("left")
    right_state = context.port_state("right")
    left_compare, right_compare = _comparison_columns(
        node.config.get("compare_columns")
    )
    diagnostics = [
        *_unknown_columns(
            node,
            left_state,
            node.config.get("left_keys"),
            field_name="left_keys",
        ),
        *_unknown_columns(
            node,
            right_state,
            node.config.get("right_keys"),
            field_name="right_keys",
        ),
        *_unknown_columns(
            node,
            left_state,
            left_compare,
            field_name="compare_columns",
        ),
        *_unknown_columns(
            node,
            right_state,
            right_compare,
            field_name="compare_columns",
        ),
    ]
    prefix = str(node.config.get("right_prefix") or "observed_")
    state = SchemaState(
        left_state.columns
        | frozenset(
            f"{prefix}{column}"
            for column in right_state.columns
        )
        | frozenset(
            (
                "_reconciliation_status",
                "_reconciliation_differences",
            )
        ),
        open=left_state.open or right_state.open,
        types={
            **left_state.types,
            **{
                f"{prefix}{column}": right_state.type_of(column)
                for column in right_state.columns
            },
            "_reconciliation_status": "string",
            "_reconciliation_differences": "array",
        },
    )
    return SchemaPropagationResult(state, tuple(diagnostics))


def _comparison_columns(value: object) -> tuple[list[str], list[str]]:
    left: list[str] = []
    right: list[str] = []
    if not isinstance(value, list):
        return left, right
    for item in value:
        if isinstance(item, str):
            left.append(item)
            right.append(item)
        elif isinstance(item, dict):
            left_name = str(item.get("left") or item.get("column") or "")
            right_name = str(
                item.get("right")
                or item.get("left")
                or item.get("column")
                or ""
            )
            left.append(left_name)
            right.append(right_name)
    return left, right


def _subflow(context: SchemaPropagationContext) -> SchemaPropagationResult:
    output_schema = _configured_schema(
        context.node.config.get("output_schema")
    )
    return SchemaPropagationResult(
        output_schema
        if output_schema.columns
        else unknown_schema()
    )


def _identity(context: SchemaPropagationContext) -> SchemaPropagationResult:
    return SchemaPropagationResult(context.input_state)


def _passthrough_with_columns(
    context: SchemaPropagationContext,
    columns: object,
    *,
    field_name: str,
) -> SchemaPropagationResult:
    diagnostics = _unknown_columns(
        context.node,
        context.input_state,
        columns,
        field_name=field_name,
    )
    return SchemaPropagationResult(
        context.input_state,
        tuple(diagnostics),
    )


def _with_column(
    state: SchemaState,
    name: str,
    data_type: str,
) -> SchemaState:
    return SchemaState(
        state.columns | frozenset((name,)),
        open=state.open,
        types={**state.types, name: data_type},
    )


def _configured_schema(value: object) -> SchemaState:
    if not isinstance(value, list):
        return unknown_schema()
    columns = {
        item
        if isinstance(item, str)
        else str(item.get("name"))
        for item in value
        if (
            isinstance(item, str)
            and item
            or isinstance(item, dict)
            and item.get("name")
        )
    }
    types = {
        str(item["name"]): str(
            item.get("data_type")
            or item.get("type")
            or "unknown"
        )
        for item in value
        if isinstance(item, dict) and item.get("name")
    }
    return SchemaState(
        frozenset(columns),
        open=not columns,
        types=types,
    )


def _inline_schema(value: object) -> SchemaState:
    if not isinstance(value, list):
        return unknown_schema()
    rows = [row for row in value if isinstance(row, dict)]
    columns = {
        str(column)
        for row in rows
        for column in row
    }
    types = {
        column: _observed_column_type(rows, column)
        for column in columns
    }
    return SchemaState(
        frozenset(columns),
        open=not columns,
        types=types,
    )


def _observed_column_type(
    rows: list[dict[str, Any]],
    column: str,
) -> str:
    observed = {
        _schema_value_type(row.get(column))
        for row in rows
        if row.get(column) is not None
    }
    return observed.pop() if len(observed) == 1 else "unknown"


def _schema_value_type(value: object) -> str:
    type_checks = (
        (bool, "boolean"),
        (int, "integer"),
        (float, "number"),
        (str, "string"),
        (list, "array"),
        (dict, "object"),
    )
    return next(
        (
            name
            for value_type, name in type_checks
            if isinstance(value, value_type)
        ),
        "unknown",
    )


def _merged_schema_types(
    states: tuple[SchemaState, ...],
) -> dict[str, str]:
    columns = frozenset().union(*(state.columns for state in states))
    return {
        column: _merged_column_type(states, column)
        for column in columns
    }


def _merged_column_type(
    states: tuple[SchemaState, ...],
    column: str,
) -> str:
    observed = {
        state.type_of(column)
        for state in states
        if (
            column in state.columns
            and state.type_of(column) != "unknown"
        )
    }
    return observed.pop() if len(observed) == 1 else "unknown"


def _derive_result_type(
    operation: str,
    source_types: list[str],
) -> str:
    if operation in {"upper", "lower", "trim", "concat"}:
        return "string"
    if operation in {"add", "subtract", "multiply", "divide"}:
        return (
            "number"
            if "number" in source_types or operation == "divide"
            else "integer"
        )
    if operation in {"copy", "coalesce"}:
        concrete = {
            item
            for item in source_types
            if item not in {"unknown", "null"}
        }
        return (
            concrete.pop()
            if len(concrete) == 1
            else "unknown"
        )
    return "unknown"


def _node_expression(
    node: GraphNode,
    field_name: str,
) -> tuple[object | None, list[DataflowDiagnostic]]:
    try:
        return (
            parse_expression(str(node.config.get(field_name) or "")),
            [],
        )
    except ExpressionError as exc:
        return (
            None,
            [
                _error(
                    "expression.invalid",
                    str(exc),
                    node_id=node.id,
                    field=field_name,
                )
            ],
        )


def _unknown_columns(
    node: GraphNode,
    state: SchemaState,
    columns: object,
    *,
    field_name: str,
) -> list[DataflowDiagnostic]:
    return [
        _error(
            "schema.unknown_column",
            f"Column {column!r} is not available at this node.",
            node_id=node.id,
            field=field_name,
        )
        for column in _text_items(columns)
        if not state.knows(column)
    ]


def _duplicate_outputs(
    node: GraphNode,
    columns: list[str],
    *,
    field_name: str,
) -> list[DataflowDiagnostic]:
    duplicates = sorted(
        column
        for column in set(columns)
        if columns.count(column) > 1
    )
    if not duplicates:
        return []
    return [
        _error(
            "schema.duplicate_output",
            f"Output column names must be unique: {', '.join(duplicates)}.",
            node_id=node.id,
            field=field_name,
        )
    ]


def _text_items(value: object) -> list[str]:
    if not isinstance(value, list):
        return []
    return [
        item
        for item in value
        if isinstance(item, str) and item
    ]


def _mapping_items(value: object) -> list[dict[str, Any]]:
    if not isinstance(value, list):
        return []
    return [
        item
        for item in value
        if isinstance(item, dict)
    ]


def unknown_schema() -> SchemaState:
    return SchemaState(frozenset(), open=True)


def _error(
    code: str,
    message: str,
    *,
    node_id: str | None = None,
    field: str | None = None,
) -> DataflowDiagnostic:
    return DataflowDiagnostic(
        severity="error",
        code=code,
        message=message,
        node_id=node_id,
        field=field,
    )


def _warning(
    code: str,
    message: str,
    *,
    node_id: str | None = None,
    field: str | None = None,
) -> DataflowDiagnostic:
    return DataflowDiagnostic(
        severity="warning",
        code=code,
        message=message,
        node_id=node_id,
        field=field,
    )


def register_schema_propagators() -> None:
    propagators = {
        "source.inline": _inline_source,
        "source.reference": _reference_source,
        "combine.union": _union,
        "combine.join": _join,
        "filter": _filter,
        "filter.expression": _filter_expression,
        "distinct": _distinct,
        "select": _select,
        "derive": _derive,
        "expression": _expression,
        "calculate": _calculate,
        "convert": _convert_or_replace,
        "replace": _convert_or_replace,
        "aggregate": _aggregate,
        "sort": _sort,
        "window.rank": _rank,
        "limit": _identity,
        "quality.rules": _quality,
        "reconcile.compare": _reconcile,
        "subflow": _subflow,
        "output": _identity,
    }
    for node_type, propagator in propagators.items():
        if OPERATOR_REGISTRY.schema_propagator(node_type) is None:
            OPERATOR_REGISTRY.register_schema_propagator(
                node_type,
                propagator,
            )


__all__ = [
    "SchemaPropagationContext",
    "SchemaPropagationResult",
    "SchemaState",
    "register_schema_propagators",
    "validate_graph_schemas",
]