# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason-manifest

# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_maco

from enum import Enum
from typing import Annotated, Any, Dict, List, Literal, Optional, Sequence, Union

from pydantic import ConfigDict, Field, StringConstraints, model_validator

from coreason_manifest.definitions.base import CoReasonBaseModel


class StateDefinition(CoReasonBaseModel):
    """Defines the internal state (memory) of the Recipe.

    Attributes:
        schema: JSON Schema of the keys available in the shared memory.
        persistence: Configuration for state durability.
    """

    model_config = ConfigDict(extra="forbid")

    schema_: Dict[str, Any] = Field(
        ..., alias="schema", description="JSON Schema of the keys available in the shared memory."
    )
    persistence: Literal["ephemeral", "persistent"] = Field(
        default="ephemeral", description="Configuration for state durability."
    )


class CouncilConfig(CoReasonBaseModel):
    """Configuration for 'Architectural Triangulation'.

    Attributes:
        strategy: The strategy for the council (e.g., 'consensus').
        voters: List of agents or models that vote.
    """

    model_config = ConfigDict(extra="forbid")

    strategy: str = Field(default="consensus", description="The strategy for the council, e.g., 'consensus'.")
    voters: List[str] = Field(..., description="List of agents or models that vote.")


class VisualMetadata(CoReasonBaseModel):
    """Data explicitly for the UI.

    Attributes:
        label: The label to display for the node.
        x_y_coordinates: The X and Y coordinates for the node on the canvas.
        icon: The icon to represent the node.
        animation_style: The animation style for the node.
    """

    model_config = ConfigDict(extra="forbid")

    label: Optional[str] = Field(None, description="The label to display for the node.")
    x_y_coordinates: Optional[List[float]] = Field(
        None, description="The X and Y coordinates for the node on the canvas."
    )
    icon: Optional[str] = Field(None, description="The icon to represent the node.")
    animation_style: Optional[str] = Field(None, description="The animation style for the node.")


class RuntimeVisualMetadata(VisualMetadata):
    """Visual metadata with runtime extensions.

    Attributes:
        color: Color override for the node (e.g., '#00FF00').
        progress: Progress indicator (e.g., '0.5' or '50%').
        animation: Animation override (e.g., 'pulse').
    """

    color: Optional[str] = Field(None, description="Color override for the node.")
    progress: Optional[str] = Field(None, description="Progress indicator.")
    animation: Optional[str] = Field(None, description="Animation override.")


class BaseNode(CoReasonBaseModel):
    """Base model for all node types.

    Attributes:
        id: Unique identifier for the node.
        council_config: Optional configuration for architectural triangulation.
        visual: Visual metadata for the UI.
        metadata: Generic metadata for operational context (e.g. cost tracking, SLAs).
    """

    model_config = ConfigDict(extra="forbid")

    id: str = Field(..., description="Unique identifier for the node.")
    council_config: Optional[CouncilConfig] = Field(
        None, description="Optional configuration for architectural triangulation."
    )
    visual: Optional[VisualMetadata] = Field(None, description="Visual metadata for the UI.")
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Generic metadata for operational context (e.g. cost tracking, SLAs).",
    )


class AgentNode(BaseNode):
    """A node that calls a specific atomic agent.

    Attributes:
        type: The type of the node (must be 'agent').
        agent_name: The name of the atomic agent to call.
        system_prompt: Overrides the registry default prompt. Required for ad-hoc/optimized agents.
        config: Runtime-specific configuration (e.g., model parameters, temperature). Merged with registry defaults.
        overrides: Runtime overrides for the agent (e.g., temperature, prompt_template_vars).
    """

    type: Literal["agent"] = Field("agent", description="Discriminator for AgentNode.")
    agent_name: str = Field(..., description="The name of the atomic agent to call.")
    system_prompt: Optional[str] = Field(
        default=None,
        description="Overrides the registry default prompt. Required for ad-hoc/optimized agents.",
    )
    config: Optional[Dict[str, Any]] = Field(
        default=None,
        description=(
            "Runtime-specific configuration (e.g., model parameters, temperature). Merged with registry defaults."
        ),
    )
    overrides: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Runtime overrides for the agent (e.g., temperature, prompt_template_vars).",
    )


class HumanNode(BaseNode):
    """A node that pauses execution for user input/approval.

    Attributes:
        type: The type of the node (must be 'human').
        timeout_seconds: Optional timeout in seconds for the user interaction.
    """

    type: Literal["human"] = Field("human", description="Discriminator for HumanNode.")
    timeout_seconds: Optional[int] = Field(None, description="Optional timeout in seconds for the user interaction.")


class LogicNode(BaseNode):
    """A node that executes pure Python logic.

    Attributes:
        type: The type of the node (must be 'logic').
        code: The Python logic code to execute.
    """

    type: Literal["logic"] = Field("logic", description="Discriminator for LogicNode.")
    code: str = Field(..., description="The Python logic code to execute.")


class DataMappingStrategy(str, Enum):
    """Strategy for mapping data."""

    DIRECT = "direct"
    JSONPATH = "jsonpath"
    LITERAL = "literal"


class DataMapping(CoReasonBaseModel):
    """Defines how to transform data between parent and child."""

    model_config = ConfigDict(extra="forbid")

    source: str = Field(..., description="The path/key source.")
    strategy: DataMappingStrategy = Field(default=DataMappingStrategy.DIRECT, description="The mapping strategy.")


