from typing import Any

from agentflow.core.state import Message
from agentflow.utils import ResponseGranularity
from pydantic import BaseModel, Field, field_validator


class GraphInputSchema(BaseModel):
    """
    Schema for graph input including messages and configuration.
    """

    messages: list[Message] = Field(
        ..., description="List of messages to process through the graph"
    )
    initial_state: dict[str, Any] | None = Field(
        default=None,
        description="Initial state for the graph execution",
    )
    config: dict[str, Any] | None = Field(
        default=None,
        description="Optional configuration for graph execution",
    )
    recursion_limit: int = Field(
        default=25,
        ge=1,
        le=100,
        description="Maximum recursion limit for graph execution",
    )

    @field_validator("messages")
    @classmethod
    def messages_must_not_be_empty(cls, v: list[Message]) -> list[Message]:
        if not v:
            raise ValueError("messages must contain at least one message")
        return v

    response_granularity: ResponseGranularity = Field(
        default=ResponseGranularity.LOW,
        description="Granularity of the response (full, partial, low)",
    )


class GraphInvokeOutputSchema(BaseModel):
    """
    Schema for graph invoke output.
    """

    messages: list[Message] = Field(
        ...,
        description="Final processed messages from the graph",
    )
    state: dict[str, Any] | None = Field(
        default=None,
        description="State information from the graph execution",
    )
    context: list[Message] | None = Field(
        default=None,
        description="Context information from the graph execution",
    )
    summary: str | None = Field(
        default=None,
        description="Summary information from the graph execution",
    )
    meta: dict[str, Any] | None = Field(
        default=None,
        description="Meta information from the graph execution",
    )


# class GraphStreamChunkSchema(BaseModel):
#     """
#     Schema for individual stream chunks from graph execution.
#     """

#     data: dict[str, Any] = Field(..., description="Chunk data")
#     metadata: dict[str, Any] | None = Field(default=None, description="Chunk metadata")


class NodeSchema(BaseModel):
    """Schema for individual graph nodes."""

    id: str = Field(..., description="Unique identifier for the node")
    name: str = Field(..., description="Name of the node")


class EdgeSchema(BaseModel):
    """Schema for individual graph edges."""

    id: str = Field(..., description="Unique identifier for the edge")
    source: str = Field(..., description="Source node identifier")
    target: str = Field(..., description="Target node identifier")


class GraphInfoSchema(BaseModel):
    """Schema for graph metadata and configuration."""

    node_count: int = Field(..., description="Number of nodes in the graph")
    edge_count: int = Field(..., description="Number of edges in the graph")
    checkpointer: bool = Field(..., description="Whether checkpointer is enabled")
    checkpointer_type: str | None = Field(None, description="Type of checkpointer if enabled")
    publisher: bool = Field(..., description="Whether publisher is enabled")
    store: bool = Field(..., description="Whether store is enabled")
    interrupt_before: list[str] | None = Field(None, description="Nodes to interrupt before")
    interrupt_after: list[str] | None = Field(None, description="Nodes to interrupt after")
    context_type: str | None = Field(None, description="Type of context for the graph")
    id_generator: str | None = Field(None, description="ID generator type for the graph")
    id_type: str | None = Field(None, description="ID type for the graph")
    state_type: str | None = Field(None, description="State type for the graph")
    state_fields: list[str] | None = Field(None, description="State fields for the graph")


class GraphSchema(BaseModel):
    """Schema for the complete graph structure."""

    info: GraphInfoSchema = Field(..., description="Graph metadata and configuration")
    nodes: list[NodeSchema] = Field(..., description="List of nodes in the graph")
    edges: list[EdgeSchema] = Field(..., description="List of edges in the graph")


class GraphStopSchema(BaseModel):
    """Schema for stopping graph execution."""

    thread_id: str = Field(..., min_length=1, description="Thread ID to stop execution for")

    @field_validator("thread_id")
    @classmethod
    def thread_id_not_blank(cls, v: str) -> str:
        if not v.strip():
            raise ValueError("thread_id cannot be empty or whitespace")
        return v.strip()

    config: dict[str, Any] | None = Field(
        default=None, description="Optional configuration for the stop operation"
    )


class RemoteToolSchema(BaseModel):
    """Schema for remote tool execution."""

    node_name: str = Field(..., description="Name of the node representing the tool")
    name: str = Field(..., description="Name of the tool to execute")
    description: str = Field(..., description="Description of the tool")
    parameters: dict[str, Any] = Field(..., description="Parameters for the tool")


class GraphSetupSchema(BaseModel):
    """Schema for setting up graph execution."""

    tools: list[RemoteToolSchema] = Field(
        ..., description="List of remote tools available for the graph"
    )


class FixGraphRequestSchema(BaseModel):
    """Schema for fixing graph state by removing messages with empty tool call content."""

    thread_id: str = Field(..., min_length=1, description="Thread ID to fix the graph state for")

    @field_validator("thread_id")
    @classmethod
    def thread_id_not_blank(cls, v: str) -> str:
        if not v.strip():
            raise ValueError("thread_id cannot be empty or whitespace")
        return v.strip()

    config: dict[str, Any] | None = Field(
        default=None, description="Optional configuration for the fix operation"
    )


class FixGraphResponseSchema(BaseModel):
    """Schema for the fix graph operation response."""

    success: bool = Field(..., description="Whether the fix operation was successful")
    message: str = Field(..., description="Status message from the fix operation")
    removed_count: int = Field(
        default=0, description="Number of messages with empty tool calls that were removed"
    )
    state: dict[str, Any] | None = Field(
        default=None, description="Updated state after fixing the graph"
    )
