from pydantic import BaseModel
from typing import Optional, List, Dict, Literal
from enum import Enum

# Language type for prompts
LanguageType = Literal["zh", "en"]


class AttributeNode(BaseModel):
    name: str
    category: Optional[str] = None  # 材料, 功能, 用途, 使用族群
    children: Optional[List["AttributeNode"]] = None


AttributeNode.model_rebuild()


class AnalyzeRequest(BaseModel):
    query: str
    model: Optional[str] = None
    temperature: Optional[float] = 0.7
    categories: Optional[List[str]] = None  # 如果為 None，使用預設類別


class AnalyzeResponse(BaseModel):
    query: str
    attributes: AttributeNode


class ModelListResponse(BaseModel):
    models: List[str]


# ===== Multi-step streaming schemas =====

class Step1Result(BaseModel):
    """Step 1 的結果：各類別屬性列表"""
    materials: List[str]
    functions: List[str]
    usages: List[str]
    users: List[str]


class CausalChain(BaseModel):
    """單條因果鏈"""
    material: str
    function: str
    usage: str
    user: str


class StreamAnalyzeRequest(BaseModel):
    """Multi-step analysis request (updated to support dynamic categories)"""
    query: str
    model: Optional[str] = None
    temperature: Optional[float] = 0.7
    chain_count: int = 5  # User can set how many causal chains to generate

    # Dynamic category support
    category_mode: Optional[str] = "dynamic_auto"  # CategoryMode enum value
    custom_categories: Optional[List[str]] = None
    suggested_category_count: int = 3  # Suggest LLM to generate this many categories

    # Language setting
    lang: LanguageType = "zh"


class StreamAnalyzeResponse(BaseModel):
    """最終完整結果"""
    query: str
    step1_result: Step1Result
    causal_chains: List[CausalChain]
    attributes: AttributeNode


# ===== Dynamic category system schemas =====

class CategoryMode(str, Enum):
    """類別模式"""
    FIXED_ONLY = "fixed_only"
    FIXED_PLUS_CUSTOM = "fixed_plus_custom"
    FIXED_PLUS_DYNAMIC = "fixed_plus_dynamic"  # Fixed + LLM suggested
    CUSTOM_ONLY = "custom_only"
    DYNAMIC_AUTO = "dynamic_auto"


class CategoryDefinition(BaseModel):
    """類別定義"""
    name: str
    description: Optional[str] = None
    is_fixed: bool = True  # LLM 生成的為 False
    order: int = 0


class Step0Result(BaseModel):
    """Step 0: LLM 分析建議類別"""
    categories: List[CategoryDefinition]


class DynamicStep1Result(BaseModel):
    """動態版本的 Step 1 結果"""
    attributes: Dict[str, List[str]]  # {類別名: [屬性列表]}


class DynamicCausalChain(BaseModel):
    """動態版本的因果鏈"""
    chain: Dict[str, str]  # {類別名: 選中屬性}


# ===== DAG (Directed Acyclic Graph) schemas =====

class DAGNode(BaseModel):
    """DAG 節點 - 每個屬性只出現一次"""
    id: str                    # 唯一 ID: "{category}_{index}"
    name: str                  # 顯示名稱
    category: str              # 所屬類別
    order: int                 # 欄位內位置


class DAGEdge(BaseModel):
    """DAG 邊 - 節點之間的連接"""
    source_id: str
    target_id: str


class AttributeDAG(BaseModel):
    """完整 DAG 結構"""
    query: str
    categories: List[CategoryDefinition]
    nodes: List[DAGNode]
    edges: List[DAGEdge]


class DAGRelationship(BaseModel):
    """Step 2 輸出 - 單一關係"""
    source_category: str
    source: str  # source attribute name
    target_category: str
    target: str  # target attribute name


# ===== Transformation Agent schemas =====

class TransformationRequest(BaseModel):
    """Transformation Agent request"""
    query: str                        # Original query (e.g., "bicycle")
    category: str                     # Category name (e.g., "Functions")
    attributes: List[str]             # Attribute list for this category
    model: Optional[str] = None
    temperature: Optional[float] = 0.7
    keyword_count: int = 3            # Number of new keywords to generate
    lang: LanguageType = "zh"         # Language for prompts


