"""Node type registry for API discovery + validation.

This module backs:
- GET /api/v1/node-types
- pathways.validation.validate_node_config()

Supported node types:
- input: Declare pathway inputs
- output: Map pathway outputs
- tool: Call any tool
- llm: Direct LLM call
- transform: Expression evaluation
- pathway: Sub-pathway composition
- conditional: Flow control branching (if/else)
- checkpoint: Persist state for resume / approval gates
- code_execute: Sandboxed Python code execution
- stage: Orchestration stage for stateful pathways
- loop: Iterate body nodes until a condition is met
"""

from __future__ import annotations

from typing import Any, ClassVar

from pydantic import BaseModel, ConfigDict

from core.pathways.nodes.checkpoint import CheckpointNode
from core.pathways.nodes.code_execute import CodeExecuteNode
from core.pathways.nodes.flow import ConditionalNode
from core.pathways.nodes.foundational import InputNode, OutputNode
from core.pathways.nodes.handoff import HandoffNode
from core.pathways.nodes.llm import LLMNode
from core.pathways.nodes.loop import LoopNode
from core.pathways.nodes.stage import StageNode
from core.pathways.nodes.sub_pathway import PathwayNode
from core.pathways.nodes.tool import ToolNode
from core.pathways.nodes.transform import TransformNode


class NodeTypeExample(BaseModel):
    """Example usage of a node type."""

    model_config = ConfigDict(frozen=True)

    name: str
    description: str
    # Config is the `config` object used in API payloads (NOT the full node dict).
    config: dict[str, Any]


class NodeTypeSpec(BaseModel):
    """Specification for a node type."""

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    type: str
    name: str
    description: str
    category: str
    node_class: type[BaseModel]
    config_schema: dict[str, Any]
    examples: list[NodeTypeExample]


def _config_schema_for(node_cls: type[BaseModel]) -> dict[str, Any]:
    """Return JSON schema for the API `config` object for a node class.

    Node classes include identity fields like id/type/name/description that are not
    part of the API `config` object. We filter those out.
    """
    schema = node_cls.model_json_schema()
    props = dict(schema.get("properties") or {})

    # Fields that are *not* part of `config` in the HTTP API.
    for k in ("id", "type", "name", "description"):
        props.pop(k, None)

    required = [r for r in (schema.get("required") or []) if r in props]

    # Keep any $defs around so nested schemas remain valid.
    out: dict[str, Any] = {
        "type": "object",
        "additionalProperties": False,
        "properties": props,
    }
    if required:
        out["required"] = required
    if "$defs" in schema:
        out["$defs"] = schema["$defs"]
    return out


