"""DuckDB storage for data persistence."""

import pandas as pd
import polars as pl
import duckdb
from pathlib import Path
from typing import Optional, List, Dict, Any, Tuple
from collections import defaultdict
from datetime import datetime
from src.config.settings import get_settings


# Type alias for DuckDB connection
DuckDBConnection = duckdb.DuckDBPyConnection


# Default column type mapping for automatic schema inference
DEFAULT_TYPE_MAPPING = {
    "ts_code": "VARCHAR(16)",
    "trade_date": "DATE",
    "open": "DOUBLE",
    "high": "DOUBLE",
    "low": "DOUBLE",
    "close": "DOUBLE",
    "pre_close": "DOUBLE",
    "change": "DOUBLE",
    "pct_chg": "DOUBLE",
    "vol": "DOUBLE",
    "amount": "DOUBLE",
    "turnover_rate": "DOUBLE",
    "volume_ratio": "DOUBLE",
    "adj_factor": "DOUBLE",
    "suspend_flag": "INTEGER",
}


class Storage:
    """DuckDB storage manager for saving and loading data.

    迁移说明：
    - 保持 API 完全兼容，调用方无需修改
    - 新增 load_polars() 方法支持 Polars 零拷贝导出
    - 使用单例模式管理数据库连接
    - 并发写入通过队列管理（见 ThreadSafeStorage）

    注意：
    - 默认使用 read_only=True 模式，允许多进程并发读取
    - 只有在数据同步时才使用 read_only=False 模式
    """

    _instance = None
    _connection: Optional[DuckDBConnection] = None

    def __new__(cls, *args, **kwargs):
        """Singleton to ensure single connection."""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, path: Optional[Path] = None, read_only: bool = True):
        """Initialize storage.

        Args:
            path: 数据库文件路径，默认为配置中的路径
            read_only: 是否为只读模式，默认为 True
                      只有在数据同步时才需要设置为 False
        """
        if hasattr(self, "_initialized"):
            return

        cfg = get_settings()
        self.base_path = path or cfg.data_path_resolved
        self.base_path.mkdir(parents=True, exist_ok=True)
        self.db_path = self.base_path / "prostock.db"
        self._read_only = read_only

        self._init_db()
        self._initialized = True

    def _init_db(self):
        """Initialize database connection and schema.

        注意：建表语句已迁移到对应的 API 文件中，
        每个同步类负责自己的表结构定义和创建。
        参见：
        - api_daily.py: DailySync.TABLE_SCHEMA
        - api_pro_bar.py: ProBarSync.TABLE_SCHEMA
        - api_bak_basic.py: BakBasicSync.TABLE_SCHEMA
        - api_financial_sync.py: FinancialSync.TABLE_SCHEMAS
        """
        self._connection = duckdb.connect(str(self.db_path), read_only=self._read_only)

    def save(
        self,
        name: str,
        data: pd.DataFrame,
        mode: str = "append",
        use_upsert: bool = True,
    ) -> dict:
        """Save data to DuckDB.

        Args:
            name: Table name
            data: DataFrame to save
            mode: 'append' or 'replace' (DELETE + INSERT)
            use_upsert: 若为 True 使用 INSERT OR REPLACE (upsert)，
                       若为 False 使用普通 INSERT (依赖外部删除逻辑)

        Returns:
            Dict with save result
        """
        if data.empty:
            return {"status": "skipped", "rows": 0}

        # 确保日期列是正确的类型 (YYYYMMDD -> date)
        # trade_date: 日线数据日期
        if "trade_date" in data.columns:
            data = data.copy()
            data["trade_date"] = pd.to_datetime(
                data["trade_date"], format="%Y%m%d"
            ).dt.date

        # ann_date: 公告日期
        if "ann_date" in data.columns:
            data = data.copy()
            data["ann_date"] = pd.to_datetime(
                data["ann_date"], format="%Y%m%d", errors="coerce"
            ).dt.date

        # f_ann_date: 最终公告日期
        if "f_ann_date" in data.columns:
            data = data.copy()
            data["f_ann_date"] = pd.to_datetime(
                data["f_ann_date"], format="%Y%m%d", errors="coerce"
            ).dt.date

        # end_date: 报告期/期末日期
        if "end_date" in data.columns:
            data = data.copy()
            data["end_date"] = pd.to_datetime(
                data["end_date"], format="%Y%m%d", errors="coerce"
            ).dt.date

        # Register DataFrame as temporary view
        self._connection.register("temp_data", data)

        try:
            if mode == "replace":
                self._connection.execute(f"DELETE FROM {name}")

            columns = ", ".join(f'"{col}"' for col in data.columns)

            if use_upsert:
                # UPSERT: INSERT OR REPLACE (需要主键约束)
                self._connection.execute(f"""
                    INSERT OR REPLACE INTO {name} ({columns})
                    SELECT {columns} FROM temp_data
                """)
            else:
                # 普通 INSERT (依赖外部删除逻辑确保无重复)
                self._connection.execute(f"""
                    INSERT INTO {name} ({columns})
                    SELECT {columns} FROM temp_data
                """)

            row_count = len(data)
            print(f"[Storage] Saved {row_count} rows to DuckDB ({name})")
            return {"status": "success", "rows": row_count}

        except Exception as e:
            print(f"[Storage] Error saving {name}: {e}")
            return {"status": "error", "error": str(e)}
        finally:
            self._connection.unregister("temp_data")

    def load(
        self,
        name: str,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        ts_code: Optional[str] = None,
    ) -> pd.DataFrame:
        """Load data from DuckDB with query pushdown.

        关键优化：
        - WHERE 条件在数据库层过滤，无需加载全表
        - 只返回匹配条件的行，大幅减少内存占用

        Args:
            name: Table name
            start_date: Start date filter (YYYYMMDD)
            end_date: End date filter (YYYYMMDD)
            ts_code: Stock code filter

        Returns:
            Filtered DataFrame
        """
        # Build WHERE clause with parameterized queries
        conditions = []
        params = []

        if start_date and end_date:
            conditions.append("trade_date BETWEEN ? AND ?")
            # Convert to DATE type
            start = pd.to_datetime(start_date, format="%Y%m%d").date()
            end = pd.to_datetime(end_date, format="%Y%m%d").date()
            params.extend([start, end])
        elif start_date:
            conditions.append("trade_date >= ?")
            params.append(pd.to_datetime(start_date, format="%Y%m%d").date())
        elif end_date:
            conditions.append("trade_date <= ?")
            params.append(pd.to_datetime(end_date, format="%Y%m%d").date())

        if ts_code:
            conditions.append("ts_code = ?")
            params.append(ts_code)

        where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
        query = f"SELECT * FROM {name} {where_clause} ORDER BY ts_code, trade_date"

        try:
            # Execute query with parameters (SQL injection safe)
            result = self._connection.execute(query, params).fetchdf()

            # Convert trade_date back to string format for compatibility
            if "trade_date" in result.columns:
                result["trade_date"] = result["trade_date"].dt.strftime("%Y%m%d")

            return result
        except Exception as e:
            print(f"[Storage] Error loading {name}: {e}")
            return pd.DataFrame()

    def load_polars(
        self,
        name: str,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        ts_code: Optional[str] = None,
    ) -> pl.DataFrame:
        """Load data as Polars DataFrame (for DataLoader).

        性能优势：
        - 零拷贝导出（DuckDB → Polars via PyArrow）
        - 需要 pyarrow 支持
        """
        # Build query
        conditions = []
        if start_date and end_date:
            start = pd.to_datetime(start_date, format="%Y%m%d").date()
            end = pd.to_datetime(end_date, format="%Y%m%d").date()
            conditions.append(f"trade_date BETWEEN '{start}' AND '{end}'")
        if ts_code:
            conditions.append(f"ts_code = '{ts_code}'")

        where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
        query = f"SELECT * FROM {name} {where_clause} ORDER BY ts_code, trade_date"

        # 使用 DuckDB 的 Polars 导出（需要 pyarrow）
        df = self._connection.sql(query).pl()

        # 将 trade_date 转换为字符串格式，保持兼容性
        if "trade_date" in df.columns:
            df = df.with_columns(
                pl.col("trade_date").dt.strftime("%Y%m%d").alias("trade_date")
            )

        return df

    def exists(self, name: str) -> bool:
        """Check if table exists."""
        result = self._connection.execute(
            """
            SELECT COUNT(*) FROM information_schema.tables
            WHERE table_name = ?
        """,
            [name],
        ).fetchone()
        return result[0] > 0

    def delete(self, name: str) -> bool:
        """Delete a table."""
        try:
            self._connection.execute(f"DROP TABLE IF EXISTS {name}")
            print(f"[Storage] Deleted table {name}")
            return True
        except Exception as e:
            print(f"[Storage] Error deleting {name}: {e}")
            return False

    def get_last_date(self, name: str) -> Optional[str]:
        """Get the latest date in storage."""
        try:
            result = self._connection.execute(f"""
                SELECT MAX(trade_date) FROM {name}
            """).fetchone()
            if result[0]:
                # Convert date back to string format
                return (
                    result[0].strftime("%Y%m%d")
                    if hasattr(result[0], "strftime")
                    else str(result[0])
                )
            return None
        except:
            return None

    def close(self):
        """Close database connection."""
        if self._connection:
            self._connection.close()
            Storage._connection = None
            Storage._instance = None

    # ======================================================================
    # 事务支持
    # ======================================================================

    def begin_transaction(self) -> None:
        """开始事务。

        使用方式:
            storage = Storage(read_only=False)
            storage.begin_transaction()
            try:
                storage.save("table1", data1)
                storage.save("table2", data2)
                storage.commit_transaction()
            except Exception:
                storage.rollback_transaction()
                raise
        """
        if self._read_only:
            raise RuntimeError("Cannot begin transaction in read-only mode")

        try:
            self._connection.execute("BEGIN TRANSACTION")
            print("[Storage] Transaction started")
        except Exception as e:
            print(f"[Storage] Error starting transaction: {e}")
            raise

    def commit_transaction(self) -> None:
        """提交事务。"""
        if self._read_only:
            return

        try:
            self._connection.execute("COMMIT")
            print("[Storage] Transaction committed")
        except Exception as e:
            print(f"[Storage] Error committing transaction: {e}")
            raise

    def rollback_transaction(self) -> None:
        """回滚事务。"""
        if self._read_only:
            return

        try:
            self._connection.execute("ROLLBACK")
            print("[Storage] Transaction rolled back")
        except Exception as e:
            print(f"[Storage] Error rolling back transaction: {e}")

    def transaction(self):
        """事务上下文管理器。

        使用方式:
            storage = Storage(read_only=False)
            with storage.transaction():
                storage.save("table1", data1)
                storage.save("table2", data2)
            # 自动提交或回滚
        """
        return _TransactionContext(self)


