"""PathwayBuilder - Fluent API for constructing Pathways.

The builder pattern enables clean, chainable pathway construction:

    pathway = (
        PathwayBuilder("my-pathway")
        .description("Analyze and summarize documents")
        .mode(PathwayMode.STATEFUL)
        .input({"document": "string", "max_length": "integer"})
        .tool("fetch", "web.fetch", args={"url": "{{input.document}}"})
        .llm("summarize", prompt="Summarize: {{fetch.output}}")
        .output({"summary": "{{summarize.output}}"})
        .connect("input", "fetch")
        .connect("fetch", "summarize")
        .connect("summarize", "output")
        .build()
    )

    # Export as JSON for version control
    json_str = pathway.model_dump_json(indent=2)

This fluent API is designed for:
- Programmatic pathway generation by agents
- Clean DSL for human developers
- Easy serialization to/from YAML
- Genetic algorithm mutation of pathway structure
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal, Self
from uuid import uuid4

if TYPE_CHECKING:
    from core.common import JsonObject
    from core.types.agents import ModelRouting

# Import node types for construction
from core.pathways.nodes import (
    CheckpointNode,
    InputNode,
    LLMNode,
    LoopNode,
    OutputNode,
    PathwayNode,
    StageNode,
    ToolNode,
    TransformNode,
)
from core.pathways.nodes.approval import ApprovalNode
from core.pathways.nodes.code_execute import CodeExecuteNode
from core.pathways.nodes.flow import ConditionalNode
from core.pathways.nodes.for_each import ForEachNode

# Import the domain models we build
from core.pathways.pathway import (
    Connection,
    Pathway,
    PathwayMode,
)


class PathwayBuilder:
    """Fluent builder for constructing Pathways programmatically."""

    def __init__(self, name: str, *, pathway_id: str | None = None) -> None:
        """Initialize builder with pathway name.

        Args:
            name: Human-readable name for the pathway
            pathway_id: Optional explicit ID (auto-generated if not provided)
        """
        self._pathway = Pathway(
            id=pathway_id or f"pathway_{uuid4().hex[:8]}",
            name=name,
        )
        self._last_node_id: str | None = None

    @property
    def pathway_id(self) -> str:
        """The ID of the pathway being built."""
        return self._pathway.id

    @property
    def node_count(self) -> int:
        """Number of nodes added so far."""
        return len(self._pathway.nodes)

    @property
    def connection_count(self) -> int:
        """Number of connections added so far."""
        return len(self._pathway.connections)

    def has_node(self, node_id: str) -> bool:
        """Check whether a node with the given ID already exists."""
        return node_id in self._pathway.nodes

    def has_connection(self, from_node: str, to_node: str) -> bool:
        """Check whether a connection between two nodes already exists."""
        return any(
            c.from_node == from_node and c.to_node == to_node for c in self._pathway.connections
        )

    def description(self, desc: str) -> Self:
        """Set pathway description."""
        self._pathway.description = desc
        return self

    def mode(self, mode: PathwayMode) -> Self:
        """Set execution mode (DAG or STATEFUL).

        - DAG: Parallel execution, no cycles (default)
        - STATEFUL: Sequential, cycles allowed, multi-turn capable
        """
        self._pathway.mode = mode
        return self

    def stateful(self) -> Self:
        """Shorthand to enable STATEFUL mode (cycles allowed)."""
        self._pathway.mode = PathwayMode.STATEFUL
        return self

    def metadata(self, **kwargs: object) -> Self:
        """Add metadata key-value pairs."""
        self._pathway.metadata.update(kwargs)
        return self

    def input(
        self,
        schema: JsonObject,
        *,
        defaults: JsonObject | None = None,
        node_id: str = "input",
    ) -> Self:
        """Add an input node with schema validation.

        Args:
            schema: JSON Schema-like dict defining input fields
            defaults: Default values for optional fields
            node_id: Node ID (default: "input")
        """
        self._pathway.add_node(
            node_id,
            InputNode(
                id=node_id,
                input_schema=schema,
                defaults=defaults or {},
            ),
        )
        self._last_node_id = node_id
        return self

    def output(
        self,
        mapping: JsonObject,
        *,
        schema: JsonObject | None = None,
        node_id: str = "output",
    ) -> Self:
        """Add an output node with mapping.

        Args:
            mapping: Dict mapping output field names to template expressions
            schema: Optional JSON Schema for output validation
            node_id: Node ID (default: "output")
        """
        self._pathway.add_node(
            node_id,
            OutputNode(
                id=node_id,
                mapping=mapping,
                output_schema=schema or {},
            ),
        )
        self._last_node_id = node_id
        return self

    def tool(
        self,
        node_id: str,
        tool_name: str,
        *,
        args: JsonObject | None = None,
    ) -> Self:
        """Add a tool node.

        Args:
            node_id: Unique ID for this node
            tool_name: Tool to call (e.g., "web.search", "workspace.read_file")
            args: Arguments with optional {{template}} placeholders
        """
        self._pathway.add_node(
            node_id,
            ToolNode(
                id=node_id,
                tool=tool_name,
                args=args or {},
            ),
        )
        self._last_node_id = node_id
        return self

    def llm(
        self,
        node_id: str,
        prompt: str,
        *,
        model: str = "reasoning",
        temperature: float = 0.7,
        max_tokens: int | None = None,
        system: str | None = None,
        response_format: Literal["text", "json"] = "text",
        json_schema: JsonObject | None = None,
        images: list[str] | None = None,
        tools: list[str] | str | None = None,
        max_tool_rounds: int = 0,
        model_routing: ModelRouting | None = None,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add an LLM node.

        Args:
            node_id: Unique ID for this node
            prompt: Prompt with {{template}} placeholders
            model: Capability name ("reasoning", "code", "fast") or explicit model
            temperature: Sampling temperature
            max_tokens: Maximum output tokens
            system: Optional system prompt
            response_format: "text" or "json"
            json_schema: JSON schema when response_format="json"
            images: Optional list of base64 images for vision models
            tools: Tool patterns (e.g., ["web.*"]) or template string
            max_tool_rounds: Maximum rounds of iterative tool calling (0 = single-shot)
            model_routing: Per-phase model routing (planning/execution/evaluation)
            name: Optional human-readable name
            description: Optional description
        """
        # Only include fields that the caller explicitly set.  Fields
        # left out of kwargs stay out of Pydantic's model_fields_set, so
        # LLMNode._effective_*() correctly falls back to RunState defaults
        # (agent model, temperature, system prompt, max_steps).
        kwargs: dict = {
            "id": node_id,
            "prompt": prompt,
            "response_format": response_format,
            "json_schema": json_schema,
            "images": images,
            "model_routing": model_routing,
            "name": name,
            "description": description,
        }
        # max_tokens has no RunState fallback — always pass it.
        kwargs["max_tokens"] = max_tokens

        # These have _effective_*() fallbacks via RunState; only pass
        # when the caller explicitly overrode the builder default.
        if model != "reasoning":
            kwargs["model"] = model
        if temperature != 0.7:
            kwargs["temperature"] = temperature
        if system is not None:
            kwargs["system"] = system
        if max_tool_rounds != 0:
            kwargs["max_tool_rounds"] = max_tool_rounds
        if tools is not None:
            kwargs["tools"] = tools
        self._pathway.add_node(node_id, LLMNode(**kwargs))
        self._last_node_id = node_id
        return self

    def transform(
        self,
        node_id: str,
        expr: str,
    ) -> Self:
        """Add a transform node for data manipulation.

        Args:
            node_id: Unique ID for this node
            expr: Safe expression to evaluate (result available in output.output)
        """
        self._pathway.add_node(
            node_id,
            TransformNode(
                id=node_id,
                expr=expr,
            ),
        )
        self._last_node_id = node_id
        return self

    def sub_pathway(
        self,
        node_id: str,
        pathway: str | Pathway,
        *,
        inputs: JsonObject | None = None,
        output_key: str | None = None,
    ) -> Self:
        """Add a sub-pathway node (nested pathway execution).

        Args:
            node_id: Unique ID for this node
            pathway: Pathway ID string or a built Pathway object
            inputs: Mapping of input values (supports {{template}} placeholders)
            output_key: Key to extract from sub-pathway outputs
        """
        pathway_id = pathway if isinstance(pathway, str) else pathway.id
        self._pathway.add_node(
            node_id,
            PathwayNode(
                id=node_id,
                pathway_id=pathway_id,
                inputs=inputs or {},
                output_key=output_key,
            ),
        )
        self._last_node_id = node_id
        return self

    def stage(
        self,
        node_id: str,
        goal: str,
        *,
        agent_id: str | None = None,
        transitions: list[tuple[str, str]] | None = None,
        on_enter: list[str] | None = None,
        on_exit: list[str] | None = None,
        max_iterations: int = 10,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a stage node for stateful pathway orchestration.

        Args:
            node_id: Unique ID for this node
            goal: What this stage should accomplish (supports {{template}})
            agent_id: Agent to run for this stage
            transitions: List of (condition, target_stage) pairs
            on_enter: Tool calls to run when entering the stage
            on_exit: Tool calls to run when leaving the stage
            max_iterations: Safety limit to prevent infinite loops
            name: Optional human-readable name
            description: Optional description
        """
        self._pathway.add_node(
            node_id,
            StageNode(
                id=node_id,
                goal=goal,
                agent_id=agent_id,
                transitions=transitions or [],
                on_enter=on_enter or [],
                on_exit=on_exit or [],
                max_iterations=max_iterations,
                name=name,
                description=description,
            ),
        )
        self._last_node_id = node_id
        return self

    def checkpoint(
        self,
        node_id: str,
        *,
        checkpoint_id: str | None = None,
        save_keys: list[str] | None = None,
        require_approval: bool = False,
        assignee_role: str | None = None,
        expiry_seconds: int | None = None,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a checkpoint node for persistence/resume.

        Args:
            node_id: Unique ID for this node
            checkpoint_id: Unique ID for the checkpoint (auto-generated if not set)
            save_keys: Which context keys to persist (None = all)
            require_approval: If true, pauses execution for human approval
            assignee_role: Role required to approve (e.g., "manager")
            expiry_seconds: How long checkpoint is valid (None = forever)
            name: Optional human-readable name
            description: Optional description
        """
        self._pathway.add_node(
            node_id,
            CheckpointNode(
                id=node_id,
                checkpoint_id=checkpoint_id,
                save_keys=save_keys,
                require_approval=require_approval,
                assignee_role=assignee_role,
                expiry_seconds=expiry_seconds,
                name=name,
                description=description,
            ),
        )
        self._last_node_id = node_id
        return self

    def conditional(
        self,
        node_id: str,
        condition: str,
        true_node: str,
        false_node: str,
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a conditional branching node.

        Args:
            node_id: Unique ID for this node
            condition: Expression to evaluate
            true_node: Node ID to go to if condition is true
            false_node: Node ID to go to if condition is false
            name: Optional human-readable name
            description: Optional description
        """
        self._pathway.add_node(
            node_id,
            ConditionalNode(
                id=node_id,
                condition=condition,
                true_node=true_node,
                false_node=false_node,
                name=name,
                description=description,
            ),
        )
        # Auto-wire the conditional to both branch targets so the DAG
        # scheduler sees the edges and can route execution correctly.
        self._pathway.connections.append(
            Connection(from_node=node_id, to_node=true_node),
        )
        self._pathway.connections.append(
            Connection(from_node=node_id, to_node=false_node),
        )
        self._last_node_id = node_id
        return self

    def approval(
        self,
        node_id: str,
        message: str = "Approval required to continue.",
        *,
        roles: list[str] | None = None,
        timeout_seconds: int = 3600,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a human approval gate.

        When reached, the DAG pauses and returns ``AWAITING_APPROVAL``.
        The caller resumes via ``PathwayVM.resume()`` after approval.

        Parameters
        ----------
        node_id
            Unique ID for this node.
        message
            What to show the approver (supports ``{{var}}`` templates).
        roles
            Roles that can approve.  Empty means anyone.
        timeout_seconds
            How long to wait for approval before timing out.
        name
            Optional human-readable name.
        description
            Optional description.
        """
        self._pathway.add_node(
            node_id,
            ApprovalNode(
                id=node_id,
                message=message,
                roles=roles or [],
                timeout_seconds=timeout_seconds,
                name=name,
                description=description,
            ),
        )
        self._last_node_id = node_id
        return self

    def for_each(
        self,
        node_id: str,
        items: str,
        body: Pathway,
        *,
        on_error: Literal["fail", "skip"] = "skip",
        max_parallel: int = 1,
        collect_results: bool = True,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a for-each iteration node.

        Iterates over a list and executes a sub-pathway for each item.

        Parameters
        ----------
        node_id
            Unique ID for this node.
        items
            Dotted expression resolving to a list from upstream inputs
            (e.g. ``"input.jobs"``).
        body
            The sub-pathway to execute for each item.
        on_error
            ``"skip"`` skips failed items, ``"fail"`` aborts the batch.
        max_parallel
            Maximum concurrent items.
        collect_results
            Whether to collect per-item results.
        name
            Optional human-readable name.
        description
            Optional description.
        """
        self._pathway.add_node(
            node_id,
            ForEachNode(
                id=node_id,
                items_expr=items,
                body_pathway_id=body.id,
                on_error=on_error,
                max_parallel=max_parallel,
                collect_results=collect_results,
                name=name,
                description=description,
            ),
        )
        self._last_node_id = node_id
        return self

    def loop_node(
        self,
        node_id: str,
        condition: str,
        body: list[str],
        *,
        max_iterations: int = 10,
        collect_results: bool = True,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a loop iteration node.

        Args:
            node_id: Unique ID for this node
            condition: Continue while this expression is true
            body: List of node IDs to execute in each iteration
            max_iterations: Safety limit
            collect_results: Whether to collect all iteration results
            name: Optional human-readable name
            description: Optional description
        """
        self._pathway.add_node(
            node_id,
            LoopNode(
                id=node_id,
                condition=condition,
                body=body,
                max_iterations=max_iterations,
                collect_results=collect_results,
                name=name,
                description=description,
            ),
        )
        self._last_node_id = node_id
        return self

    def code_execute(
        self,
        node_id: str,
        code: str,
        *,
        sandbox: Literal["local", "docker"] = "local",
        timeout_ms: int = 30000,
        memory_mb: int = 512,
        input_data: JsonObject | None = None,
        name: str | None = None,
        description: str | None = None,
    ) -> Self:
        """Add a sandboxed code execution node.

        Args:
            node_id: Unique ID for this node
            code: Python code to execute (supports {{template}})
            sandbox: Sandbox type ("local" or "docker")
            timeout_ms: Execution timeout in milliseconds
            memory_mb: Memory limit for Docker sandbox
            input_data: Data passed to the code as input_obj
            name: Optional human-readable name
            description: Optional description
        """
        self._pathway.add_node(
            node_id,
            CodeExecuteNode(
                id=node_id,
                code=code,
                sandbox=sandbox,
                timeout_ms=timeout_ms,
                memory_mb=memory_mb,
                input_data=input_data or {},
                name=name,
                description=description,
            ),
        )
        self._last_node_id = node_id
        return self

    def connect(
        self,
        from_node: str,
        to_node: str,
        *,
        from_output: str = "output",
        to_input: str = "input",
        when: str | None = None,
    ) -> Self:
        """Connect two nodes.

        Args:
            from_node: Source node ID
            to_node: Target node ID
            from_output: Output field to extract (default: "output")
            to_input: Input field to set (default: "input")
            when: Optional condition for stateful transitions
        """
        self._pathway.connections.append(
            Connection(
                from_node=from_node,
                to_node=to_node,
                from_output=from_output,
                to_input=to_input,
                when=when,
            ),
        )
        return self

    def chain(self, to_node: str) -> Self:
        """Connect the last added node to another node.

        Convenience method for linear pathways:
            builder.tool("a", "..").chain("b").tool("b", "..").chain("c")
        """
        if self._last_node_id is None:
            raise ValueError("No previous node to chain from")
        return self.connect(self._last_node_id, to_node)

    def build(self) -> Pathway:
        """Build and return the pathway.

        Validates node/scheduler compatibility before returning.

        Returns
        -------
            The constructed Pathway

        Raises
        ------
            ValueError
                If no nodes have been added or if nodes are incompatible
                with the selected pathway mode.
        """
        pathway = self._pathway

        if not pathway.nodes:
            raise ValueError("PathwayBuilder requires at least one node")

        # Validate node/scheduler compatibility.
        stateful_only: set[type] = {StageNode, CheckpointNode}

        stateful_ids = [n.id for n in pathway.nodes.values() if type(n) in stateful_only]
        if stateful_ids and pathway.mode != PathwayMode.STATEFUL:
            raise ValueError(
                f"Nodes {stateful_ids} require stateful mode. "
                "Use .mode(PathwayMode.STATEFUL) or .stateful()."
            )

        return pathway

    def to_json(self, *, indent: int = 2) -> str:
        """Build and serialize to JSON in one step."""
        return self._pathway.model_dump_json(indent=indent, exclude_none=True)


def pathway_builder(name: str, *, pathway_id: str | None = None) -> Self:
    """Create a new pathway builder.

    Args:
        name: Human-readable name for the pathway
        pathway_id: Optional explicit ID (auto-generated if not provided)

    Returns
    -------
        A fresh PathwayBuilder instance
    """
    return PathwayBuilder(name, pathway_id=pathway_id)


__all__ = [
    "PathwayBuilder",
    "pathway_builder",
]
