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


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):
    """多步驟分析請求（更新為支持動態類別）"""
    query: str
    model: Optional[str] = None
    temperature: Optional[float] = 0.7
    chain_count: int = 5  # 用戶可設定要生成多少條因果鏈

    # 新增：動態類別支持
    category_mode: Optional[str] = "dynamic_auto"  # CategoryMode enum 值
    custom_categories: Optional[List[str]] = None
    suggested_category_count: int = 3  # 建議 LLM 生成的類別數量


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"
    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]  # {類別名: 選中屬性}