from enum import StrEnum
from typing import Literal

from litellm.types.utils import ChatCompletionMessageToolCall, Usage
from pydantic import BaseModel

from judy_core.messages import Task
from judy_core.types import ContentUnion


class OutputEventKind(StrEnum):
    TOOL_CALL = "tool_call"
    TASK_FAIL = "task_fail"
    TASK_SUCCESS = "task_success"
    CREATE_AGENT_FOLDER = "create_agent_folder"
    DEBUG_MODEL_INPUT = "debug_llm_input"
    DEBUG_MODEL_OUTPUT = "debug_llm_output"
    DEBUG_TOOL_RESULT = "debug_tool_result"
    DEBUG_TASK_LOG = "debug_task_log"


class ToolCall(BaseModel):
    kind: Literal[OutputEventKind.TOOL_CALL] = OutputEventKind.TOOL_CALL
    source: str
    tool_call_id: str
    name: str
    arguments: str


class ToolResult(BaseModel):
    id: str
    name: str
    content: ContentUnion


class TaskFailResult(BaseModel):
    kind: Literal[OutputEventKind.TASK_FAIL] = OutputEventKind.TASK_FAIL
    source: str = ""
    content: str


class TaskSuccessResult(BaseModel):
    kind: Literal[OutputEventKind.TASK_SUCCESS] = OutputEventKind.TASK_SUCCESS
    content: str


class CreateAgentFolder(BaseModel):
    kind: Literal[OutputEventKind.CREATE_AGENT_FOLDER] = (
        OutputEventKind.CREATE_AGENT_FOLDER
    )
    folder_name: str


class DebugModelInput(BaseModel):
    kind: Literal[OutputEventKind.DEBUG_MODEL_INPUT] = OutputEventKind.DEBUG_MODEL_INPUT
    source: str
    role: str
    content: str


class DebugModelOutput(BaseModel):
    kind: Literal[OutputEventKind.DEBUG_MODEL_OUTPUT] = (
        OutputEventKind.DEBUG_MODEL_OUTPUT
    )
    source: str
    response_id: str
    model: str | None
    tool_calls: list[ChatCompletionMessageToolCall] | None
    content: str | None
    reasoning_content: str | None
    usage: Usage | None
    hidden_params: dict | None
    cost_money: float | None
    cost_time: float


class DebugToolResult(BaseModel):
    kind: Literal[OutputEventKind.DEBUG_TOOL_RESULT] = OutputEventKind.DEBUG_TOOL_RESULT
    source: str
    tool_call_id: str
    name: str
    content: ContentUnion
    time_cost: float


class DebugTaskLog(BaseModel):
    kind: Literal[OutputEventKind.DEBUG_TASK_LOG] = OutputEventKind.DEBUG_TASK_LOG
    level: str = "warning"
    content: str


type InputEvent = Task | ToolResult
type OutputEvent = (
    ToolCall
    | TaskFailResult
    | TaskSuccessResult
    | CreateAgentFolder
    | DebugModelInput
    | DebugModelOutput
    | DebugToolResult
    | DebugTaskLog
)
