#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Text Extractor
文本提取器 - 重构版本，支持从多种格式文档中提取文本内容
"""

import os
import logging
import mimetypes
from typing import List, Dict, Any, Optional, Union
from pathlib import Path
from dataclasses import dataclass
from datetime import datetime

# 导入依赖库
try:
    import PyPDF2
    import pdfplumber
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False

try:
    from docx import Document
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

try:
    import openpyxl
    from openpyxl import load_workbook
    EXCEL_AVAILABLE = True
except ImportError:
    EXCEL_AVAILABLE = False

try:
    import pandas as pd
    PANDAS_AVAILABLE = True
except ImportError:
    PANDAS_AVAILABLE = False

from ..config import DocumentProcessingConfig
from ..exceptions import DocumentProcessingError, ResourceNotFoundError

logger = logging.getLogger(__name__)


@dataclass
class ExtractedDocument:
    """提取的文档数据"""
    filename: str
    file_type: str
    content: str  # 纯文本内容
    metadata: Dict[str, Any]  # 文档元数据
    extracted_at: datetime
    file_size: int
    page_count: Optional[int] = None
    extraction_method: Optional[str] = None
    error_info: Optional[str] = None
    
    def __post_init__(self):
        """初始化后处理"""
        if not self.extracted_at:
            self.extracted_at = datetime.now()
        
        # 计算内容统计
        self.metadata.setdefault('content_length', len(self.content))
        self.metadata.setdefault('word_count', len(self.content.split()))
        self.metadata.setdefault('line_count', len(self.content.splitlines()))
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'filename': self.filename,
            'file_type': self.file_type,
            'content': self.content,
            'metadata': self.metadata,
            'extracted_at': self.extracted_at.isoformat(),
            'file_size': self.file_size,
            'page_count': self.page_count,
            'extraction_method': self.extraction_method,
            'error_info': self.error_info
        }
    
    def get_summary(self) -> Dict[str, Any]:
        """获取文档摘要信息"""
        return {
            'filename': self.filename,
            'file_type': self.file_type,
            'file_size': self.file_size,
            'page_count': self.page_count,
            'content_length': len(self.content),
            'word_count': len(self.content.split()),
            'extracted_at': self.extracted_at.isoformat(),
            'extraction_method': self.extraction_method,
            'has_error': bool(self.error_info)
        }


class TextExtractor:
    """
    文本提取器 - 重构版本
    支持从多种文档格式中提取文本内容
    """
    
    SUPPORTED_EXTENSIONS = {
        '.pdf': 'PDF Document',
        '.docx': 'Microsoft Word Document',
        '.doc': 'Microsoft Word Document (Legacy)',
        '.xlsx': 'Microsoft Excel Spreadsheet',
        '.xls': 'Microsoft Excel Spreadsheet (Legacy)',
        '.txt': 'Plain Text',
        '.md': 'Markdown',
        '.csv': 'Comma Separated Values',
        '.json': 'JSON Document',
        '.xml': 'XML Document',
        '.html': 'HTML Document',
        '.htm': 'HTML Document'
    }
    
    def __init__(self, config: DocumentProcessingConfig):
        """
        初始化文本提取器
        
        Args:
            config: 文档处理配置
        """
        self.config = config
        self.extraction_methods = {
            '.pdf': self._extract_from_pdf,
            '.docx': self._extract_from_docx,
            '.doc': self._extract_from_docx,
            '.xlsx': self._extract_from_excel,
            '.xls': self._extract_from_excel,
            '.txt': self._extract_from_text,
            '.md': self._extract_from_text,
            '.csv': self._extract_from_csv,
            '.json': self._extract_from_json,
            '.xml': self._extract_from_xml,
            '.html': self._extract_from_html,
            '.htm': self._extract_from_html
        }
        
        logger.info(f"文本提取器初始化完成，支持格式: {list(self.SUPPORTED_EXTENSIONS.keys())}")
    
    def extract_from_file(self, file_path: Union[str, Path]) -> ExtractedDocument:
        """
        从文件中提取文本
        
        Args:
            file_path: 文件路径
            
        Returns:
            ExtractedDocument: 提取的文档数据
            
        Raises:
            DocumentProcessingError: 提取失败时抛出
        """
        file_path = Path(file_path)
        
        if not file_path.exists():
            raise ResourceNotFoundError(f"文件不存在: {file_path}")
        
        if not file_path.is_file():
            raise DocumentProcessingError(f"路径不是文件: {file_path}")
        
        file_ext = file_path.suffix.lower()
        
        if file_ext not in self.SUPPORTED_EXTENSIONS:
            raise DocumentProcessingError(f"不支持的文件格式: {file_ext}")
        
        try:
            logger.info(f"开始提取文件: {file_path}")
            
            # 获取文件基本信息
            file_size = file_path.stat().st_size
            
            # 检查文件大小限制
            if self.config.max_file_size > 0 and file_size > self.config.max_file_size:
                raise DocumentProcessingError(
                    f"文件大小 {file_size} 超过限制 {self.config.max_file_size}"
                )
            
            # 根据文件扩展名选择提取方法
            extraction_method = self.extraction_methods.get(file_ext)
            if not extraction_method:
                raise DocumentProcessingError(f"未找到对应的提取方法: {file_ext}")
            
            # 执行文本提取
            content, metadata, page_count = extraction_method(file_path)
            
            # 内容长度检查
            if self.config.max_content_length > 0 and len(content) > self.config.max_content_length:
                logger.warning(f"内容长度 {len(content)} 超过限制，将被截断")
                content = content[:self.config.max_content_length]
                metadata['content_truncated'] = True
            
            # 创建提取结果
            extracted_doc = ExtractedDocument(
                filename=file_path.name,
                file_type=self.SUPPORTED_EXTENSIONS[file_ext],
                content=content,
                metadata=metadata,
                extracted_at=datetime.now(),
                file_size=file_size,
                page_count=page_count,
                extraction_method=extraction_method.__name__
            )
            
            logger.info(f"文件提取完成: {file_path}, 内容长度: {len(content)}")
            return extracted_doc
            
        except Exception as e:
            error_msg = f"文件提取失败 {file_path}: {str(e)}"
            logger.error(error_msg, exc_info=True)
            
            # 创建错误文档
            return ExtractedDocument(
                filename=file_path.name,
                file_type=self.SUPPORTED_EXTENSIONS.get(file_ext, "Unknown"),
                content="",
                metadata={'extraction_error': str(e)},
                extracted_at=datetime.now(),
                file_size=file_path.stat().st_size if file_path.exists() else 0,
                page_count=None,
                extraction_method=None,
                error_info=error_msg
            )
    
    def extract_from_directory(self, directory_path: Union[str, Path]) -> List[ExtractedDocument]:
        """
        从目录中提取所有支持的文档
        
        Args:
            directory_path: 目录路径
            
        Returns:
            List[ExtractedDocument]: 提取的文档列表
        """
        directory_path = Path(directory_path)
        
        if not directory_path.exists():
            raise ResourceNotFoundError(f"目录不存在: {directory_path}")
        
        if not directory_path.is_dir():
            raise DocumentProcessingError(f"路径不是目录: {directory_path}")
        
        extracted_docs = []
        processed_count = 0
        
        # 递归搜索文件
        for file_path in directory_path.rglob("*"):
            if file_path.is_file():
                file_ext = file_path.suffix.lower()
                
                if file_ext in self.SUPPORTED_EXTENSIONS:
                    try:
                        doc = self.extract_from_file(file_path)
                        extracted_docs.append(doc)
                        processed_count += 1
                        
                        # 检查处理数量限制
                        if (self.config.max_documents > 0 and 
                            processed_count >= self.config.max_documents):
                            logger.warning(f"达到文档数量限制 {self.config.max_documents}")
                            break
                            
                    except Exception as e:
                        logger.error(f"处理文件失败 {file_path}: {e}")
                        continue
        
        logger.info(f"目录提取完成: {directory_path}, 处理文件数: {len(extracted_docs)}")
        return extracted_docs
    
    def _extract_from_pdf(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从PDF文件中提取文本"""
        if not PDF_AVAILABLE:
            raise DocumentProcessingError("PDF支持库未安装，请安装: pip install PyPDF2 pdfplumber")
        
        content = ""
        metadata = {}
        page_count = 0
        
        try:
            # 优先使用pdfplumber，回退到PyPDF2
            with pdfplumber.open(file_path) as pdf:
                page_count = len(pdf.pages)
                metadata['page_count'] = page_count
                
                for page in pdf.pages:
                    text = page.extract_text()
                    if text:
                        content += text + "\n"
                
                # 提取表格数据
                tables = []
                for page in pdf.pages:
                    page_tables = page.extract_tables()
                    if page_tables:
                        tables.extend(page_tables)
                
                if tables:
                    metadata['table_count'] = len(tables)
                    # 将表格转换为文本
                    for i, table in enumerate(tables):
                        content += f"\n\n表格 {i+1}:\n"
                        for row in table:
                            if row:
                                content += "\t".join(str(cell) if cell else "" for cell in row) + "\n"
                
        except Exception as e:
            logger.warning(f"pdfplumber提取失败，尝试PyPDF2: {e}")
            
            try:
                with open(file_path, 'rb') as file:
                    pdf_reader = PyPDF2.PdfReader(file)
                    page_count = len(pdf_reader.pages)
                    metadata['page_count'] = page_count
                    
                    for page in pdf_reader.pages:
                        content += page.extract_text() + "\n"
                        
            except Exception as e2:
                raise DocumentProcessingError(f"PDF提取失败: {e2}")
        
        metadata['extraction_library'] = 'pdfplumber' if 'pdfplumber' in str(type(pdf)) else 'PyPDF2'
        return content.strip(), metadata, page_count
    
    def _extract_from_docx(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从Word文档中提取文本"""
        if not DOCX_AVAILABLE:
            raise DocumentProcessingError("Word文档支持库未安装，请安装: pip install python-docx")
        
        try:
            doc = Document(file_path)
            content = ""
            metadata = {}
            
            # 提取段落文本
            paragraph_count = 0
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    content += paragraph.text + "\n"
                    paragraph_count += 1
            
            # 提取表格文本
            table_count = len(doc.tables)
            for table in doc.tables:
                content += "\n\n表格:\n"
                for row in table.rows:
                    row_text = "\t".join(cell.text.strip() for cell in row.cells)
                    content += row_text + "\n"
            
            metadata.update({
                'paragraph_count': paragraph_count,
                'table_count': table_count,
                'extraction_library': 'python-docx'
            })
            
            return content.strip(), metadata, None
            
        except Exception as e:
            raise DocumentProcessingError(f"Word文档提取失败: {e}")
    
    def _extract_from_excel(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从Excel文件中提取文本"""
        if not EXCEL_AVAILABLE:
            raise DocumentProcessingError("Excel支持库未安装，请安装: pip install openpyxl")
        
        try:
            workbook = load_workbook(file_path, data_only=True)
            content = ""
            metadata = {}
            
            sheet_count = len(workbook.sheetnames)
            total_rows = 0
            
            for sheet_name in workbook.sheetnames:
                sheet = workbook[sheet_name]
                content += f"\n\n工作表: {sheet_name}\n"
                
                sheet_rows = 0
                for row in sheet.iter_rows(values_only=True):
                    if any(cell is not None for cell in row):
                        row_text = "\t".join(str(cell) if cell is not None else "" for cell in row)
                        content += row_text + "\n"
                        sheet_rows += 1
                        total_rows += 1
            
            metadata.update({
                'sheet_count': sheet_count,
                'total_rows': total_rows,
                'extraction_library': 'openpyxl'
            })
            
            return content.strip(), metadata, sheet_count
            
        except Exception as e:
            raise DocumentProcessingError(f"Excel文件提取失败: {e}")
    
    def _extract_from_text(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从文本文件中提取内容"""
        try:
            # 尝试多种编码
            encodings = ['utf-8', 'gbk', 'gb2312', 'latin-1']
            content = ""
            used_encoding = ""
            
            for encoding in encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as file:
                        content = file.read()
                        used_encoding = encoding
                        break
                except UnicodeDecodeError:
                    continue
            
            if not content and not used_encoding:
                raise DocumentProcessingError("无法解码文本文件，尝试了多种编码方式")
            
            metadata = {
                'encoding': used_encoding,
                'line_count': len(content.splitlines())
            }
            
            return content, metadata, None
            
        except Exception as e:
            raise DocumentProcessingError(f"文本文件提取失败: {e}")
    
    def _extract_from_csv(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从CSV文件中提取内容"""
        try:
            if PANDAS_AVAILABLE:
                # 使用pandas处理CSV
                df = pd.read_csv(file_path)
                content = df.to_string(index=False)
                metadata = {
                    'rows': len(df),
                    'columns': len(df.columns),
                    'column_names': list(df.columns),
                    'extraction_library': 'pandas'
                }
            else:
                # 使用内置csv模块
                import csv
                content = ""
                with open(file_path, 'r', encoding='utf-8') as file:
                    csv_reader = csv.reader(file)
                    rows = list(csv_reader)
                    
                for row in rows:
                    content += "\t".join(row) + "\n"
                
                metadata = {
                    'rows': len(rows),
                    'columns': len(rows[0]) if rows else 0,
                    'extraction_library': 'csv'
                }
            
            return content, metadata, None
            
        except Exception as e:
            raise DocumentProcessingError(f"CSV文件提取失败: {e}")
    
    def _extract_from_json(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从JSON文件中提取内容"""
        try:
            import json
            
            with open(file_path, 'r', encoding='utf-8') as file:
                data = json.load(file)
            
            # 将JSON转换为格式化的文本
            content = json.dumps(data, ensure_ascii=False, indent=2)
            
            metadata = {
                'json_keys': list(data.keys()) if isinstance(data, dict) else [],
                'json_type': type(data).__name__,
                'extraction_library': 'json'
            }
            
            return content, metadata, None
            
        except Exception as e:
            raise DocumentProcessingError(f"JSON文件提取失败: {e}")
    
    def _extract_from_xml(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从XML文件中提取内容"""
        try:
            import xml.etree.ElementTree as ET
            
            tree = ET.parse(file_path)
            root = tree.getroot()
            
            # 提取所有文本内容
            content = ""
            def extract_text(element, level=0):
                nonlocal content
                indent = "  " * level
                if element.text and element.text.strip():
                    content += f"{indent}{element.tag}: {element.text.strip()}\n"
                else:
                    content += f"{indent}{element.tag}\n"
                
                for child in element:
                    extract_text(child, level + 1)
            
            extract_text(root)
            
            metadata = {
                'root_tag': root.tag,
                'element_count': len(list(root.iter())),
                'extraction_library': 'xml.etree.ElementTree'
            }
            
            return content, metadata, None
            
        except Exception as e:
            raise DocumentProcessingError(f"XML文件提取失败: {e}")
    
    def _extract_from_html(self, file_path: Path) -> tuple[str, Dict[str, Any], Optional[int]]:
        """从HTML文件中提取内容"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                html_content = file.read()
            
            # 简单的HTML文本提取（移除标签）
            import re
            # 移除script和style标签及其内容
            html_content = re.sub(r'<script.*?</script>', '', html_content, flags=re.DOTALL)
            html_content = re.sub(r'<style.*?</style>', '', html_content, flags=re.DOTALL)
            
            # 移除HTML标签
            content = re.sub(r'<[^>]+>', '', html_content)
            
            # 清理多余的空白
            content = re.sub(r'\s+', ' ', content).strip()
            
            metadata = {
                'original_size': len(html_content),
                'extraction_library': 'regex'
            }
            
            return content, metadata, None
            
        except Exception as e:
            raise DocumentProcessingError(f"HTML文件提取失败: {e}")
    
    def get_supported_formats(self) -> Dict[str, str]:
        """获取支持的文件格式"""
        return self.SUPPORTED_EXTENSIONS.copy()
    
    def is_supported_format(self, file_path: Union[str, Path]) -> bool:
        """检查文件格式是否支持"""
        file_path = Path(file_path)
        return file_path.suffix.lower() in self.SUPPORTED_EXTENSIONS
    
    def get_extraction_stats(self) -> Dict[str, Any]:
        """获取提取器统计信息"""
        available_libraries = []
        
        if PDF_AVAILABLE:
            available_libraries.append("PDF (PyPDF2, pdfplumber)")
        if DOCX_AVAILABLE:
            available_libraries.append("Word (python-docx)")
        if EXCEL_AVAILABLE:
            available_libraries.append("Excel (openpyxl)")
        if PANDAS_AVAILABLE:
            available_libraries.append("CSV (pandas)")
        
        return {
            'supported_formats': list(self.SUPPORTED_EXTENSIONS.keys()),
            'available_libraries': available_libraries,
            'max_file_size': self.config.max_file_size,
            'max_content_length': self.config.max_content_length,
            'max_documents': self.config.max_documents
        } 