class RecipeNode(BaseNode):
    """A node that executes another Recipe as a sub-graph.

    Attributes:
        type: The type of the node (must be 'recipe').
        recipe_id: The ID of the recipe to execute.
        input_mapping: How parent state maps to child inputs (parent_key -> child_key).
        output_mapping: How child result maps back to parent state (child_key -> parent_key).
    """

    type: Literal["recipe"] = Field("recipe", description="Discriminator for RecipeNode.")
    recipe_id: str = Field(..., description="The ID of the recipe to execute.")
    input_mapping: Dict[str, Union[str, DataMapping]] = Field(
        ..., description="Mapping of parent state keys to child input keys."
    )
    output_mapping: Dict[str, Union[str, DataMapping]] = Field(
        ..., description="Mapping of child output keys to parent state keys."
    )


class MapNode(BaseNode):
    """A node that spawns multiple parallel executions of a sub-branch.

    Attributes:
        type: The type of the node (must be 'map').
        items_path: Dot-notation path to the list in the state.
        processor_node_id: The node (or subgraph) to run for each item.
        concurrency_limit: Max parallel executions.
    """

    type: Literal["map"] = Field("map", description="Discriminator for MapNode.")
    items_path: str = Field(..., description="Dot-notation path to the list in the state.")
    processor_node_id: str = Field(..., description="The node (or subgraph) to run for each item.")
    concurrency_limit: int = Field(..., description="Max parallel executions.")


# Discriminated Union for polymorphism
Node = Annotated[
    Union[AgentNode, HumanNode, LogicNode, RecipeNode, MapNode],
    Field(discriminator="type", description="Polymorphic node definition."),
]


class Edge(CoReasonBaseModel):
    """Represents a connection between two nodes.

    Attributes:
        source_node_id: The ID of the source node.
        target_node_id: The ID of the target node.
        condition: Optional Python expression for conditional branching.
    """

    model_config = ConfigDict(extra="forbid")

    source_node_id: str = Field(..., description="The ID of the source node.")
    target_node_id: str = Field(..., description="The ID of the target node.")
    condition: Optional[str] = Field(None, description="Optional Python expression for conditional branching.")


RouterRef = Annotated[
    str,
    StringConstraints(
        pattern=r"^[a-zA-Z_][a-zA-Z0-9_]*(\.[a-zA-Z_][a-zA-Z0-9_]*)*$",
        strip_whitespace=True,
    ),
]


class RouterExpression(CoReasonBaseModel):
    """A structured expression for routing logic (e.g., CEL or JSONLogic)."""

    model_config = ConfigDict(extra="forbid")

    operator: str = Field(..., description="The operator (e.g., 'eq', 'gt').")
    args: List[Any] = Field(..., description="Arguments for the expression.")


RouterDefinition = Annotated[
    Union[RouterRef, RouterExpression],
    Field(description="A reference to a python function or a logic expression."),
]


class ConditionalEdge(CoReasonBaseModel):
    """Represents a dynamic routing connection from one node to multiple potential targets.

    Attributes:
        source_node_id: The ID of the source node.
        router_logic: A reference to a python function or a logic expression that returns the next node ID.
        mapping: A dictionary mapping the router's output (e.g., "approve", "reject") to target Node IDs.
    """

    model_config = ConfigDict(extra="forbid")

    source_node_id: str = Field(..., description="The ID of the source node.")
    router_logic: RouterDefinition = Field(
        ..., description="A reference to a python function or logic expression that determines the path."
    )
    mapping: Dict[str, str] = Field(..., description="Map of router output values to target node IDs.")


def validate_edge_integrity(
    nodes: Sequence[Node],
    edges: Sequence[Union[Edge, ConditionalEdge]],
) -> None:
    """Ensures that all edges point to valid nodes.

    Args:
        nodes: List of nodes in the graph.
        edges: List of edges connecting the nodes.

    Raises:
        ValueError: If an edge points to a non-existent node.
    """
    node_ids = {node.id for node in nodes}

    for edge in edges:
        if edge.source_node_id not in node_ids:
            raise ValueError(f"Edge source node '{edge.source_node_id}' not found in nodes.")

        # ConditionalEdge mapping targets
        if isinstance(edge, ConditionalEdge):
            for target_id in edge.mapping.values():
                if target_id not in node_ids:
                    raise ValueError(f"ConditionalEdge target node '{target_id}' not found in nodes.")
        else:
            # Regular Edge
            if edge.target_node_id not in node_ids:
                raise ValueError(f"Edge target node '{edge.target_node_id}' not found in nodes.")


class GraphTopology(CoReasonBaseModel):
    """The topology definition of the recipe.

    Attributes:
        nodes: List of nodes in the graph.
        edges: List of edges connecting the nodes.
        state_schema: Optional schema definition for the graph state.
    """

    model_config = ConfigDict(extra="forbid")

    nodes: List[Node] = Field(..., description="List of nodes in the graph.")
    edges: List[Union[Edge, ConditionalEdge]] = Field(..., description="List of edges connecting the nodes.")
    state_schema: Optional[StateDefinition] = Field(default=None, description="Schema definition for the graph state.")

    @model_validator(mode="after")
    def validate_graph_integrity(self) -> "GraphTopology":
        """Ensures that all edges point to valid nodes."""
        validate_edge_integrity(self.nodes, self.edges)
        return self


Topology = GraphTopology
