"""
统一数据加载器，基于 DuckDB 支持多种数据格式
"""

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

import duckdb
import pandas as pd

logger = logging.getLogger("super_datascope")


class DataLoader:
    """统一数据加载器，基于 DuckDB 支持多种数据格式"""

    SUPPORTED_FORMATS = {
        '.parquet': 'parquet',
        '.csv': 'csv',
        '.sqlite': 'sqlite',
        '.db': 'sqlite',
    }

    @staticmethod
    def load_file(path: str, table_name: Optional[str] = None) -> pd.DataFrame:
        """根据文件扩展名自动选择加载方式"""
        path = str(path)
        ext = Path(path).suffix.lower()
        fmt = DataLoader.SUPPORTED_FORMATS.get(ext)

        if fmt is None:
            raise ValueError(f"不支持的文件格式: {ext}，支持: {list(DataLoader.SUPPORTED_FORMATS.keys())}")

        if fmt == 'parquet':
            return DataLoader.load_parquet(path)
        elif fmt == 'csv':
            return DataLoader.load_csv(path)
        elif fmt == 'sqlite':
            return DataLoader.load_sqlite(path, table_name)
        else:
            raise ValueError(f"未知格式: {fmt}")

    @staticmethod
    def load_parquet(path: str) -> pd.DataFrame:
        """使用 DuckDB 读取 parquet 文件"""
        conn = duckdb.connect()
        try:
            df = conn.execute(f"SELECT * FROM read_parquet('{path}')").df()
            logger.info(f"加载 parquet: {path}, {df.shape[0]} 行 x {df.shape[1]} 列")
            return df
        finally:
            conn.close()

    @staticmethod
    def load_csv(path: str) -> pd.DataFrame:
        """使用 DuckDB 读取 CSV 文件"""
        conn = duckdb.connect()
        try:
            df = conn.execute(f"SELECT * FROM read_csv_auto('{path}')").df()
            logger.info(f"加载 CSV: {path}, {df.shape[0]} 行 x {df.shape[1]} 列")
            return df
        finally:
            conn.close()

    @staticmethod
    def load_sqlite(path: str, table_name: Optional[str] = None) -> pd.DataFrame:
        """使用 DuckDB 读取 SQLite 文件中的指定表"""
        conn = duckdb.connect()
        try:
            conn.execute(f"ATTACH '{path}' AS sqlite_db (TYPE SQLITE)")

            if table_name is None:
                tables = DataLoader.list_sqlite_tables(path)
                if not tables:
                    raise ValueError(f"SQLite 文件中没有数据表: {path}")
                table_name = tables[0]
                logger.info(f"未指定表名，使用第一个表: {table_name}")

            df = conn.execute(f"SELECT * FROM sqlite_db.{table_name}").df()
            logger.info(f"加载 SQLite: {path} -> {table_name}, {df.shape[0]} 行 x {df.shape[1]} 列")
            return df
        finally:
            conn.close()

    @staticmethod
    def list_sqlite_tables(path: str) -> list[str]:
        """列出 SQLite 文件中的所有表名"""
        conn = duckdb.connect()
        try:
            conn.execute(f"ATTACH '{path}' AS sqlite_db (TYPE SQLITE)")
            result = conn.execute(
                "SELECT table_name FROM information_schema.tables "
                "WHERE table_catalog='sqlite_db' AND table_schema='main' "
                "ORDER BY table_name"
            ).fetchall()
            return [row[0] for row in result]
        finally:
            conn.close()

    @staticmethod
    def get_file_info(path: str) -> dict:
        """获取文件基本信息"""
        p = Path(path)
        ext = p.suffix.lower()
        fmt = DataLoader.SUPPORTED_FORMATS.get(ext, 'unknown')
        size = p.stat().st_size if p.exists() else 0

        info = {
            'path': str(p),
            'name': p.name,
            'size': size,
            'size_human': _format_size(size),
            'format': fmt,
            'extension': ext,
        }

        # 对 SQLite 文件列出表名
        if fmt == 'sqlite' and p.exists():
            try:
                info['tables'] = DataLoader.list_sqlite_tables(str(p))
            except Exception:
                info['tables'] = []

        return info


def scan_directory(path: str, recursive: bool = True) -> list[dict]:
    """深度遍历目录，返回所有支持的数据文件

    Returns:
        list[dict]: 每个文件的信息，包含 path, name, size, format
    """
    root = Path(path)
    if not root.is_dir():
        raise ValueError(f"路径不是目录: {path}")

    supported_exts = set(DataLoader.SUPPORTED_FORMATS.keys())
    files = []

    pattern = '**/*' if recursive else '*'
    for p in root.glob(pattern):
        if p.is_file() and p.suffix.lower() in supported_exts:
            files.append(DataLoader.get_file_info(str(p)))

    # 按文件名排序
    files.sort(key=lambda f: f['name'])
    logger.info(f"扫描目录 {path}: 找到 {len(files)} 个数据文件")
    return files


def _format_size(size_bytes: int) -> str:
    """格式化文件大小"""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"
