"""
文件查询工具模块 (File Query Tool)

该模块提供文件查询功能，支持读取 CSV 和 Excel 文件的表头信息、
统计信息和样本数据，而不需要读取完整数据。
"""

import os
import pandas as pd
from dataclasses import dataclass, asdict
from typing import List, Optional, Dict, Any
from pathlib import Path


@dataclass
class ColumnInfo:
    """列信息数据类"""
    name: str  # 列名
    dtype: str  # 数据类型
    non_null_count: int  # 非空值数量
    null_count: int  # 空值数量
    unique_count: Optional[int] = None  # 唯一值数量（可选）
    sample_values: Optional[List[Any]] = None  # 样本值（可选）
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return asdict(self)


@dataclass
class SampleData:
    """样本数据数据类"""
    rows: List[Dict[str, Any]]  # 样本行数据
    row_count: int  # 样本行数
    columns: List[str]  # 列名列表
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "rows": self.rows,
            "row_count": self.row_count,
            "columns": self.columns
        }


@dataclass
class FileMetadata:
    """文件元数据"""
    file_path: str  # 文件路径
    file_name: str  # 文件名
    file_size: int  # 文件大小（字节）
    file_type: str  # 文件类型（csv/excel）
    extension: str  # 文件扩展名
    total_rows: Optional[int] = None  # 总行数
    total_columns: Optional[int] = None  # 总列数
    columns: Optional[List[ColumnInfo]] = None  # 列信息列表
    sample_data: Optional[SampleData] = None  # 样本数据
    read_rows: int = 100  # 实际读取的行数
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        result = asdict(self)
        if self.columns:
            result["columns"] = [col.to_dict() for col in self.columns]
        if self.sample_data:
            result["sample_data"] = self.sample_data.to_dict()
        return result


