#!/usr/bin/env python3
"""Exercise the connector -> datasource -> dataflow capability path."""

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.datasources import datasource_catalogue, datasource_lifecycle
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,
    PipelinePreviewRequest,
)
from govoplan_dataflow.backend.service import preview_pipeline
from govoplan_datasources.backend.db.models import (
    DatasourceMaterializationRecord,
    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__,
            DatasourceMaterializationRecord.__table__,
            DatasourceStageRecord.__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)
        if writer is None or lifecycle is None or catalogue 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.")
    engine.dispose()
    print("Connector -> Datasources -> Dataflow 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())