"""Workflow graph executor (build plan Section 8.3).

Executes a workflow's nodes in topological order, following edges and condition
branches, persisting each node as a :class:`Step`. Approval and sensitive-tool
nodes pause the run (creating an :class:`Approval`); :meth:`resume` continues once
a human decides. Failed nodes retry with backoff per their node policy.
"""

from __future__ import annotations

import json
import time
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta

from legate.agent.enums import ApprovalStatus, RunStatus, StepType
from legate.agent.events import run_event_bus
from legate.agent.runner import _step_event
from legate.agent.tooling import ToolContext, ToolRegistry, ToolResult
from legate.audit.service import record as audit_record
from legate.config import get_settings
from legate.db.models.approval import Approval
from legate.db.models.run import Run
from legate.db.models.step import Step
from legate.db.models.workflow import Workflow
from legate.observability.logging import get_logger
from legate.policy.engine import ApprovalRules, PolicyEffect, evaluate_tool_call
from legate.workflow.templating import evaluate_condition, render_value

logger = get_logger(__name__)

_NODE_STEP_TYPE = {
    "trigger": StepType.TRIGGER,
    "agent": StepType.AGENT_REASON,
    "tool": StepType.TOOL_CALL,
    "condition": StepType.CONDITION,
    "transform": StepType.TRANSFORM,
    "approval": StepType.APPROVAL,
    "delay": StepType.TRANSFORM,
}


class WorkflowGraphError(ValueError):
    """Raised when a workflow graph is malformed (e.g. contains a cycle)."""


class NodeError(Exception):
    """A retryable node execution failure."""


@dataclass
class _Pause:
    kind: str  # "approval_node" | "tool_node"
    summary: str = ""
    tool: str | None = None
    args: dict = field(default_factory=dict)


@dataclass
class _NodeOutcome:
    result: dict | None = None
    pause: _Pause | None = None
    failed: bool = False
    error: str = ""


def _topo_order(nodes: list[dict], edges: list[dict]) -> list[str]:
    ids: list[str] = [str(n["id"]) for n in nodes]
    indegree: dict[str, int] = dict.fromkeys(ids, 0)
    adj: dict[str, list[str]] = {nid: [] for nid in ids}
    for edge in edges:
        src, dst = edge.get("from"), edge.get("to")
        if isinstance(src, str) and isinstance(dst, str) and src in adj and dst in indegree:
            adj[src].append(dst)
            indegree[dst] += 1
    queue = [nid for nid in ids if indegree[nid] == 0]
    order: list[str] = []
    while queue:
        nid = queue.pop(0)
        order.append(nid)
        for nxt in adj[nid]:
            indegree[nxt] -= 1
            if indegree[nxt] == 0:
                queue.append(nxt)
    if len(order) != len(ids):
        raise WorkflowGraphError("workflow graph contains a cycle")
    return order


