"""Lattice sub-client — workflow DAGs, graph traversal, and cross-workflow edges."""

from __future__ import annotations

from typing import Any, Mapping, Sequence

from sbn._http import HttpTransport

PREFIX = "/internal/lattice"

# ---------------------------------------------------------------------------
# Field normalization — accept common aliases from pgdag / sonic-pay
# ---------------------------------------------------------------------------

# Canonical top-level fields accepted by the backend NodeSchema / EdgeSchema
_NODE_FIELDS = {"step_id", "name", "skill", "dept", "description", "estimated_seconds", "meta"}
_EDGE_FIELDS = {"source_id", "target_id", "edge_type", "weight", "meta"}

# Short-form aliases used by pgdag and consumer DAG definitions
_NODE_ALIASES: dict[str, str] = {"id": "step_id", "label": "name"}
_EDGE_ALIASES: dict[str, str] = {"source": "source_id", "target": "target_id", "type": "edge_type"}


def _normalize_node(raw: Mapping[str, Any]) -> dict[str, Any]:
    """Map common node aliases to canonical field names.

    Unknown fields (e.g. ``step_type``, ``hard_stop``) are folded into
    the ``meta`` sub-dict so they survive the round-trip without breaking
    the backend's strict Pydantic validation.
    """
    node: dict[str, Any] = {}
    extras: dict[str, Any] = {}
    for k, v in raw.items():
        canonical = _NODE_ALIASES.get(k, k)
        if canonical in _NODE_FIELDS:
            node[canonical] = v
        else:
            extras[k] = v
    if extras:
        merged = dict(node.get("meta") or {})
        merged.update(extras)
        node["meta"] = merged
    return node


def _normalize_edge(raw: Mapping[str, Any]) -> dict[str, Any]:
    """Map common edge aliases to canonical field names."""
    edge: dict[str, Any] = {}
    extras: dict[str, Any] = {}
    for k, v in raw.items():
        canonical = _EDGE_ALIASES.get(k, k)
        if canonical in _EDGE_FIELDS:
            edge[canonical] = v
        else:
            extras[k] = v
    if extras:
        merged = dict(edge.get("meta") or {})
        merged.update(extras)
        edge["meta"] = merged
    return edge