class _TransactionContext:
    """事务上下文管理器。"""

    def __init__(self, storage: Storage):
        self.storage = storage

    def __enter__(self):
        self.storage.begin_transaction()
        return self.storage

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.storage.commit_transaction()
        else:
            self.storage.rollback_transaction()
        return False


class ThreadSafeStorage:
    """线程安全的 DuckDB 写入包装器。

    DuckDB 写入时不支持并发，使用队列收集写入请求，
    在 sync 结束时统一批量写入。

    注意：
    - 此类自动使用 read_only=False 模式，用于数据同步
    - 不要在多进程中同时使用此类，只应在单进程中用于批量写入
    - 支持事务模式：在事务中批量写入，确保原子性
    """

    def __init__(self):
        # 使用 read_only=False 模式创建 Storage，用于写入操作
        self.storage = Storage(read_only=False)
        self._pending_writes: List[tuple] = []  # [(name, data, use_upsert), ...]
        self._in_transaction: bool = False

    def queue_save(self, name: str, data: pd.DataFrame, use_upsert: bool = True):
        """将数据放入写入队列（不立即写入）

        Args:
            name: 表名
            data: DataFrame 数据
            use_upsert: 若为 True 使用 INSERT OR REPLACE，若为 False 使用普通 INSERT
        """
        if not data.empty:
            self._pending_writes.append((name, data, use_upsert))

    def flush(self, use_transaction: bool = True):
        """批量写入所有队列数据。

        调用时机：在 sync 结束时统一调用，避免并发写入冲突。

        Args:
            use_transaction: 若为 True，使用事务包装批量写入
        """
        if not self._pending_writes:
            return

        # 按表名和 use_upsert 分组
        table_data = defaultdict(list)

        for name, data, use_upsert in self._pending_writes:
            key = (name, use_upsert)
            table_data[key].append(data)

        # 批量写入每个表
        if use_transaction:
            # 使用事务确保原子性
            try:
                self.storage.begin_transaction()
                total_rows = 0

                for (name, use_upsert), data_list in table_data.items():
                    combined = pd.concat(data_list, ignore_index=True)
                    # 在批量数据中先去重
                    if (
                        "ts_code" in combined.columns
                        and "trade_date" in combined.columns
                    ):
                        combined = combined.drop_duplicates(
                            subset=["ts_code", "trade_date"], keep="last"
                        )
                    result = self.storage.save(
                        name, combined, mode="append", use_upsert=use_upsert
                    )
                    if result.get("status") == "success":
                        total_rows += result.get("rows", 0)

                self.storage.commit_transaction()
                print(
                    f"[ThreadSafeStorage] Transaction committed: {total_rows} rows saved"
                )

            except Exception as e:
                self.storage.rollback_transaction()
                print(f"[ThreadSafeStorage] Transaction rolled back due to error: {e}")
                raise
        else:
            # 不使用事务，逐表写入
            for (name, use_upsert), data_list in table_data.items():
                combined = pd.concat(data_list, ignore_index=True)
                # 在批量数据中先去重
                if "ts_code" in combined.columns and "trade_date" in combined.columns:
                    combined = combined.drop_duplicates(
                        subset=["ts_code", "trade_date"], keep="last"
                    )
                self.storage.save(name, combined, mode="append", use_upsert=use_upsert)

        self._pending_writes.clear()

    def begin_transaction(self) -> None:
        """开始事务模式。

        在事务模式下，flush 会被延迟到 commit_transaction 或 rollback_transaction。
        """
        self._in_transaction = True
        self.storage.begin_transaction()

    def commit_transaction(self) -> None:
        """提交事务。"""
        if self._pending_writes:
            # 先写入队列中的数据
            self.flush(use_transaction=False)
        self.storage.commit_transaction()
        self._in_transaction = False

    def rollback_transaction(self) -> None:
        """回滚事务。"""
        self._pending_writes.clear()
        self.storage.rollback_transaction()
        self._in_transaction = False

    def __getattr__(self, name):
        """代理其他方法到 Storage 实例"""
        return getattr(self.storage, name)
