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

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

Schema generation is lazy — computed on first access, not at import time.
This prevents importing a single node class from triggering schema
resolution for every node in the system.
"""

import functools
from typing import Any, ClassVar

from pydantic import BaseModel, ConfigDict


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

    model_config = ConfigDict(frozen=True)

    name: str
    description: str
    config: dict[str, Any]


class _NodeTypeSpecData:
    """Raw spec data — no schema generation at definition time."""

    __slots__ = ("category", "description", "examples", "name", "node_class", "type")

    def __init__(
        self,
        *,
        type: str,
        name: str,
        description: str,
        category: str,
        node_class: type[BaseModel],
        examples: list[NodeTypeExample],
    ) -> None:
        self.type = type
        self.name = name
        self.description = description
        self.category = category
        self.node_class = node_class
        self.examples = examples


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

    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]


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

    Cached so each node class is only resolved once.
    """
    schema = node_cls.model_json_schema()
    props = dict(schema.get("properties") or {})

    for k in ("id", "type", "name", "description"):
        props.pop(k, None)

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

    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


def _build_spec(data: _NodeTypeSpecData) -> NodeTypeSpec:
    """Build a full NodeTypeSpec with lazily-generated config schema."""
    return NodeTypeSpec(
        type=data.type,
        name=data.name,
        description=data.description,
        category=data.category,
        node_class=data.node_class,
        config_schema=_config_schema_for(data.node_class),
        examples=data.examples,
    )


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

    Specs are defined as lightweight data. Schema generation happens
    on first access via list_all(), get(), or by_category().
    """

    _RAW: ClassVar[list[_NodeTypeSpecData]] = []
    _resolved: ClassVar[list[NodeTypeSpec] | None] = None

    @classmethod
    def _register_defaults(cls) -> None:
        """Lazy-import node classes and populate _RAW on first access."""
        if cls._RAW:
            return

        from core.pathways.nodes.approval import ApprovalNode
        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.for_each import ForEachNode
        from core.pathways.nodes.foundational import InputNode, OutputNode
        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

        cls._RAW = [
            _NodeTypeSpecData(
                type="input",
                name="Input",
                description=(
                    "Declare pathway inputs with schema. Makes inputs available as {{input.field}}."
                ),
                category="io",
                node_class=InputNode,
                examples=[
                    NodeTypeExample(
                        name="Query input",
                        description="Declare a required query input.",
                        config={"input_schema": {"query": {"type": "string"}}},
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="output",
                name="Output",
                description="Map pathway outputs from upstream nodes.",
                category="io",
                node_class=OutputNode,
                examples=[
                    NodeTypeExample(
                        name="Map result",
                        description="Map upstream output to pathway result.",
                        config={"mapping": {"result": "{{summarize.response}}"}},
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="tool",
                name="Tool",
                description="Call any tool. Use for Pattern 01 (Pure Pipeline).",
                category="compute",
                node_class=ToolNode,
                examples=[
                    NodeTypeExample(
                        name="Web search",
                        description="Call `web.search` with a templated query.",
                        config={
                            "tool": "web.search",
                            "args": {"query": "{{input.topic}}"},
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="transform",
                name="Transform",
                description="Evaluate a safe expression over inputs.",
                category="compute",
                node_class=TransformNode,
                examples=[
                    NodeTypeExample(
                        name="Math",
                        description="Compute a derived value from upstream inputs.",
                        config={"expr": "x * 2"},
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="llm",
                name="LLM",
                description=(
                    "Pattern 02: Pipeline + Generation. Direct LLM call — "
                    "generates text but doesn't choose tools."
                ),
                category="compute",
                node_class=LLMNode,
                examples=[
                    NodeTypeExample(
                        name="Summarize",
                        description="Summarize upstream search results.",
                        config={
                            "prompt": "Summarize: {{search.output.results}}",
                            "temperature": 0.3,
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="pathway",
                name="Pathway (Sub-pathway)",
                description=(
                    "Execute a registered sub-pathway. Enables pathway "
                    "composition — pathways in pathways."
                ),
                category="composition",
                node_class=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",
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="conditional",
                name="Conditional",
                description=("Branch based on a condition. Routes to true_node or false_node."),
                category="flow_control",
                node_class=ConditionalNode,
                examples=[
                    NodeTypeExample(
                        name="Check approval",
                        description="Route based on review approval.",
                        config={
                            "condition": "{{review.output.approved}} == true",
                            "true_node": "deploy",
                            "false_node": "revise",
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="checkpoint",
                name="Checkpoint",
                description=("Persist state for resume. Optionally pauses for human approval."),
                category="flow_control",
                node_class=CheckpointNode,
                examples=[
                    NodeTypeExample(
                        name="Approval gate",
                        description="Pause pathway for human approval.",
                        config={
                            "require_approval": True,
                            "assignee_role": "manager",
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="approval",
                name="Approval",
                description=(
                    "Human-in-the-loop approval gate. Pauses the DAG workflow "
                    "until a human approves, then resumes from checkpoint."
                ),
                category="flow_control",
                node_class=ApprovalNode,
                examples=[
                    NodeTypeExample(
                        name="Manager approval",
                        description="Require manager approval before escalation.",
                        config={
                            "message": "Approve escalation for {{input.description}}?",
                            "roles": ["manager"],
                            "timeout_seconds": 3600,
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="code_execute",
                name="Code Execute",
                description="Run sandboxed Python code with input data.",
                category="compute",
                node_class=CodeExecuteNode,
                examples=[
                    NodeTypeExample(
                        name="Transform data",
                        description="Process data with Python code.",
                        config={
                            "code": ("result = [x * 2 for x in input_obj.get('values', [])]"),
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="for_each",
                name="ForEach",
                description=(
                    "Iterate over a list and execute a sub-pathway for each item. "
                    "Supports parallel execution and per-item error handling (skip/fail)."
                ),
                category="flow_control",
                node_class=ForEachNode,
                examples=[
                    NodeTypeExample(
                        name="Process jobs",
                        description="Process each completed job with a billing sub-pathway.",
                        config={
                            "items_expr": "fetch.output.data.records",
                            "body_pathway_id": "create_invoice",
                            "on_error": "skip",
                            "max_parallel": 3,
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="stage",
                name="Stage",
                description=(
                    "Orchestration stage for stateful multi-turn pathways with transitions."
                ),
                category="flow_control",
                node_class=StageNode,
                examples=[
                    NodeTypeExample(
                        name="Research stage",
                        description="Research phase that transitions to analysis.",
                        config={
                            "goal": "Research the topic thoroughly",
                            "transitions": [["research_complete == true", "analysis"]],
                        },
                    ),
                ],
            ),
            _NodeTypeSpecData(
                type="loop",
                name="Loop",
                description="Iterate body nodes until a condition is met.",
                category="flow_control",
                node_class=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 _ensure_resolved(cls) -> list[NodeTypeSpec]:
        """Build full specs on first access."""
        if cls._resolved is None:
            cls._register_defaults()
            cls._resolved = [_build_spec(d) for d in cls._RAW]
        return cls._resolved

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

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

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

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


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