from __future__ import annotations

from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field


class RouteName(StrEnum):
    DIRECT_ANSWER = "direct_answer"
    MILVUS_SEARCH = "milvus_search"
    SQL_QUERY = "sql_query"
    WEB_SEARCH = "web_search"
    MULTI_SOURCE = "multi_source"
    CLARIFY = "clarify"
    REFUSE = "refuse"


class KnowledgeType(StrEnum):
    """每条证据所属的知识来源类型，用于路由决策与 Grader 来源覆盖检查。"""

    STRUCTURED = "structured"  # SQLite / 工单结构化数据
    INTERNAL_DOC = "internal_doc"  # Milvus / 内部排障手册
    REALTIME = "realtime"  # Web / 厂商官网最新公告


class TaskType(StrEnum):
    """用户问题的高层意图分类，与路由选择对齐。"""

    DIAGNOSE = "diagnose"  # 综合诊断：需要三源证据
    HISTORICAL_QUERY = "historical_query"  # 历史工单统计
    DOCS_QA = "docs_qa"  # 内部手册问答
    VENDOR_STATUS = "vendor_status"  # 厂商当前状态
    TEXT_TRANSFORM = "text_transform"  # 文本改写
    UNKNOWN = "unknown"


# 把数据源到 KnowledgeType 的映射固化在 schemas 层，避免下游再硬编码字符串。
SOURCE_TO_KNOWLEDGE: dict[RouteName, KnowledgeType] = {
    RouteName.SQL_QUERY: KnowledgeType.STRUCTURED,
    RouteName.MILVUS_SEARCH: KnowledgeType.INTERNAL_DOC,
    RouteName.WEB_SEARCH: KnowledgeType.REALTIME,
}


class RouteDecision(BaseModel):
    needs_retrieval: bool
    intent: str
    task_type: TaskType = TaskType.UNKNOWN
    routes: list[RouteName] = Field(default_factory=list)
    knowledge_types: list[KnowledgeType] = Field(default_factory=list)
    requires_decomposition: bool = False
    confidence: str = "medium"
    reason_code: str
    filters: dict[str, Any] = Field(default_factory=dict)
    max_rounds: int = 2
    fallback: str = "clarify"


class PlanStep(BaseModel):
    id: str
    tool: RouteName
    query: str
    arguments: dict[str, Any] = Field(default_factory=dict)
    depends_on: list[str] = Field(default_factory=list)


class RetrievalPlan(BaseModel):
    goal: str
    steps: list[PlanStep]


class Evidence(BaseModel):
    source_type: str
    source: str
    content: str
    score: float | None = None
    knowledge_type: KnowledgeType | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class ToolResult(BaseModel):
    status: str
    tool: str
    data: list[dict[str, Any]] = Field(default_factory=list)
    evidence: list[Evidence] = Field(default_factory=list)
    latency_ms: int = 0
    retryable: bool = False
    error_code: str | None = None
    error_message: str | None = None


class QualityGrade(BaseModel):
    relevant: bool
    sufficient: bool
    missing_aspects: list[str] = Field(default_factory=list)
    conflict: bool = False
    recommended_action: str
    reason: str


class QueryRequest(BaseModel):
    query: str = Field(min_length=1, max_length=4000)
    session_id: str = "default"
    debug: bool = False


class QueryResponse(BaseModel):
    answer: str
    citations: list[Evidence]
    route: RouteDecision
    executed_queries: list[dict[str, Any]] = Field(default_factory=list)
    trace: list[dict[str, Any]] = Field(default_factory=list)
    termination_reason: str