from typing import Union, Optional, List

from pydantic import BaseModel, ConfigDict
from pydantic.alias_generators import to_camel


class HttpTransport(BaseModel):
    url: str
    enable_persona_auth: bool = True

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)


class StdioTransport(BaseModel):
    command: str
    args: list[str]
    env: Optional[dict[str, str]] = None
    cwd: Optional[str] = None


McpTransport = Union[HttpTransport, StdioTransport]


class ToolInvokeRequest(BaseModel):
    mcp_server_name: str
    transport: McpTransport
    tool_name: str
    args: dict
    session_id: str
    credentials_id: str | None = None
    user_id: str | None = None
    billing_disabled: bool = False

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class ToolInvokeResponse(BaseModel):
    results: list = []

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class RetrieveResourceRequest(BaseModel):
    mcp_server_name: str
    transport: McpTransport
    uri: str
    session_id: str
    credentials_id: str | None = None
    user_id: str | None = None
    billing_disabled: bool = False

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class RetrieveResourceResponse(BaseModel):
    contents: list = []

    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True)


class GenerationRequest(BaseModel):
    agent_id: str | None = None
    user_message: str | dict | list = ""
    initial_context: dict | None = None
    response_schema: dict | None = None
    parent_session_id: str | None = None
    user_id: str | None = None

    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
