from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from schemas.domain.divination import DerivedDivinationData


class RunStatus(str, Enum):
    SUCCESS = "success"
    FAILED = "failed"


class ToolStatus(str, Enum):
    SUCCESS = "success"
    FAILURE = "failure"
    PARTIAL = "partial"


class ErrorInfo(BaseModel):
    model_config = ConfigDict(extra="forbid")

    code: str
    message: str
    retryable: bool = False
    details: dict[str, Any] | None = None


class ToolAgentOutput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: str
    tool_call_id: str
    tool_call_args: dict[str, Any] | None = None
    status: ToolStatus
    result: str
    error: ErrorInfo | None = None


class WorkerAgentOutputLite(BaseModel):
    model_config = ConfigDict(extra="forbid")

    status: RunStatus = RunStatus.SUCCESS
    sign_level: Literal["上上签", "中上签", "中下签", "下下签"]
    conclusion: list[str] = Field(min_length=1, max_length=6)
    focus_points: list[str] = Field(default_factory=list, max_length=6)
    advice: list[str] = Field(min_length=1, max_length=6)
    keywords: list[str] = Field(min_length=3, max_length=8)
    answer: str = Field(min_length=1, max_length=4000)
    error: ErrorInfo | None = None


class FollowUpOutput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    status: RunStatus = RunStatus.SUCCESS
    answer: str = Field(min_length=1, max_length=4000)
    error: ErrorInfo | None = None


class AgentOutput(WorkerAgentOutputLite):
    model_config = ConfigDict(extra="forbid")

    divination_derived: DerivedDivinationData | None = None


WorkerAgentOutput = WorkerAgentOutputLite
RuntimeAgentOutput = AgentOutput | FollowUpOutput


def resolve_worker_output_model(
    *, runtime_mode: str
) -> type[WorkerAgentOutputLite | FollowUpOutput]:
    normalized = runtime_mode.strip().lower()
    if normalized == "follow_up":
        return FollowUpOutput
    return WorkerAgentOutputLite