class TransformationDescription(BaseModel):
    """單一轉換描述"""
    keyword: str                      # 新關鍵字
    description: str                  # 與 query 結合的描述


class TransformationCategoryResult(BaseModel):
    """單一類別的轉換結果"""
    category: str
    original_attributes: List[str]    # 原始屬性
    new_keywords: List[str]           # 新生成的關鍵字
    descriptions: List[TransformationDescription]


class TransformationDAGResult(BaseModel):
    """完整 Transformation 結果"""
    query: str
    results: List[TransformationCategoryResult]


# ===== Expert Transformation Agent schemas =====

class ExpertProfile(BaseModel):
    """專家檔案"""
    id: str                           # e.g., "expert-0"
    name: str                         # e.g., "藥師"
    domain: str                       # e.g., "醫療與健康"
    perspective: Optional[str] = None # e.g., "從藥物與健康管理角度思考"


class ExpertKeyword(BaseModel):
    """專家視角生成的關鍵字"""
    keyword: str                      # 關鍵字本身
    expert_id: str                    # 哪個專家生成的
    expert_name: str                  # 專家名稱（冗餘，方便前端）
    source_attribute: str             # 來自哪個原始屬性


class ExpertTransformationDescription(BaseModel):
    """專家關鍵字的描述"""
    keyword: str
    expert_id: str
    expert_name: str
    description: str


class ExpertTransformationCategoryResult(BaseModel):
    """單一類別的轉換結果（專家版）"""
    category: str
    original_attributes: List[str]
    expert_keywords: List[ExpertKeyword]  # 所有專家生成的關鍵字
    descriptions: List[ExpertTransformationDescription]


class ExpertTransformationDAGResult(BaseModel):
    """完整轉換結果（專家版）"""
    query: str
    experts: List[ExpertProfile]      # 使用的專家列表
    results: List[ExpertTransformationCategoryResult]


class ExpertSource(str, Enum):
    """專家來源類型"""
    LLM = "llm"
    CURATED = "curated"  # 精選職業（210筆，含具體領域）
    DBPEDIA = "dbpedia"
    WIKIDATA = "wikidata"


class ExpertTransformationRequest(BaseModel):
    """Expert Transformation Agent request"""
    query: str
    category: str
    attributes: List[str]

    # Expert parameters
    expert_count: int = 3                        # Number of experts (2-8)
    keywords_per_expert: int = 1                 # Keywords per expert per attribute (1-3)
    custom_experts: Optional[List[str]] = None   # User-specified experts

    # Expert source parameters
    expert_source: ExpertSource = ExpertSource.LLM  # Expert source
    expert_language: str = "en"                     # Language for external sources

    # LLM parameters
    model: Optional[str] = None
    temperature: Optional[float] = 0.7

    # Prompt language
    lang: LanguageType = "zh"


# ===== Deduplication Agent schemas =====

class DeduplicationMethod(str, Enum):
    """去重方法"""
    EMBEDDING = "embedding"  # 向量相似度
    LLM = "llm"              # LLM 成對判斷


class DeduplicationRequest(BaseModel):
    """Deduplication request"""
    descriptions: List[ExpertTransformationDescription]
    method: DeduplicationMethod = DeduplicationMethod.EMBEDDING  # Deduplication method
    similarity_threshold: float = 0.85  # Cosine similarity threshold (0.0-1.0), only for Embedding
    model: Optional[str] = None  # Embedding/LLM model
    lang: LanguageType = "zh"  # Prompt language (for LLM method)


class DescriptionGroup(BaseModel):
    """相似描述分組"""
    group_id: str                          # "group-0", "group-1"...
    representative: ExpertTransformationDescription  # 代表描述
    duplicates: List[ExpertTransformationDescription]  # 相似描述
    similarity_scores: List[float]         # 每個重複項的相似度分數


class DeduplicationResult(BaseModel):
    """去重結果"""
    total_input: int           # 輸入描述總數
    total_groups: int          # 分組數量
    total_duplicates: int      # 重複項數量
    groups: List[DescriptionGroup]
    threshold_used: float
    method_used: DeduplicationMethod       # 使用的去重方法
    model_used: str                        # 使用的模型