"""
Configuration-related Pydantic schemas
"""
from typing import Dict, Optional, Any, List
from pydantic import BaseModel, Field

class DatabaseConfig(BaseModel):
    url: str = Field(..., description="数据库连接URL")

class NewApiConfig(BaseModel):
    api_key: str = Field(..., description="New API Key")
    base_url: Optional[str] = None

class AlphaEngineConfig(BaseModel):
    api_url: str = Field(..., description="AlphaEngine API URL")
    api_key: str = Field(..., description="AlphaEngine API Key")
    token: str = Field(..., description="AlphaEngine Token")
    user_id: int = Field(999041, description="User ID")
    model: str = Field("deepseek-r1", description="Model name")
    using_indicator: bool = Field(True, description="Whether to use indicators")
    start_time: str = Field("2024-01-01", description="Start time for data query")
    doc_show_type: List[str] = Field(["A001", "A002", "A003", "A004"], description="Document types")
    simple_tracking: bool = Field(True, description="Whether to enable simple tracking")

class DataSourceConfig(BaseModel):
    api_key: str = Field(..., description="数据源API Key")

class ConfigResponse(BaseModel):
    database: DatabaseConfig
    new_api: NewApiConfig
    alpha_engine: Optional[AlphaEngineConfig] = None
    data_sources: Dict[str, DataSourceConfig]

class ConfigUpdateRequest(BaseModel):
    database: Optional[DatabaseConfig] = None
    new_api: Optional[NewApiConfig] = None
    alpha_engine: Optional[AlphaEngineConfig] = None
    data_sources: Optional[Dict[str, DataSourceConfig]] = None

class ConfigTestRequest(BaseModel):
    config_type: str
    config_data: Dict[str, Any]

class ConfigTestResponse(BaseModel):
    success: bool
    message: str