class NodeTypeRegistry:
    """Central registry of supported node types."""

    _SPECS: ClassVar[list[NodeTypeSpec]] = [
        NodeTypeSpec(
            type="input",
            name="Input",
            description="Declare pathway inputs with schema. Makes inputs available as {{input.field}}.",
            category="io",
            node_class=InputNode,
            config_schema=_config_schema_for(InputNode),
            examples=[
                NodeTypeExample(
                    name="Query input",
                    description="Declare a required query input.",
                    config={"input_schema": {"query": {"type": "string"}}},
                ),
            ],
        ),
        NodeTypeSpec(
            type="output",
            name="Output",
            description="Map pathway outputs from upstream nodes.",
            category="io",
            node_class=OutputNode,
            config_schema=_config_schema_for(OutputNode),
            examples=[
                NodeTypeExample(
                    name="Map result",
                    description="Map upstream output to pathway result.",
                    config={"mapping": {"result": "{{summarize.response}}"}},
                ),
            ],
        ),
        NodeTypeSpec(
            type="tool",
            name="Tool",
            description="Call any tool. Use for Pattern 01 (Pure Pipeline).",
            category="compute",
            node_class=ToolNode,
            config_schema=_config_schema_for(ToolNode),
            examples=[
                NodeTypeExample(
                    name="Web search",
                    description="Call `web.search` with a templated query.",
                    config={"tool": "web.search", "args": {"query": "{{input.topic}}"}},
                ),
            ],
        ),
        NodeTypeSpec(
            type="transform",
            name="Transform",
            description="Evaluate a safe expression over inputs.",
            category="compute",
            node_class=TransformNode,
            config_schema=_config_schema_for(TransformNode),
            examples=[
                NodeTypeExample(
                    name="Math",
                    description="Compute a derived value from upstream inputs.",
                    config={"expr": "x * 2"},
                )
            ],
        ),
        NodeTypeSpec(
            type="llm",
            name="LLM",
            description="Pattern 02: Pipeline + Generation. Direct LLM call — generates text but doesn't choose tools.",
            category="compute",
            node_class=LLMNode,
            config_schema=_config_schema_for(LLMNode),
            examples=[
                NodeTypeExample(
                    name="Summarize",
                    description="Summarize upstream search results.",
                    config={"prompt": "Summarize: {{search.output.results}}", "temperature": 0.3},
                ),
            ],
        ),
        NodeTypeSpec(
            type="pathway",
            name="Pathway (Sub-pathway)",
            description="Execute a registered sub-pathway. Enables pathway composition — pathways in pathways.",
            category="composition",
            node_class=PathwayNode,
            config_schema=_config_schema_for(PathwayNode),
            examples=[
                NodeTypeExample(
                    name="Diagnostic sub-pathway",
                    description="Run a reusable diagnostic pathway.",
                    config={
                        "pathway_id": "diagnostic_deep_dive_v2",
                        "inputs": {
                            "issue": "{{state.issue_details}}",
                            "location": "{{state.area}}",
                        },
                    },
                ),
                NodeTypeExample(
                    name="Approval pathway",
                    description="Embed a standard approval process.",
                    config={
                        "pathway_id": "standard_approval_flow",
                        "inputs": {"item": "{{state.item}}", "cost": "{{state.total_cost}}"},
                        "output_key": "approval_result",
                    },
                ),
            ],
        ),
        NodeTypeSpec(
            type="handoff",
            name="Handoff",
            description="Route + delegate to the best-matching agent. Core multi-agent orchestration primitive.",
            category="orchestration",
            node_class=HandoffNode,
            config_schema=_config_schema_for(HandoffNode),
            examples=[
                NodeTypeExample(
                    name="Route by approach",
                    description="Route to specialist agent based on classification.",
                    config={
                        "candidates": [
                            {
                                "agent_id": "system.researcher",
                                "description": "Information gathering",
                                "when": "{{understand.output.approach}} == 'research'",
                            },
                            {
                                "agent_id": "system.coder",
                                "description": "Code tasks",
                                "when": "{{understand.output.approach}} == 'code'",
                            },
                        ],
                        "fallback": "direct_answer",
                        "message": "{{input.message}}",
                    },
                ),
                NodeTypeExample(
                    name="Simple delegation",
                    description="Delegate to a single agent unconditionally.",
                    config={
                        "candidates": [
                            {
                                "agent_id": "system.planner",
                                "description": "Task decomposition",
                            },
                        ],
                        "message": "Break down this task: {{input.message}}",
                    },
                ),
            ],
        ),
        NodeTypeSpec(
            type="conditional",
            name="Conditional",
            description="Branch based on a condition. Routes to true_node or false_node.",
            category="flow_control",
            node_class=ConditionalNode,
            config_schema=_config_schema_for(ConditionalNode),
            examples=[
                NodeTypeExample(
                    name="Check approval",
                    description="Route based on review approval.",
                    config={
                        "condition": "{{review.output.approved}} == true",
                        "true_node": "deploy",
                        "false_node": "revise",
                    },
                ),
            ],
        ),
        NodeTypeSpec(
            type="checkpoint",
            name="Checkpoint",
            description="Persist state for resume. Optionally pauses for human approval.",
            category="flow_control",
            node_class=CheckpointNode,
            config_schema=_config_schema_for(CheckpointNode),
            examples=[
                NodeTypeExample(
                    name="Approval gate",
                    description="Pause pathway for human approval.",
                    config={
                        "require_approval": True,
                        "assignee_role": "manager",
                    },
                ),
            ],
        ),
        NodeTypeSpec(
            type="code_execute",
            name="Code Execute",
            description="Run sandboxed Python code with input data.",
            category="compute",
            node_class=CodeExecuteNode,
            config_schema=_config_schema_for(CodeExecuteNode),
            examples=[
                NodeTypeExample(
                    name="Transform data",
                    description="Process data with Python code.",
                    config={
                        "code": "result = [x * 2 for x in input_obj.get('values', [])]",
                    },
                ),
            ],
        ),
        NodeTypeSpec(
            type="stage",
            name="Stage",
            description="Orchestration stage for stateful multi-turn pathways with transitions.",
            category="flow_control",
            node_class=StageNode,
            config_schema=_config_schema_for(StageNode),
            examples=[
                NodeTypeExample(
                    name="Research stage",
                    description="Research phase that transitions to analysis.",
                    config={
                        "goal": "Research the topic thoroughly",
                        "transitions": [["research_complete == true", "analysis"]],
                    },
                ),
            ],
        ),
        NodeTypeSpec(
            type="loop",
            name="Loop",
            description="Iterate body nodes until a condition is met.",
            category="flow_control",
            node_class=LoopNode,
            config_schema=_config_schema_for(LoopNode),
            examples=[
                NodeTypeExample(
                    name="Retry loop",
                    description="Retry a step up to 3 times.",
                    config={
                        "condition": "success != true",
                        "body": ["attempt"],
                        "max_iterations": 3,
                    },
                ),
            ],
        ),
    ]

    @classmethod
    def list_all(cls) -> list[NodeTypeSpec]:
        return list(cls._SPECS)

    @classmethod
    def categories(cls) -> list[str]:
        return sorted({s.category for s in cls._SPECS})

    @classmethod
    def by_category(cls, category: str) -> list[NodeTypeSpec]:
        c = (category or "").strip().lower()
        return [s for s in cls._SPECS if s.category.lower() == c]

    @classmethod
    def get(cls, node_type: str) -> NodeTypeSpec | None:
        t = (node_type or "").strip().lower()
        for s in cls._SPECS:
            if s.type.lower() == t:
                return s
        return None


__all__ = ["NodeTypeExample", "NodeTypeRegistry", "NodeTypeSpec"]