class FileQueryTool:
    """
    文件查询工具
    
    用于读取 CSV 和 Excel 文件的表头信息、统计信息和样本数据。
    支持只读取前 N 行，避免加载完整数据。
    
    Attributes:
        default_read_rows: 默认读取的行数（100 行）
    
    Example:
        >>> tool = FileQueryTool()
        >>> metadata = tool.query("data.csv")
        >>> print(metadata.columns)
        >>> print(metadata.sample_data)
    """
    
    def __init__(self, default_read_rows: int = 100):
        """
        初始化文件查询工具
        
        Args:
            default_read_rows: 默认读取的行数，用于获取表头和样本数据
        """
        self.default_read_rows = default_read_rows
    
    def query(self, file_path: str, read_rows: Optional[int] = None) -> FileMetadata:
        """
        查询文件信息
        
        读取文件的前 N 行，提取表头信息、统计信息和样本数据。
        
        Args:
            file_path: 文件路径
            read_rows: 要读取的行数，如果为 None 则使用默认值
            
        Returns:
            FileMetadata: 文件元数据对象
            
        Raises:
            FileNotFoundError: 文件不存在
            ValueError: 不支持的文件格式
            Exception: 其他读取错误
        """
        if read_rows is None:
            read_rows = self.default_read_rows
        
        # 验证文件存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在：{file_path}")
        
        # 获取文件信息
        file_path = os.path.abspath(file_path)
        file_name = os.path.basename(file_path)
        file_size = os.path.getsize(file_path)
        extension = os.path.splitext(file_name)[1].lower()
        
        # 确定文件类型
        if extension == ".csv":
            file_type = "csv"
            df = self._read_csv(file_path, read_rows)
        elif extension in [".xlsx", ".xls"]:
            file_type = "excel"
            df = self._read_excel(file_path, read_rows)
        else:
            raise ValueError(f"不支持的文件格式：{extension}，仅支持 CSV 和 Excel 文件")
        
        # 提取列信息
        columns = self._extract_column_info(df)
        
        # 提取样本数据
        sample_data = self._extract_sample_data(df, read_rows)
        
        # 创建元数据对象
        metadata = FileMetadata(
            file_path=file_path,
            file_name=file_name,
            file_size=file_size,
            file_type=file_type,
            extension=extension,
            total_rows=len(df),
            total_columns=len(df.columns),
            columns=columns,
            sample_data=sample_data,
            read_rows=min(len(df), read_rows)
        )
        
        return metadata
    
    def _read_csv(self, file_path: str, read_rows: int) -> pd.DataFrame:
        """
        读取 CSV 文件
        
        Args:
            file_path: CSV 文件路径
            read_rows: 要读取的行数
            
        Returns:
            pd.DataFrame: 读取的数据框
        """
        df = pd.read_csv(file_path, nrows=read_rows)
        return df
    
    def _read_excel(self, file_path: str, read_rows: int) -> pd.DataFrame:
        """
        读取 Excel 文件
        
        Args:
            file_path: Excel 文件路径
            read_rows: 要读取的行数
            
        Returns:
            pd.DataFrame: 读取的数据框
        """
        df = pd.read_excel(file_path, nrows=read_rows)
        return df
    
    def _extract_column_info(self, df: pd.DataFrame) -> List[ColumnInfo]:
        """
        提取列信息
        
        Args:
            df: 数据框
            
        Returns:
            List[ColumnInfo]: 列信息列表
        """
        columns = []
        
        for col in df.columns:
            # 获取数据类型
            dtype = str(df[col].dtype)
            
            # 统计非空值和空值
            non_null_count = df[col].notna().sum()
            null_count = df[col].isna().sum()
            
            # 计算唯一值数量（仅对非空值）
            unique_count = df[col].dropna().nunique()
            
            # 获取样本值（最多 5 个）
            sample_values = df[col].dropna().head(5).tolist()
            
            column_info = ColumnInfo(
                name=col,
                dtype=dtype,
                non_null_count=int(non_null_count),
                null_count=int(null_count),
                unique_count=int(unique_count),
                sample_values=sample_values
            )
            columns.append(column_info)
        
        return columns
    
    def _extract_sample_data(self, df: pd.DataFrame, read_rows: int) -> SampleData:
        """
        提取样本数据
        
        Args:
            df: 数据框
            read_rows: 实际读取的行数
            
        Returns:
            SampleData: 样本数据对象
        """
        # 获取实际行数（可能小于 read_rows）
        actual_rows = min(len(df), read_rows)
        
        # 转换为字典列表
        rows = df.head(actual_rows).to_dict(orient='records')
        
        return SampleData(
            rows=rows,
            row_count=actual_rows,
            columns=df.columns.tolist()
        )
    
    def get_column_names(self, file_path: str, read_rows: Optional[int] = None) -> List[str]:
        """
        快速获取列名列表
        
        Args:
            file_path: 文件路径
            read_rows: 要读取的行数
            
        Returns:
            List[str]: 列名列表
        """
        metadata = self.query(file_path, read_rows)
        return [col.name for col in metadata.columns]
    
    def get_column_dtype(self, file_path: str, column_name: str, 
                         read_rows: Optional[int] = None) -> str:
        """
        获取指定列的数据类型
        
        Args:
            file_path: 文件路径
            column_name: 列名
            read_rows: 要读取的行数
            
        Returns:
            str: 数据类型
        """
        metadata = self.query(file_path, read_rows)
        for col in metadata.columns:
            if col.name == column_name:
                return col.dtype
        raise ValueError(f"列 '{column_name}' 不存在")
    
    def get_file_info(self, file_path: str, read_rows: Optional[int] = None) -> Dict[str, Any]:
        """
        获取文件信息的字典格式
        
        Args:
            file_path: 文件路径
            read_rows: 要读取的行数
            
        Returns:
            Dict[str, Any]: 文件信息字典
        """
        metadata = self.query(file_path, read_rows)
        return metadata.to_dict()