from typing import Any, Dict, List, NewType, Optional, Union
from uuid import UUID

from fastapi_utils.api_model import APIModel
from pydantic import BaseModel

ApplicationID = NewType("ApplicationID", UUID)
InteractionID = NewType("InteractionID", UUID)


class Node(BaseModel):
    id: str
    name: str
    alias: Optional[str] = None
    slots: Optional[dict] = None


class Edge(BaseModel):
    src_block: Optional[str] = None
    dst_block: Optional[str] = None
    dst_port: Optional[str] = None
    alias: Optional[str] = None
    case: Union[Optional[bool], Optional[str], Optional[int]] = None


class Metadata(BaseModel):
    name: str = "Untitled_Application"
    owner: str = "anonymous"
    created_at: int = 1
    updated_at: int = 1
    deploy: bool = False
    version: Optional[str] = None
    description: Optional[str] = None
    custom_fields: Dict[str, Any] = {}


class ApplicationConfiguration(BaseModel):
    metadata: Metadata
    nodes: List[Node]
    edges: List[Edge]


class Parameter(BaseModel):
    name: str
    class_name: str
    default: Optional[Any]
    is_variable_keyword: bool = False


class TypeInfo(BaseModel):
    name: str
    alias: str
    candidates: List[str]
    slots: Optional[List[Parameter]]


class BlockInfo(BaseModel):
    name: str
    alias: str
    dir: str
    slots: List[Parameter]
    inports: List[Parameter]
    outports: List[Parameter]


class DebugInfo(BaseModel):
    block: str
    output: str


class UsageInfo(BaseModel):
    block: str
    prompt_tokens: int
    completion_tokens: int
    total_cost: float


class ApplicationCreate(APIModel):
    configuration: ApplicationConfiguration


class ApplicationRun(APIModel):
    input: Union[str, List[str], Dict[str, str]]


class ApplicationCreateResponse(APIModel):
    id: ApplicationID


class ApplicationRunResponse(APIModel):
    id: InteractionID
    output: Optional[str] = None
    trace: Optional[dict] = None
    usage: Optional[List[UsageInfo]] = None


class ApplicationUsageResponse(APIModel):
    id: ApplicationID
    successful_interactions: int
    successful_llm_calls: int
    prompt_tokens: int
    completion_tokens: int
    total_cost: float


class InteractionDebugResponse(APIModel):
    application: ApplicationID
    interaction: InteractionID
    debug: List[DebugInfo]
    usage: Optional[List[UsageInfo]] = None


class ApplicationTypesResponse(APIModel):
    types: List[TypeInfo]


class ApplicationBlocksResponse(APIModel):
    blocks: List[BlockInfo]


class ItemDeleteResponse(APIModel):
    success: bool
    message: str


class ItemUpdateResponse(APIModel):
    success: bool
    message: str


class User(APIModel):
    user: str


class AppWithLatestCall(BaseModel):
    application_id: str
    name: str
    owner: str
    created_at: int
    description: Optional[str] = None
    deploy: bool
    version: Optional[str] = None
    custom_fields: dict
    latest_called_at: Optional[int] = None


class ApplicationListResponse(APIModel):
    applications: List[AppWithLatestCall]
    total: int