backend/app/database.py

"""
DuckDB 数据库管理模块
实现数据库初始化、Hugging Face Dataset 同步、单例连接管理
"""

import os
import logging
from pathlib import Path
from typing import Optional
from functools import lru_cache

from dotenv import load_dotenv

# 加载 .env 文件
load_dotenv()

import duckdb
from huggingface_hub import hf_hub_download, upload_file, login

logger = logging.getLogger(__name__)

# 环境变量配置
DUCKDB_PATH = os.getenv("DUCKDB_PATH", "/app/data/stock_data.duckdb")
HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO_ID = os.getenv("DATASET_REPO_ID", "")
HF_HOME = os.getenv("HF_HOME", "/tmp/.huggingface")


class DatabaseManager:
    """DuckDB 数据库管理器 - 单例模式"""
    
    _instance: Optional['DatabaseManager'] = None
    _connection: Optional[duckdb.DuckDBPyConnection] = None
    
    def __new__(cls) -> 'DatabaseManager':
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    @property
    def conn(self) -> duckdb.DuckDBPyConnection:
        """获取数据库连接"""
        if self._connection is None:
            # 确保目录存在
            os.makedirs(os.path.dirname(DUCKDB_PATH), exist_ok=True)
            self._connection = duckdb.connect(DUCKDB_PATH)
            
            # 如果在 HF 环境，配置远程访问插件
            if HF_TOKEN and DATASET_REPO_ID:
                try:
                    self._connection.execute("INSTALL httpfs; LOAD httpfs;")
                    self._connection.execute(f"SET s3_region='us-east-1';") # HF 默认区域
                    # 这里的逻辑可以根据是否切换到 Parquet 进一步扩展
                except Exception as e:
                    logger.warning(f"Failed to load httpfs extension: {e}")
                    
            logger.info(f"Database connection established: {DUCKDB_PATH}")
        return self._connection
    
    def close(self) -> None:
        """关闭数据库连接"""
        if self._connection is not None:
            self._connection.close()
            self._connection = None
            logger.info("Database connection closed")
    
    def init_db(self) -> None:
        """
        初始化数据库
        - 在本地环境：使用本地 DuckDB
        - 在 HF 环境：挂载云端 Parquet 文件
        """
        conn = self.conn
        
        # 检查是否在 HF 环境
        is_hf = os.getenv("SPACES_REPO_NAME") is not None or (HF_TOKEN and DATASET_REPO_ID)
        
        if is_hf:
            logger.info("HF environment detected. Downloading remote Parquet files...")
            try:
                from huggingface_hub import list_repo_files, hf_hub_download
                
                # 1. 动态获取文件列表
                all_files = list_repo_files(repo_id=DATASET_REPO_ID, repo_type="dataset")
                
                # 2. 股票列表：下载并加载
                list_file = "data/stock_list.parquet"
                if list_file in all_files:
                    local_list_path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=list_file, repo_type="dataset")
                    conn.execute(f"CREATE OR REPLACE TABLE stock_list AS SELECT * FROM read_parquet('{local_list_path}')")
                    logger.info("Local stock_list table created from remote parquet")
                
                # 3. 日线数据：下载并加载 (下载模式替代挂载模式)
                # 注意：如果文件很多，这步可能会花点时间，但比下载整个 .duckdb 快得多
                parquet_files = [f for f in all_files if f.startswith("data/parquet/") and f.endswith(".parquet")]
                
                if parquet_files:
                    local_paths = []
                    for f in parquet_files:
                        # HF SDK 会自动处理缓存，已下载的文件不会重复下载
                        path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=f, repo_type="dataset")
                        local_paths.append(f"'{path}'")
                    
                    files_sql = ", ".join(local_paths)
                    conn.execute("DROP VIEW IF EXISTS stock_daily")
                    conn.execute("DROP TABLE IF EXISTS stock_daily")
                    conn.execute(f"CREATE OR REPLACE VIEW stock_daily AS SELECT * FROM read_parquet([{files_sql}])")
                    logger.info(f"Remote stock_daily view created with {len(parquet_files)} partitions")
                else:
                    logger.warning("No parquet data files found in dataset")
                    self._create_tables()
                
                # 验证
                count = conn.execute("SELECT COUNT(*) FROM stock_list").fetchone()[0]
                logger.info(f"Verification successful: {count} stocks found")
                
            except Exception as e:
                logger.error(f"Failed to load remote Parquet: {e}")
                self._create_tables()
        else:
            self._create_tables()
            logger.info("Local database initialized")
    
    def _download_from_hf(self) -> None:
        """从 Hugging Face Dataset 下载数据库文件"""
        if not HF_TOKEN or not DATASET_REPO_ID:
            logger.warning("HF_TOKEN or DATASET_REPO_ID not set, skipping download")
            return
        
        try:
            # 登录 Hugging Face
            login(token=HF_TOKEN)
            
            # 下载数据库文件
            downloaded_path = hf_hub_download(
                repo_id=DATASET_REPO_ID,
                filename="stock_data.duckdb",
                repo_type="dataset",
                local_dir=os.path.dirname(DUCKDB_PATH),
            )
            
            # 如果下载路径不同，重命名
            if downloaded_path != DUCKDB_PATH:
                os.rename(downloaded_path, DUCKDB_PATH)
            
            logger.info(f"Database downloaded from HF Dataset: {DATASET_REPO_ID}")
        except Exception as e:
            logger.error(f"Failed to download database from HF: {e}")
    
    def upload_db(self) -> None:
        """上传 Parquet 分区到 Hugging Face Dataset (不再上传全量 .duckdb)"""
        if not HF_TOKEN or not DATASET_REPO_ID:
            logger.warning("HF_TOKEN or DATASET_REPO_ID not set, skipping upload")
            return
        
        try:
            # 先关闭连接
            self.close()
            login(token=HF_TOKEN)
            
            # 1. 上传股票列表 (这个很小，可以全量覆盖)
            if Path(DUCKDB_PATH).exists():
                # 导出列表为 parquet 方便云端读取
                conn = duckdb.connect(DUCKDB_PATH)
                list_path = os.path.join(os.path.dirname(DUCKDB_PATH), "stock_list.parquet")
                conn.execute(f"COPY stock_list TO '{list_path}' (FORMAT PARQUET)")
                conn.close()
                
                upload_file(
                    path_or_fileobj=list_path,
                    path_in_repo="data/stock_list.parquet",
                    repo_id=DATASET_REPO_ID,
                    repo_type="dataset",
                )
            
            # 2. 上传所有 Parquet 行情文件 (增量核心)
            parquet_dir = Path(os.path.dirname(DUCKDB_PATH)) / "parquet"
            if parquet_dir.exists():
                for p_file in parquet_dir.glob("*.parquet"):
                    upload_file(
                        path_or_fileobj=str(p_file),
                        path_in_repo=f"data/parquet/{p_file.name}",
                        repo_id=DATASET_REPO_ID,
                        repo_type="dataset",
                    )
            
            logger.info(f"Parquet files uploaded to HF Dataset: {DATASET_REPO_ID}")
        except Exception as e:
            logger.error(f"Failed to upload to HF: {e}")
        finally:
            _ = self.conn
    
    def _create_tables(self) -> None:
        """创建数据库表结构"""
        conn = self.conn
        
        # 日线行情表
        conn.execute("""
            CREATE TABLE IF NOT EXISTS stock_daily (
                code VARCHAR,
                trade_date DATE,
                open DOUBLE,
                high DOUBLE,
                low DOUBLE,
                close DOUBLE,
                volume BIGINT,
                amount DOUBLE,
                pct_chg DOUBLE,
                turnover_rate DOUBLE,
                PRIMARY KEY (code, trade_date)
            )
        """)
        
        # 股票基础信息表
        conn.execute("""
            CREATE TABLE IF NOT EXISTS stock_list (
                code VARCHAR PRIMARY KEY,
                name VARCHAR,
                market VARCHAR,
                list_date DATE
            )
        """)
        
        # 创建索引
        conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_code_date 
            ON stock_daily (code, trade_date)
        """)
        
        logger.info("Database tables created/verified")


# 全局单例实例
db_manager = DatabaseManager()


def get_db() -> DatabaseManager:
    """获取数据库管理器实例"""
    return db_manager