"""CSV file handling module"""

import csv
from pathlib import Path
from typing import List, Dict, Any, Optional, Iterator
import pandas as pd
from datetime import datetime


class CSVHandler:
    """CSV文件处理器"""
    
    def __init__(self, encoding: str = 'utf-8'):
        """初始化CSV处理器
        
        Args:
            encoding: 文件编码
        """
        self.encoding = encoding
    
    def read_file(self, file_path: str, 
                  has_header: bool = True) -> List[Dict[str, Any]]:
        """读取CSV文件
        
        Args:
            file_path: 文件路径
            has_header: 是否有标题行
            
        Returns:
            数据行列表
        """
        path = Path(file_path)
        if not path.exists():
            print(f"文件不存在: {file_path}")
            return []
        
        try:
            with open(path, 'r', encoding=self.encoding) as f:
                if has_header:
                    reader = csv.DictReader(f)
                    return list(reader)
                else:
                    reader = csv.reader(f)
                    return [{'col_' + str(i): val for i, val in enumerate(row)} 
                           for row in reader]
        except Exception as e:
            print(f"读取CSV文件失败 {file_path}: {e}")
            return []
    
    def read_file_iter(self, file_path: str,
                      has_header: bool = True,
                      chunk_size: Optional[int] = None) -> Iterator[Dict[str, Any]]:
        """迭代读取CSV文件
        
        Args:
            file_path: 文件路径
            has_header: 是否有标题行
            chunk_size: 分块大小（如果指定则返回块）
            
        Yields:
            数据行或数据块
        """
        path = Path(file_path)
        if not path.exists():
            print(f"文件不存在: {file_path}")
            return
        
        try:
            with open(path, 'r', encoding=self.encoding) as f:
                if has_header:
                    reader = csv.DictReader(f)
                else:
                    reader = csv.reader(f)
                
                if chunk_size:
                    chunk = []
                    for row in reader:
                        if not has_header:
                            row = {'col_' + str(i): val for i, val in enumerate(row)}
                        chunk.append(row)
                        
                        if len(chunk) >= chunk_size:
                            yield chunk
                            chunk = []
                    
                    if chunk:
                        yield chunk
                else:
                    for row in reader:
                        if not has_header:
                            row = {'col_' + str(i): val for i, val in enumerate(row)}
                        yield row
                        
        except Exception as e:
            print(f"读取CSV文件失败 {file_path}: {e}")
    
    def write_file(self, file_path: str, 
                  data: List[Dict[str, Any]],
                  fieldnames: Optional[List[str]] = None,
                  append: bool = False):
        """写入CSV文件
        
        Args:
            file_path: 文件路径
            data: 数据行列表
            fieldnames: 字段名列表，如果为None则从数据推断
            append: 是否追加模式
        """
        if not data:
            print("没有数据要写入")
            return
        
        path = Path(file_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        # 推断字段名
        if fieldnames is None:
            fieldnames = list(data[0].keys())
        
        mode = 'a' if append else 'w'
        write_header = not append or not path.exists()
        
        try:
            with open(path, mode, encoding=self.encoding, newline='') as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                
                if write_header:
                    writer.writeheader()
                
                writer.writerows(data)
                
            print(f"✓ 数据已写入: {file_path} ({len(data)} 行)")
        except Exception as e:
            print(f"写入CSV文件失败 {file_path}: {e}")
            raise
    
    def write_file_incremental(self, file_path: str,
                              data_iterator: Iterator[Dict[str, Any]],
                              fieldnames: Optional[List[str]] = None,
                              batch_size: int = 1000):
        """增量写入CSV文件
        
        Args:
            file_path: 文件路径
            data_iterator: 数据迭代器
            fieldnames: 字段名列表
            batch_size: 批次大小
        """
        path = Path(file_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        first_row = None
        batch = []
        total_rows = 0
        
        try:
            for row in data_iterator:
                if first_row is None:
                    first_row = row
                    if fieldnames is None:
                        fieldnames = list(row.keys())
                
                batch.append(row)
                
                if len(batch) >= batch_size:
                    # 写入批次
                    mode = 'w' if total_rows == 0 else 'a'
                    write_header = total_rows == 0
                    
                    with open(path, mode, encoding=self.encoding, newline='') as f:
                        writer = csv.DictWriter(f, fieldnames=fieldnames)
                        if write_header:
                            writer.writeheader()
                        writer.writerows(batch)
                    
                    total_rows += len(batch)
                    print(f"  已写入 {total_rows} 行...")
                    batch = []
            
            # 写入剩余数据
            if batch:
                mode = 'w' if total_rows == 0 else 'a'
                write_header = total_rows == 0
                
                with open(path, mode, encoding=self.encoding, newline='') as f:
                    writer = csv.DictWriter(f, fieldnames=fieldnames)
                    if write_header:
                        writer.writeheader()
                    writer.writerows(batch)
                
                total_rows += len(batch)
            
            print(f"✓ 数据已写入: {file_path} (共 {total_rows} 行)")
            
        except Exception as e:
            print(f"增量写入CSV文件失败 {file_path}: {e}")
            raise
    
    def merge_csv_files(self, input_files: List[str], 
                       output_file: str,
                       remove_duplicates: bool = False,
                       duplicate_keys: Optional[List[str]] = None):
        """合并多个CSV文件
        
        Args:
            input_files: 输入文件列表
            output_file: 输出文件路径
            remove_duplicates: 是否去重
            duplicate_keys: 去重键列表
        """
        all_data = []
        
        for file_path in input_files:
            if Path(file_path).exists():
                data = self.read_file(file_path)
                all_data.extend(data)
        
        if not all_data:
            print("没有数据可合并")
            return
        
        # 去重处理
        if remove_duplicates and duplicate_keys:
            seen = set()
            unique_data = []
            
            for row in all_data:
                key = tuple(row.get(k) for k in duplicate_keys)
                if key not in seen:
                    seen.add(key)
                    unique_data.append(row)
            
            all_data = unique_data
            print(f"去重后保留 {len(all_data)} 行")
        
        # 写入合并后的文件
        self.write_file(output_file, all_data)
    
    def read_with_pandas(self, file_path: str, **kwargs) -> pd.DataFrame:
        """使用pandas读取CSV文件
        
        Args:
            file_path: 文件路径
            **kwargs: pandas read_csv参数
            
        Returns:
            DataFrame对象
        """
        try:
            return pd.read_csv(file_path, encoding=self.encoding, **kwargs)
        except Exception as e:
            print(f"使用pandas读取CSV失败 {file_path}: {e}")
            return pd.DataFrame()
    
    def write_with_pandas(self, df: pd.DataFrame, 
                         file_path: str, 
                         **kwargs):
        """使用pandas写入CSV文件
        
        Args:
            df: DataFrame对象
            file_path: 文件路径
            **kwargs: pandas to_csv参数
        """
        try:
            path = Path(file_path)
            path.parent.mkdir(parents=True, exist_ok=True)
            
            df.to_csv(file_path, encoding=self.encoding, 
                     index=False, **kwargs)
            
            print(f"✓ 数据已写入: {file_path} ({len(df)} 行)")
        except Exception as e:
            print(f"使用pandas写入CSV失败 {file_path}: {e}")
            raise
    
    def get_file_info(self, file_path: str) -> Dict[str, Any]:
        """获取CSV文件信息
        
        Args:
            file_path: 文件路径
            
        Returns:
            文件信息字典
        """
        path = Path(file_path)
        
        if not path.exists():
            return {'exists': False}
        
        try:
            # 获取基本信息
            stat = path.stat()
            
            # 读取第一行获取列信息
            with open(path, 'r', encoding=self.encoding) as f:
                reader = csv.DictReader(f)
                first_row = next(reader, None)
                columns = list(reader.fieldnames) if reader.fieldnames else []
                
                # 统计行数
                row_count = 1 if first_row else 0
                for _ in reader:
                    row_count += 1
            
            return {
                'exists': True,
                'size_bytes': stat.st_size,
                'size_mb': round(stat.st_size / (1024 * 1024), 2),
                'modified': datetime.fromtimestamp(stat.st_mtime),
                'columns': columns,
                'column_count': len(columns),
                'row_count': row_count
            }
        except Exception as e:
            return {
                'exists': True,
                'error': str(e)
            }