class WorkflowExecutor:
    def __init__(self, *, db, settings=None, bus=None) -> None:
        self.db = db
        self.settings = settings or get_settings()
        self.bus = bus or run_event_bus
        self._registry: ToolRegistry | None = None
        self._approval_rules = ApprovalRules()

    # -- entry points ------------------------------------------------------

    async def execute(self, run: Run, workflow: Workflow) -> Run:
        self._prepare(run, workflow)
        await self._start(run)
        context: dict = {"trigger": run.input or {}}
        return await self._run_pass(run, workflow, context)

    async def resume(
        self, run: Run, workflow: Workflow, approval: Approval, *, approved: bool
    ) -> Run:
        self._prepare(run, workflow)
        await self._start(run)
        state = run.state_json or {}
        context: dict = state.get("context", {"trigger": run.input or {}})
        node_id = approval.requested_args.get("node_id")
        kind = approval.requested_args.get("kind")

        if kind == "approval_node":
            context[node_id] = {"approved": approved, "result": approved}
            await self._record_step(
                run,
                StepType.APPROVAL,
                node_id,
                {},
                {"approved": approved},
                status="approved" if approved else "rejected",
            )
        elif kind == "tool_node":
            if approved:
                context[node_id] = await self._invoke_tool(
                    run, approval.requested_tool, approval.requested_args.get("args", {})
                )
            else:
                context[node_id] = {"ok": False, "rejected": True}
                await self._record_step(
                    run,
                    StepType.TOOL_CALL,
                    node_id,
                    {"tool": approval.requested_tool},
                    {"rejected": True},
                    status="rejected",
                )
        return await self._run_pass(run, workflow, context)

    # -- core pass ---------------------------------------------------------

    async def _run_pass(self, run: Run, workflow: Workflow, context: dict) -> Run:
        definition = workflow.definition_json or {}
        nodes = definition.get("nodes", [])
        edges = definition.get("edges", [])
        node_by_id = {n["id"]: n for n in nodes}
        try:
            order = _topo_order(nodes, edges)
        except WorkflowGraphError as exc:
            return await self._fail(run, str(exc))

        last_result: dict | None = None
        for node_id in order:
            node = node_by_id[node_id]
            if node_id in context:
                last_result = context[node_id] if node_id != "trigger" else last_result
                continue
            if not self._reached(node_id, node_by_id, edges, context):
                continue

            outcome = await self._run_node(run, workflow, node, context)
            if outcome.pause is not None:
                await self._pause(run, node_id, outcome.pause, context)
                return self._reload(run)
            if outcome.failed:
                run.state_json = {"context": context}
                return await self._fail(run, outcome.error)
            context[node_id] = outcome.result or {}
            last_result = context[node_id]

        run.output = {"final": last_result}
        run.state_json = {"context": context}
        await self._set_status(run, RunStatus.SUCCEEDED)
        return run

    def _reached(self, node_id, node_by_id, edges, context) -> bool:
        node = node_by_id[node_id]
        if node.get("type") == "trigger":
            return True
        incoming = [e for e in edges if e.get("to") == node_id]
        if not incoming:
            return False
        return any(self._edge_taken(e, node_by_id, context) for e in incoming)

    def _edge_taken(self, edge, node_by_id, context) -> bool:
        src = edge.get("from")
        if src not in context:
            return False
        branch = self._branch_value(node_by_id.get(src, {}), context.get(src))
        when = edge.get("when")
        if branch is None:
            return when is None or str(when).lower() == "true"
        target = "true" if branch else "false"
        if when is None:
            return branch is True
        return str(when).lower() == target

    @staticmethod
    def _branch_value(node: dict, output) -> bool | None:
        node_type = node.get("type")
        if not isinstance(output, dict):
            return None
        if node_type == "condition":
            return bool(output.get("result"))
        if node_type == "approval":
            return bool(output.get("approved"))
        if node_type == "tool":
            if output.get("rejected") or output.get("denied"):
                return False
            return bool(output.get("ok"))
        return None

    # -- node execution ----------------------------------------------------

    async def _run_node(self, run, workflow, node, context) -> _NodeOutcome:
        retry = node.get("config", {}).get("retry", {})
        max_attempts = max(1, int(retry.get("max_attempts", 1)))
        backoff = float(retry.get("backoff_sec", 0))
        last_error = ""
        for attempt in range(max_attempts):
            try:
                return await self._dispatch_node(run, node, context)
            except NodeError as exc:
                last_error = str(exc)
                logger.warning("node %s attempt %d failed: %s", node.get("id"), attempt + 1, exc)
                if backoff and attempt < max_attempts - 1:
                    time.sleep(min(backoff, 0.01))  # tests keep this tiny
        await self._record_step(
            run,
            _NODE_STEP_TYPE.get(node.get("type"), StepType.TRANSFORM),
            node.get("id"),
            {"attempts": max_attempts},
            {"error": last_error},
            status="error",
            error=last_error,
        )
        return _NodeOutcome(failed=True, error=f"node '{node.get('id')}' failed: {last_error}")

    async def _dispatch_node(self, run, node, context) -> _NodeOutcome:
        node_type = node.get("type")
        cfg = node.get("config", {})
        node_id = node["id"]

        if node_type == "trigger":
            result = context.get("trigger", {})
            await self._record_step(run, StepType.TRIGGER, node_id, {}, {"payload": result})
            return _NodeOutcome(result=result if isinstance(result, dict) else {"value": result})

        if node_type == "condition":
            result = evaluate_condition(str(cfg.get("expr", "false")), context)
            await self._record_step(
                run, StepType.CONDITION, node_id, {"expr": cfg.get("expr")}, {"result": result}
            )
            return _NodeOutcome(result={"result": result})

        if node_type == "transform":
            rendered = render_value(cfg.get("template", cfg.get("value", {})), context)
            output = rendered if isinstance(rendered, dict) else {"value": rendered}
            await self._record_step(run, StepType.TRANSFORM, node_id, {}, output)
            return _NodeOutcome(result=output)

        if node_type == "delay":
            # Production schedules a delayed resume; inline execution continues now.
            await self._record_step(
                run,
                StepType.TRANSFORM,
                node_id,
                {"seconds": cfg.get("seconds", 0)},
                {"delayed": True},
            )
            return _NodeOutcome(result={"delayed": True})

        if node_type == "approval":
            return _NodeOutcome(
                pause=_Pause(kind="approval_node", summary=str(cfg.get("summary", "")))
            )

        if node_type == "agent":
            return await self._run_agent_node(run, node, context)

        if node_type == "tool":
            return await self._run_tool_node(run, node, context)

        raise NodeError(f"unsupported node type: {node_type}")

    async def _run_agent_node(self, run, node, context) -> _NodeOutcome:
        from legate.agent.runner import AgentRunner
        from legate.agent.service import _build_embedder, _build_vector_store, _resolve_llm
        from legate.db.models.agent import Agent
        from legate.repositories import runs as runs_repo

        cfg = node.get("config", {})
        agent = self.db.get(Agent, uuid.UUID(str(cfg["agent_id"]))) if cfg.get("agent_id") else None
        if agent is None or agent.workspace_id != run.workspace_id:
            raise NodeError("agent node references an unknown agent")
        rendered_input = render_value(cfg.get("input", ""), context)
        child = runs_repo.create(
            self.db,
            workspace_id=run.workspace_id,
            agent_id=agent.id,
            input={
                "input": (
                    rendered_input
                    if isinstance(rendered_input, str)
                    else json.dumps(rendered_input)
                )
            },
            created_by=run.created_by,
        )
        self.db.commit()
        runner = AgentRunner(
            db=self.db,
            llm=_resolve_llm(agent.provider, agent.model),
            registry=self._registry,
            embedder=_build_embedder(),
            vector_store=_build_vector_store(),
        )
        await runner.execute(child, agent)
        if child.status != RunStatus.SUCCEEDED:
            raise NodeError(f"agent run ended {child.status.value}")

        raw = (child.output or {}).get("output")
        output = raw if isinstance(raw, dict) else {"text": raw}
        result = {"output": output, "run_id": str(child.id)}
        await self._record_step(
            run, StepType.AGENT_REASON, node["id"], {"agent_id": str(agent.id)}, result
        )
        return _NodeOutcome(result=result)

    async def _run_tool_node(self, run, node, context) -> _NodeOutcome:
        cfg = node.get("config", {})
        tool_name = str(cfg.get("tool", ""))
        args = render_value(cfg.get("args", {}), context)
        registered = self._registry.get(tool_name) if self._registry else None
        if registered is None:
            raise NodeError(f"unknown tool: {tool_name}")

        decision = evaluate_tool_call(
            tool=registered.spec, allowed_tools=[tool_name], rules=self._approval_rules
        )
        if decision.effect is PolicyEffect.DENY:
            await self._record_step(
                run,
                StepType.TOOL_CALL,
                node["id"],
                {"tool": tool_name},
                {"denied": True, "reason": decision.reason},
                status="denied",
                error=decision.reason,
            )
            self._audit(run, "policy.deny", target=tool_name, metadata={"reason": decision.reason})
            return _NodeOutcome(result={"ok": False, "denied": True, "error": decision.reason})
        if decision.effect is PolicyEffect.REQUIRE_APPROVAL:
            return _NodeOutcome(
                pause=_Pause(
                    kind="tool_node",
                    tool=tool_name,
                    args=args if isinstance(args, dict) else {},
                    summary=f"Workflow wants to call {tool_name}",
                )
            )
        if decision.auto_approved:
            self._audit(run, "policy.auto_approve", target=tool_name)
        result = await self._invoke_tool(run, tool_name, args if isinstance(args, dict) else {})
        return _NodeOutcome(result=result)

    async def _invoke_tool(self, run, tool_name: str, args: dict) -> dict:
        registered = self._registry.get(tool_name) if self._registry else None
        if registered is None:
            await self._record_step(
                run,
                StepType.TOOL_CALL,
                tool_name,
                {"args": args},
                {"error": "unknown tool"},
                status="error",
                error="unknown tool",
            )
            return {"ok": False, "error": "unknown tool"}
        ctx = ToolContext(workspace_id=run.workspace_id, settings=self.settings, db=self.db)
        try:
            result = await registered.handler(args, ctx)
        except Exception as exc:  # a tool must never crash the run
            logger.exception("tool %s raised", tool_name)
            result = ToolResult(ok=False, error=f"tool raised: {exc}")
        await self._record_step(
            run,
            StepType.TOOL_CALL,
            tool_name,
            {"args": args},
            {"ok": result.ok, "data": result.data, "error": result.error},
            status="ok" if result.ok else "error",
            error=result.error,
        )
        return {"ok": result.ok, "data": result.data, "error": result.error}

    # -- pause / persistence ----------------------------------------------

    async def _pause(self, run, node_id, pause: _Pause, context: dict) -> None:
        run.state_json = {"context": context}
        self.db.add(run)
        ttl = self.settings.approval_ttl_hours
        expires_at = datetime.now(UTC) + timedelta(hours=ttl) if ttl > 0 else None
        approval = Approval(
            workspace_id=run.workspace_id,
            run_id=run.id,
            step_id=None,
            status=ApprovalStatus.PENDING,
            requested_tool=pause.tool or f"approval:{node_id}",
            requested_args={"node_id": node_id, "kind": pause.kind, "args": pause.args},
            summary=pause.summary or f"Workflow node {node_id} needs approval",
            pending_transcript=[],
            required_role=None,
            expires_at=expires_at,
        )
        self.db.add(approval)
        self.db.commit()
        self._audit(
            run,
            "policy.require_approval",
            target=pause.tool or node_id,
            metadata={"approval_id": str(approval.id), "node_id": node_id},
        )
        await self._set_status(run, RunStatus.WAITING_APPROVAL)

    def _prepare(self, run: Run, workflow: Workflow) -> None:
        from legate.connectors.resolution import build_workspace_tool_registry
        from legate.db.models.workspace import Workspace

        self._registry = build_workspace_tool_registry(self.db, run.workspace_id)
        workspace = self.db.get(Workspace, run.workspace_id)
        self._approval_rules = ApprovalRules.from_settings(
            workspace.settings_json if workspace else {}
        )

    async def _start(self, run: Run) -> None:
        if run.status != RunStatus.RUNNING:
            run.status = RunStatus.RUNNING
        if run.started_at is None:
            run.started_at = datetime.now(UTC)
        self.db.add(run)
        self.db.commit()
        await self.bus.publish(run.id, {"type": "run.status", "status": run.status.value})

    async def _record_step(
        self, run, step_type, name, input_data, output_data, *, status="ok", error=None
    ) -> Step:
        step = Step(
            run_id=run.id,
            workspace_id=run.workspace_id,
            idx=run.total_steps,
            type=step_type,
            name=str(name),
            status=status,
            input=input_data,
            output=output_data,
            error=error,
        )
        run.total_steps += 1
        self.db.add(step)
        self.db.add(run)
        self.db.commit()
        self.db.refresh(step)
        await self.bus.publish(run.id, {"type": "step", "step": _step_event(step)})
        return step

    def _audit(self, run, action: str, *, target=None, metadata=None) -> None:
        audit_record(
            self.db,
            workspace_id=run.workspace_id,
            actor=f"workflow:{run.workflow_id}",
            action=action,
            target=target,
            metadata=metadata or {},
        )
        self.db.commit()

    async def _set_status(self, run, status: RunStatus) -> None:
        run.status = status
        if status in {RunStatus.SUCCEEDED, RunStatus.FAILED, RunStatus.CANCELLED}:
            run.ended_at = datetime.now(UTC)
        self.db.add(run)
        self.db.commit()
        event = {"type": "run.status", "status": status.value}
        if run.error:
            event["error"] = run.error
        await self.bus.publish(run.id, event)

    async def _fail(self, run, reason: str) -> Run:
        run.error = reason
        await self._set_status(run, RunStatus.FAILED)
        return run

    def _reload(self, run: Run) -> Run:
        self.db.refresh(run)
        return run
