"""WebUI-specific models for Agentbyte."""

from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.context import ToolApprovalResponse


class EntityInfo(BaseModel):
    """Common discovery metadata for a WebUI entity."""

    id: str = Field(description="Stable entity identifier")
    name: str = Field(description="Human-readable name")
    description: str | None = Field(default=None, description="Entity description")
    type: Literal["agent", "orchestrator", "workflow"] = Field(
        description="Entity surface type"
    )
    source: Literal["memory", "directory", "github"] = Field(
        description="How the entity was discovered"
    )
    module_path: str | None = Field(
        default=None,
        description="Source module path when directory-discovered",
    )
    has_env: bool = Field(
        default=False,
        description="Whether a colocated .env file exists for the entity",
    )
    example_tasks: list[str] = Field(
        default_factory=list,
        description="Suggested example tasks for the UI",
    )


class AgentInfo(EntityInfo):
    """Discovery metadata for agents."""

    type: Literal["agent"] = "agent"
    model: str | None = Field(default=None, description="Configured model hint")
    tools: list[str] = Field(default_factory=list, description="Available tool names")
    memory_type: str | None = Field(
        default=None,
        description="Configured memory backend type when available",
    )
    middleware_count: int = Field(
        default=0,
        description="Number of configured middlewares",
    )


class OrchestratorInfo(EntityInfo):
    """Discovery metadata for orchestrators."""

    type: Literal["orchestrator"] = "orchestrator"
    orchestrator_type: str = Field(description="Concrete orchestrator type")
    agents: list[str] = Field(
        default_factory=list,
        description="Participating agent names",
    )
    termination_conditions: list[str] = Field(
        default_factory=list,
        description="Termination condition summaries",
    )


class WorkflowInfo(EntityInfo):
    """Discovery metadata for workflows."""

    type: Literal["workflow"] = "workflow"
    steps: list[str] = Field(default_factory=list, description="Workflow step ids")
    start_step: str | None = Field(
        default=None,
        description="Workflow start step id",
    )
    end_step_ids: list[str] = Field(
        default_factory=list,
        description="Workflow end step ids",
    )
    input_schema: dict[str, Any] | None = Field(
        default=None,
        description="JSON schema for the start step input when available",
    )
    start_step_id: str | None = Field(
        default=None,
        description="Backward-compatible alias for the workflow start step id",
    )


class RunEntityRequest(BaseModel):
    """Request payload for WebUI execution endpoints."""

    messages: list[dict[str, Any]] | None = Field(
        default=None,
        description="Serialized Agentbyte messages for agent/orchestrator runs",
    )
    input_data: Any = Field(
        default=None,
        description="Structured input data for workflow execution",
    )
    session_id: str | None = Field(
        default=None,
        description="Optional session id to continue",
    )
    user_id: str | None = Field(
        default=None,
        description="Optional user id that owns the session",
    )
    stream_tokens: bool = Field(
        default=True,
        description="Enable token-level streaming when the agent path supports it",
    )
    approval_responses: list[ToolApprovalResponse] | None = Field(
        default=None,
        description="Approval responses to inject before resuming an agent run",
    )


class WebUIStreamEvent(BaseModel):
    """Envelope used for SSE events sent to the frontend."""

    session_id: str = Field(description="Owning session id")
    timestamp: datetime = Field(
        default_factory=datetime.now,
        description="Envelope creation time",
    )
    event: Any = Field(description="Serialized framework event payload")

    model_config = ConfigDict(arbitrary_types_allowed=True)


class CreateSessionRequest(BaseModel):
    """Request payload for explicit session creation."""

    entity_id: str = Field(description="Entity identifier to bind the session to")
    entity_type: str = Field(default="agent", description="Entity type for the session")
    user_id: str | None = Field(default=None, description="Optional session owner id")


class HealthResponse(BaseModel):
    """Basic health response for the WebUI backend."""

    status: str = Field(description="Health status")
    entities_dir: str | None = Field(
        default=None,
        description="Entity discovery directory when configured",
    )
    entities_count: int = Field(default=0, description="Number of visible entities")


class SessionInfo(BaseModel):
    """Session metadata returned by the WebUI API."""

    id: str = Field(description="Session identifier")
    entity_id: str = Field(description="Linked entity id")
    entity_type: str = Field(description="Linked entity type")
    created_at: datetime = Field(description="When the session was created")
    last_activity: datetime = Field(description="When the session was last touched")
    message_count: int = Field(default=0, description="Stored message count")
    user_id: str | None = Field(default=None, description="Owning user identifier")


class SessionMessagesResponse(BaseModel):
    """Session transcript payload returned by the WebUI API."""

    session_id: str = Field(description="Session identifier")
    messages: list[dict[str, Any]] = Field(
        default_factory=list,
        description="Serialized session messages",
    )
    final_result: str | None = Field(
        default=None,
        description="Latest persisted orchestrator final result when available",
    )
    stop_reason: str | None = Field(
        default=None,
        description="Latest persisted orchestrator stop reason when available",
    )


class StatsResponse(BaseModel):
    """Aggregated entity and session statistics."""

    entities: dict[str, Any] = Field(description="Entity counters")
    sessions: dict[str, Any] = Field(description="Session counters")


Entity = AgentInfo | OrchestratorInfo | WorkflowInfo


__all__ = [
    "AgentInfo",
    "CreateSessionRequest",
    "Entity",
    "EntityInfo",
    "HealthResponse",
    "OrchestratorInfo",
    "RunEntityRequest",
    "SessionMessagesResponse",
    "SessionInfo",
    "StatsResponse",
    "WebUIStreamEvent",
    "WorkflowInfo",
]
