from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field

from nexus.resilience.errors import OrchestratorError


class WorkflowDefinition(BaseModel):
    """Serializable workflow definition from fluent API."""

    agents: dict[str, dict] = Field(default_factory=dict)
    tasks: list[dict] = Field(default_factory=list)
    forks: list[dict] = Field(default_factory=list)
    resilience_config: dict = Field(default_factory=dict)
    memory_config: dict = Field(default_factory=dict)
    mcp_servers: list[str] = Field(default_factory=list)


class WorkflowBuilder:
    """Fluent builder for NEXUS workflow definitions."""

    def __init__(self) -> None:
        self._agents: dict[str, dict] = {}
        self._tasks: list[dict] = []
        self._forks: list[dict] = []
        self._resilience_config: dict = {}
        self._memory_config: dict = {}
        self._mcp_servers: list[str] = []

    def add_agent(self, agent_id: str, role: str, goal: str, tools: list | None = None, **kwargs) -> WorkflowBuilder:
        """Add agent definition."""
        self._agents[agent_id] = {"agent_id": agent_id, "role": role, "goal": goal, "tools": tools or [], **kwargs}
        return self

    def add_task(
        self,
        task_id: str,
        agent: str,
        description: str,
        depends_on: list[str] | None = None,
        **kwargs,
    ) -> WorkflowBuilder:
        """Add task definition."""
        self._tasks.append(
            {
                "task_id": task_id,
                "agent": agent,
                "description": description,
                "depends_on": depends_on or [],
                **kwargs,
            }
        )
        return self

    def add_fork(
        self,
        fork_id: str,
        branches: list[str],
        merge_strategy: str = "first_wins",
        evaluator_agent: str | None = None,
    ) -> WorkflowBuilder:
        """Add fork definition."""
        self._forks.append(
            {
                "fork_id": fork_id,
                "branches": branches,
                "merge_strategy": merge_strategy,
                "evaluator_agent": evaluator_agent,
            }
        )
        return self

    def connect_mcp_server(self, url: str) -> WorkflowBuilder:
        """Register MCP server URL."""
        self._mcp_servers.append(url)
        return self

    def with_resilience(self, **kwargs) -> WorkflowBuilder:
        """Set resilience configuration entries."""
        self._resilience_config.update(kwargs)
        return self

    def with_memory(self, **kwargs) -> WorkflowBuilder:
        """Set memory configuration entries."""
        self._memory_config.update(kwargs)
        return self

    def build(self) -> WorkflowDefinition:
        """Validate references and return workflow definition."""
        agent_ids = set(self._agents.keys())
        task_ids = {t["task_id"] for t in self._tasks}

        for task in self._tasks:
            if task["agent"] not in agent_ids:
                raise OrchestratorError(f"Task agent '{task['agent']}' is not registered")

        for fork in self._forks:
            for branch in fork["branches"]:
                if branch not in task_ids:
                    raise OrchestratorError(f"Fork branch '{branch}' is not a known task")

        return WorkflowDefinition(
            agents=self._agents,
            tasks=self._tasks,
            forks=self._forks,
            resilience_config=self._resilience_config,
            memory_config=self._memory_config,
            mcp_servers=self._mcp_servers,
        )
