"""Pydantic models for launcher configuration and state."""

from pathlib import Path
from typing import Optional

from pydantic import BaseModel, Field


class AgentConfig(BaseModel):
    """Configuration for a single Mercury agent."""

    name: str
    mcp_api_key: str
    mercury_url: str = "https://mercury-studio-production.up.railway.app/api/v1/mcp"
    working_directory: Path = Field(default_factory=lambda: Path.home())
    model: Optional[str] = None
    idle_timeout_seconds: int = 300
    cooldown_seconds: int = 5
    max_budget_usd: float = 2.0


class LauncherConfig(BaseModel):
    """Top-level launcher configuration."""

    agents: list[AgentConfig] = Field(default_factory=list)
    log_level: str = "INFO"
    log_dir: Optional[Path] = None


class AgentContext(BaseModel):
    """Agent identity context returned by mercury_get_agent_context."""

    name: str = ""
    agent_role: str = ""
    system_prompt: str = ""
    connected_agents: list[dict] = Field(default_factory=list)
    open_tasks: list[dict] = Field(default_factory=list)
    composio_toolkits: list[str] = Field(default_factory=list)
    org_id: str = ""
    org_name: str = ""


class DiscoveredAgent(BaseModel):
    """An agent discovered during startup via mercury_get_agent_context."""

    config: AgentConfig
    context: AgentContext


class OrgInfo(BaseModel):
    """Organization grouping for discovered agents."""

    org_name: str
    agents: list[DiscoveredAgent] = Field(default_factory=list)
