#!/usr/bin/env python3
"""Exercise connector -> datasource -> dataflow publication capabilities."""

from __future__ import annotations

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from govoplan_connectors.backend.db.models import ConnectorTabularSource
from govoplan_core.auth import ApiPrincipal
from govoplan_core.core.access import PrincipalRef
from govoplan_core.core.dataflows import (
    DataflowPublicationTarget,
    DataflowRunRequest,
    dataflow_run_lifecycle,
)
from govoplan_core.core.datasources import (
    DatasourceReadRequest,
    datasource_catalogue,
    datasource_lifecycle,
    datasource_publication,
)
from govoplan_core.core.modules import ModuleContext
from govoplan_core.core.tabular_sources import (
    TabularSnapshotInput,
    tabular_snapshot_writer,
)
from govoplan_core.db.base import Base
from govoplan_core.server.registry import build_platform_registry
from govoplan_dataflow.backend.schemas import (
    GraphEdge,
    GraphNode,
    GraphPosition,
    PipelineGraph,
    PipelineCreateRequest,
    PipelinePreviewRequest,
)
from govoplan_dataflow.backend.db.models import (
    DataflowPipeline,
    DataflowPipelineRevision,
    DataflowRun,
)
from govoplan_dataflow.backend.service import create_pipeline, preview_pipeline
from govoplan_datasources.backend.db.models import (
    DatasourceMaterializationRecord,
    DatasourcePayloadRecord,
    DatasourcePayloadRowRecord,
    DatasourcePublicationRecord,
    DatasourceRecord,
    DatasourceStageRecord,
)


def main() -> int:
    registry = build_platform_registry(
        ("connectors", "datasources", "dataflow", "workflow")
    )
    registry.configure_capability_context(
        ModuleContext(registry=registry, settings=object())
    )
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(
        engine,
        tables=[
            ConnectorTabularSource.__table__,
            DatasourceRecord.__table__,
            DatasourcePayloadRecord.__table__,
            DatasourcePayloadRowRecord.__table__,
            DatasourceMaterializationRecord.__table__,
            DatasourceStageRecord.__table__,
            DatasourcePublicationRecord.__table__,
            DataflowPipeline.__table__,
            DataflowPipelineRevision.__table__,
            DataflowRun.__table__,
        ],
    )
    session_factory = sessionmaker(bind=engine)
    with session_factory() as session:
        principal = _principal()
        writer = tabular_snapshot_writer(registry)
        lifecycle = datasource_lifecycle(registry)
        catalogue = datasource_catalogue(registry)
        publisher = datasource_publication(registry)
        runner = dataflow_run_lifecycle(registry)
        if (
            writer is None
            or lifecycle is None
            or catalogue is None
            or publisher is None
            or runner is None
        ):
            raise RuntimeError("Datasource composition capabilities are incomplete.")

        origin = writer.create_snapshot(
            session,
            principal,
            snapshot=TabularSnapshotInput(
                name="Monthly cases",
                source_name="connector_monthly_cases",
                rows=(
                    {"id": 1, "amount": 5},
                    {"id": 2, "amount": 15},
                ),
            ),
        )
        datasource = lifecycle.register_origin(
            session,
            principal,
            origin_ref=origin.ref,
            name="Monthly cases cache",
            source_name="monthly_cases",
            mode="cached",
        )
        result = preview_pipeline(
            session,
            tenant_id="tenant-1",
            actor_id="account-1",
            payload=PipelinePreviewRequest(
                graph=_graph(
                    datasource_ref=datasource.ref,
                    fingerprint=datasource.fingerprint,
                ),
                row_limit=100,
            ),
            principal=principal,
            registry=registry,
        )
        expected_rows = [
            {"id": 1, "amount": 5},
            {"id": 2, "amount": 15},
        ]
        if result.status != "succeeded":
            raise RuntimeError(f"Dataflow preview failed: {result.diagnostics}")
        if result.rows != expected_rows:
            raise RuntimeError(f"Unexpected Dataflow rows: {result.rows!r}")
        if result.source_fingerprints[0]["source_ref"] != datasource.ref:
            raise RuntimeError("Dataflow lineage did not retain the datasource reference.")
        pipeline = create_pipeline(
            session,
            tenant_id="tenant-1",
            actor_id="account-1",
            payload=PipelineCreateRequest(
                name="Monthly case output",
                status="active",
                graph=_graph(
                    datasource_ref=datasource.ref,
                    fingerprint=datasource.fingerprint,
                ),
                editor_mode="graph",
            ),
        )
        run_request = DataflowRunRequest(
            pipeline_ref=f"pipeline:{pipeline.id}",
            revision=1,
            idempotency_key="composition-run-1",
            publication=DataflowPublicationTarget(
                name="Monthly case result",
                source_name="monthly_case_result",
                freeze=True,
                frozen_label="Composition evidence",
            ),
        )
        published = runner.start_run(
            session,
            principal,
            request=run_request,
        )
        replayed = runner.start_run(
            session,
            principal,
            request=run_request,
        )
        if published.status != "succeeded":
            raise RuntimeError(f"Dataflow publication failed: {published.error}")
        if replayed.ref != published.ref or not replayed.replayed:
            raise RuntimeError("Dataflow run idempotency did not replay the prior run.")
        if (
            not published.output_datasource_ref
            or not published.output_materialization_ref
        ):
            raise RuntimeError("Dataflow publication did not retain output references.")
        output = catalogue.read_datasource(
            session,
            principal,
            request=DatasourceReadRequest(
                datasource_ref=published.output_datasource_ref,
            ),
        )
        if list(output.rows) != expected_rows:
            raise RuntimeError(
                f"Unexpected published Dataflow rows: {list(output.rows)!r}"
            )
        if (
            output.materialization is None
            or output.materialization.ref != published.output_materialization_ref
            or output.materialization.frozen_at is None
        ):
            raise RuntimeError(
                "Published Datasource materialization is not pinned and frozen."
            )
    engine.dispose()
    print(
        "Connector -> Datasources -> pinned Dataflow publication composition passed."
    )
    return 0


def _principal() -> ApiPrincipal:
    return ApiPrincipal(
        principal=PrincipalRef(
            account_id="account-1",
            membership_id="membership-1",
            tenant_id="tenant-1",
            scopes=frozenset(
                {
                    "connectors:source:read",
                    "connectors:source:write",
                    "datasources:catalogue:read",
                    "datasources:source:write",
                    "datasources:stage:write",
                    "dataflow:pipeline:run",
                }
            ),
        ),
        account=object(),
        user=object(),
    )


def _graph(*, datasource_ref: str, fingerprint: str) -> PipelineGraph:
    return PipelineGraph(
        nodes=[
            GraphNode(
                id="source",
                type="source.reference",
                label="Cases",
                position=GraphPosition(x=0, y=0),
                config={
                    "source_ref": datasource_ref,
                    "source_name": "monthly_cases",
                    "expected_fingerprint": fingerprint,
                    "consistency": "current",
                },
            ),
            GraphNode(
                id="output",
                type="output",
                label="Output",
                position=GraphPosition(x=200, y=0),
                config={},
            ),
        ],
        edges=[
            GraphEdge(
                id="source-output",
                source="source",
                target="output",
            )
        ],
    )


if __name__ == "__main__":
    raise SystemExit(main())