FA3-Datafetch/backend/app/schemas.py

"""
Pydantic Schemas for FA3 Refactored Architecture
API 请求和响应的数据模型
"""
from pydantic import BaseModel, Field
from typing import Optional, Dict, List
from datetime import datetime

#  =============================================================================
# Company Schemas
# =============================================================================
class CompanyBase(BaseModel):
    market: str
    symbol: str
    company_name: str

class CompanyCreate(CompanyBase):
    pass

class CompanyResponse(CompanyBase):
    id: int
    created_at: datetime
    updated_at: datetime
    
    class Config:
        from_attributes = True


# =============================================================================
# Data Check & Fetch Schemas
# =============================================================================
class DataCheckRequest(BaseModel):
    market: str
    symbol: str
    data_source: str

class DataCheckResponse(BaseModel):
    has_data: bool
    company_id: Optional[int] = None
    data_source: str
    last_update: Optional[Dict] = None
    message: Optional[str] = None

class FetchDataRequest(BaseModel):
    market: str
    symbol: str
    company_name: str
    data_source: str
    force_refresh: bool = False
    currency: Optional[str] = None
    frequency: Optional[str] = "Annual"

class FetchDataResponse(BaseModel):
    update_id: int
    data_source: str
    status: str
    message: str


# =============================================================================
# Data Update Schemas
# =============================================================================
class DataUpdateBase(BaseModel):
    company_id: int
    data_source: str
    update_type: str
    status: str

class DataUpdateCreate(DataUpdateBase):
    pass

class DataUpdateResponse(DataUpdateBase):
    id: int
    started_at: datetime
    completed_at: Optional[datetime] = None
    error_message: Optional[str] = None
    data_start_date: Optional[str] = None
    data_end_date: Optional[str] = None
    fetched_tables: Optional[List[str]] = None
    row_counts: Optional[Dict[str, int]] = None
    progress_message: Optional[str] = None
    progress_percentage: Optional[int] = 0
    
    class Config:
        from_attributes = True



# =============================================================================
# Financial Data Schemas
# =============================================================================
class FinancialDataResponse(BaseModel):
    company: CompanyResponse
    data_source: str
    income_statement: List[Dict]
    balance_sheet: List[Dict]
    cash_flow: List[Dict]
    daily_basic: List[Dict]
    dividend: Optional[List[Dict]] = None
    repurchase: Optional[List[Dict]] = None
    employee: Optional[List[Dict]] = None
    unified_data: Optional[List[Dict]] = None


# =============================================================================
# AI Analysis Schemas
# =============================================================================
class AnalysisStartRequest(BaseModel):
    company_id: int
    data_source: str
    model: Optional[str] = "gemini-2.0-flash"

class AnalysisStartResponse(BaseModel):
    analysis_id: int
    company_id: int
    data_source: str
    status: str
    message: str

class AnalysisStatusResponse(BaseModel):
    id: int
    company_id: int
    data_source: str
    ai_model: str
    status: str
    created_at: datetime
    completed_at: Optional[datetime] = None
    error_message: Optional[str] = None
    
    class Config:
        from_attributes = True

class AnalysisResultResponse(BaseModel):
    id: int
    company_id: int
    data_source: str
    ai_model: str
    status: str
    company_profile: Optional[str] = None
    fundamental_analysis: Optional[str] = None
    insider_analysis: Optional[str] = None
    bullish_analysis: Optional[str] = None
    bearish_analysis: Optional[str] = None
    total_tokens: int
    tokens_by_section: Optional[Dict[str, int]] = None
    created_at: datetime
    completed_at: Optional[datetime] = None
    
    class Config:
        from_attributes = True


# =============================================================================
# Data Source Schemas
# =============================================================================
class DataSourceInfo(BaseModel):
    source: str
    available: bool
    description: str
    supported_markets: List[str]

class DataSourceListResponse(BaseModel):
    market: str
    sources: List[DataSourceInfo]


# =============================================================================
# Configuration Schemas
# =============================================================================
class ConfigUpdateRequest(BaseModel):
    key: str
    value: str

class ConfigResponse(BaseModel):
    key: str
    value: str
    updated_at: datetime
    
    class Config:
        from_attributes = True