from pydantic import BaseModel
from typing import Optional, List


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  # 用戶可設定要生成多少條因果鏈


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