class LatticeClient:
    """Workflow DAG registration, step orchestration, and graph queries."""

    def __init__(self, transport: HttpTransport) -> None:
        self._t = transport

    # ── Workflows ──────────────────────────────────────────────────────

    def create_workflow(
        self,
        *,
        domain: str,
        name: str,
        nodes: Sequence[Mapping[str, Any]],
        edges: Sequence[Mapping[str, Any]] | None = None,
        topology: str = "series",
        gec_frontier: Mapping[str, Any] | None = None,
        meta: Mapping[str, Any] | None = None,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {
            "domain": domain,
            "name": name,
            "nodes": [_normalize_node(n) for n in nodes],
            "edges": [_normalize_edge(e) for e in edges] if edges else [],
            "topology": topology,
        }
        if gec_frontier:
            body["gec_frontier"] = dict(gec_frontier)
        if meta:
            body["meta"] = dict(meta)
        return self._t.post(f"{PREFIX}/workflows", json=body).json()

    def list_workflows(
        self,
        *,
        domain: str | None = None,
        state: str | None = None,
        scopes: Sequence[str] | None = None,
        limit: int = 50,
    ) -> list[dict[str, Any]]:
        """List workflows with optional scope filtering.

        Parameters
        ----------
        scopes:
            Glob-style patterns like ``["finance.*", "treasury.*"]``.
        """
        params: dict[str, Any] = {"limit": limit}
        if domain:
            params["domain"] = domain
        if state:
            params["state"] = state
        if scopes:
            params["scopes"] = ",".join(scopes)
        data = self._t.get(f"{PREFIX}/workflows", params=params).json()
        return data.get("workflows", [])

    def discover_workflows(
        self,
        scopes: Sequence[str],
        *,
        limit: int = 50,
    ) -> list[dict[str, Any]]:
        """Discover active workflows from other projects matching scope patterns.

        Returns sanitized metadata (no full node/edge details).
        """
        params: dict[str, Any] = {
            "scopes": ",".join(scopes),
            "limit": limit,
        }
        data = self._t.get(f"{PREFIX}/workflows/discover", params=params).json()
        return data.get("workflows", [])

    def get_workflow(self, workflow_id: str) -> dict[str, Any]:
        return self._t.get(f"{PREFIX}/workflows/{workflow_id}").json()

    def activate_workflow(self, workflow_id: str) -> dict[str, Any]:
        return self._t.post(f"{PREFIX}/workflows/{workflow_id}/activate", json={}).json()

    def clone_workflow(
        self,
        workflow_id: str,
        target_project_id: str,
    ) -> dict[str, Any]:
        return self._t.post(
            f"{PREFIX}/workflows/{workflow_id}/clone",
            json={"target_project_id": target_project_id},
        ).json()

    # ── Export / Import ─────────────────────────────────────────────────

    def export_workflow(self, workflow_id: str) -> dict[str, Any]:
        """Export a workflow to a portable JSON dict."""
        return self._t.get(f"{PREFIX}/workflows/{workflow_id}/export").json()

    def export_workflows(
        self,
        *,
        workflow_ids: Sequence[str] | None = None,
        domain: str | None = None,
    ) -> dict[str, Any]:
        """Export multiple workflows as a portable bundle."""
        body: dict[str, Any] = {}
        if workflow_ids:
            body["workflow_ids"] = list(workflow_ids)
        if domain:
            body["domain"] = domain
        return self._t.post(f"{PREFIX}/workflows/export", json=body).json()

    def import_workflow(
        self,
        data: Mapping[str, Any],
        *,
        activate: bool = False,
    ) -> dict[str, Any]:
        """Import a workflow from a portable export dict."""
        payload = dict(data)
        # Normalise nodes/edges inside the workflow envelope if present
        wf = payload.get("workflow", payload)
        if "nodes" in wf:
            wf["nodes"] = [_normalize_node(n) for n in wf["nodes"]]
        if "edges" in wf:
            wf["edges"] = [_normalize_edge(e) for e in wf["edges"]]
        return self._t.post(
            f"{PREFIX}/workflows/import",
            json={"data": payload, "activate": activate},
        ).json()

    def import_workflows(
        self,
        data: Mapping[str, Any],
        *,
        activate: bool = False,
    ) -> dict[str, Any]:
        """Import a bundle of workflows from a bulk export."""
        payload = dict(data)
        for wf_entry in payload.get("workflows", []):
            wf = wf_entry.get("workflow", wf_entry)
            if "nodes" in wf:
                wf["nodes"] = [_normalize_node(n) for n in wf["nodes"]]
            if "edges" in wf:
                wf["edges"] = [_normalize_edge(e) for e in wf["edges"]]
        return self._t.post(
            f"{PREFIX}/workflows/import-bulk",
            json={"data": payload, "activate": activate},
        ).json()

    # ── Graph traversal ────────────────────────────────────────────────

    def get_execution_order(self, workflow_id: str) -> dict[str, Any]:
        """Get topological execution order and ready steps."""
        return self._t.get(f"{PREFIX}/workflows/{workflow_id}/order").json()

    def get_ready_steps(
        self,
        workflow_id: str,
        completed: Sequence[str] | None = None,
    ) -> list[dict[str, Any]]:
        """Get steps ready for execution given completed set."""
        params: dict[str, str] = {}
        if completed:
            params["completed"] = ",".join(completed)
        data = self._t.get(
            f"{PREFIX}/workflows/{workflow_id}/ready-steps", params=params
        ).json()
        return data.get("data", {}).get("ready_steps", data.get("ready_steps", []))

    # ── Cross-workflow edges ───────────────────────────────────────────

    def create_edge(
        self,
        *,
        source_id: str,
        target_id: str,
        edge_type: str,
        weight: float = 1.0,
        meta: Mapping[str, Any] | None = None,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {
            "source_id": source_id,
            "target_id": target_id,
            "edge_type": edge_type,
            "weight": weight,
        }
        if meta:
            body["meta"] = dict(meta)
        return self._t.post(f"{PREFIX}/edges", json=body).json()

    def list_edges(
        self,
        *,
        source_id: str | None = None,
        target_id: str | None = None,
        edge_type: str | None = None,
        limit: int = 100,
    ) -> list[dict[str, Any]]:
        params: dict[str, Any] = {"limit": limit}
        if source_id:
            params["source_id"] = source_id
        if target_id:
            params["target_id"] = target_id
        if edge_type:
            params["edge_type"] = edge_type
        data = self._t.get(f"{PREFIX}/edges", params=params).json()
        return data.get("edges", [])

    # ── Step outcomes & analytics ──────────────────────────────────────

    def record_step_outcome(
        self,
        workflow_id: str,
        step_id: str,
        *,
        agent_id: str | None = None,
        outcome: str = "completed",
        duration_seconds: float | None = None,
        error_message: str | None = None,
        meta: Mapping[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Record step execution outcome."""
        body: dict[str, Any] = {"outcome": outcome}
        if agent_id:
            body["agent_id"] = agent_id
        if duration_seconds is not None:
            body["duration_seconds"] = duration_seconds
        if error_message:
            body["error_message"] = error_message
        if meta:
            body["meta"] = dict(meta)
        return self._t.post(
            f"{PREFIX}/workflows/{workflow_id}/steps/{step_id}/outcome",
            json=body,
        ).json()

    def mark_workflow_complete(self, workflow_id: str) -> dict[str, Any]:
        """Mark workflow run as complete."""
        return self._t.post(
            f"{PREFIX}/workflows/{workflow_id}/complete", json={},
        ).json()

    def get_workflow_analytics(self, workflow_id: str) -> dict[str, Any]:
        """Get aggregated analytics for workflow execution."""
        return self._t.get(f"{PREFIX}/workflows/{workflow_id}/analytics").json()

    def get_workflow_benchmarks(self, workflow_id: str) -> dict[str, Any]:
        """Get cross-project benchmark data for workflow execution."""
        return self._t.get(f"{PREFIX}/workflows/{workflow_id}/benchmarks").json()
