"""Node I/O models — the canonical Pydantic wrappers for all node data flow.

Every node's compute() takes NodeInputs and returns NodeOutputs.
These are Pydantic models, not dicts. The VM validates at boundaries.

NodeInputs uses extra="allow" because upstream nodes produce arbitrary
fields that flow through the graph. Specific nodes can subclass and
add typed fields for validation.

Usage:
    class MyNodeOutputs(NodeOutputs):
        response: str
        confidence: float

    class MyNode(NodeBase):
        async def compute(self, inputs: NodeInputs, ctx: Context) -> MyNodeOutputs:
            prompt = inputs.get("prompt", "")
            ...
            return MyNodeOutputs(response="...", confidence=0.95)
"""

from typing import Any, overload

from pydantic import BaseModel, Field

from core.common import JsonObject, JSONValue


class NodeInputs(BaseModel):
    """Base model for all node inputs.

    Accepts arbitrary fields from upstream nodes via extra="allow".
    Access fields with attribute access or .get() / .items().

    This is a Pydantic model, not a dict. Use .to_dict() when you need a
    plain dict (e.g. for Jinja2 template resolution). Prefer .get() for
    individual field access — it preserves original typed objects (Pydantic
    models, lists of models, etc.) without serialising them to dicts.
    """

    model_config = {"extra": "allow"}

    @overload
    def get(self, key: str) -> Any | None: ...

    @overload
    def get(self, key: str, default: Any) -> Any: ...

    def get(self, key: str, default: Any | None = None) -> Any:
        """Dict-like get that preserves the *original* typed value.

        Unlike model_dump(), this returns the Python object as-is — Pydantic
        models, typed lists, etc. stay intact.
        """
        # Check extra fields first (most dynamic data lives here)
        extra = self.__pydantic_extra__ or {}
        if key in extra:
            return extra[key]
        # Then declared model fields
        if key in type(self).model_fields and hasattr(self, key):
            val = getattr(self, key)
            if val is not None:
                return val
        return default

    def items(self) -> list[tuple[str, Any]]:
        """Iterate over all fields *without* serialisation.

        Returns the original Python objects, not model_dump() dicts.
        """
        pairs = [(name, getattr(self, name)) for name in type(self).model_fields]
        extra = self.__pydantic_extra__ or {}
        pairs.extend(extra.items())
        return pairs

    def keys(self) -> list[str]:
        """List all field names without serialisation."""
        names = list(type(self).model_fields.keys())
        extra = self.__pydantic_extra__ or {}
        names.extend(extra.keys())
        return names

    def to_dict(self) -> JsonObject:
        """Convert to a plain JsonObject for template resolution.

        This is the *only* place that calls model_dump().  Use .get() or
        attribute access everywhere else.
        """
        return self.model_dump()


class NodeOutputs(BaseModel):
    """Base model for all node outputs.

    Provides dict-like access so the VM scheduler can introspect
    results without converting to raw dicts.

    Subclass this for typed outputs per node type.

    Error fields are populated by the scheduler when a node fails:
    - ``error``: Human-readable error message
    - ``error_type``: Exception class name
    - ``error_traceback``: Truncated traceback string (debug builds)
    - ``error_node_id``: ID of the node that produced the error
    """

    model_config = {"extra": "allow"}

    # Structured error fields — populated by scheduler on failure
    error: str | None = None
    error_type: str | None = None
    error_traceback: str | None = None
    error_node_id: str | None = None

    @overload
    def get(self, key: str) -> Any | None: ...

    @overload
    def get(self, key: str, default: Any) -> Any: ...

    def get(self, key: str, default: Any | None = None) -> Any:
        """Dict-like get that preserves the *original* typed value."""
        extra = self.__pydantic_extra__ or {}
        if key in extra:
            return extra[key]
        if key in type(self).model_fields and hasattr(self, key):
            return getattr(self, key)
        return default

    def __contains__(self, key: str) -> bool:
        extra = self.__pydantic_extra__ or {}
        return key in extra or key in type(self).model_fields

    def __getitem__(self, key: str) -> Any:
        extra = self.__pydantic_extra__ or {}
        if key in extra:
            return extra[key]
        if key in type(self).model_fields:
            return getattr(self, key)
        raise KeyError(key)

    def items(self) -> list[tuple[str, Any]]:
        """Iterate over all fields *without* serialisation."""
        pairs = [(name, getattr(self, name)) for name in type(self).model_fields]
        extra = self.__pydantic_extra__ or {}
        pairs.extend(extra.items())
        return pairs

    def keys(self) -> list[str]:
        """List all field names without serialisation."""
        names = list(type(self).model_fields.keys())
        extra = self.__pydantic_extra__ or {}
        names.extend(extra.keys())
        return names

    def to_dict(self) -> JsonObject:
        """Convert to a plain JsonObject — only call at serialisation boundaries."""
        return self.model_dump()


# =============================================================================
# GENERIC I/O (for nodes that don't need specific typing)
# =============================================================================


class GenericInput(BaseModel):
    """Generic input when you don't need specific typing."""

    model_config = {"extra": "allow"}

    data: JSONValue = None


class GenericOutput(BaseModel):
    """Generic output when you don't need specific typing."""

    model_config = {"extra": "allow"}

    data: JSONValue = None


# =============================================================================
# VALIDATION I/O
# =============================================================================


class ValidationInput(BaseModel):
    """Input for validation nodes."""

    model_config = {"extra": "forbid"}

    data: JSONValue
    validation_schema: JsonObject | None = None
    rules: list[str] = Field(default_factory=list)


class ValidationOutput(BaseModel):
    """Output from validation nodes."""

    model_config = {"extra": "forbid"}

    valid: bool
    errors: list[str] = Field(default_factory=list)
    warnings: list[str] = Field(default_factory=list)
    cleaned_data: JSONValue = None


__all__ = [
    "GenericInput",
    "GenericOutput",
    "NodeInputs",
    "NodeOutputs",
    "ValidationInput",
    "ValidationOutput",
]
