"""
Document processor for handling various file types and extracting content.

This module provides functionality to process PDFs, Word documents, Excel files,
and other document types to extract text content for training and inference.
"""

import json
import csv
from pathlib import Path
from typing import Dict, List, Optional, Any, Union
import pandas as pd

from ..core.config import ToxoConfig
from ..utils.logger import get_logger
from ..utils.exceptions import ValidationError


class DocumentProcessor:
    """
    Processor for various document types.
    
    Handles extraction of text content from PDFs, Word docs, Excel files,
    and other formats commonly used in training data.
    """
    
    def __init__(self, config: ToxoConfig):
        """
        Initialize document processor.
        
        Args:
            config: Toxo configuration
        """
        self.config = config
        self.logger = get_logger("toxo.document_processor")
        
        # Supported file types
        self.supported_extensions = {
            '.txt': self._process_text,
            '.md': self._process_text,
            '.json': self._process_json,
            '.jsonl': self._process_jsonl,
            '.csv': self._process_csv,
            '.xlsx': self._process_excel,
            '.xls': self._process_excel,
        }
        
        # Try to import optional dependencies
        self.pdf_available = self._check_pdf_support()
        self.docx_available = self._check_docx_support()
        
        if self.pdf_available:
            self.supported_extensions['.pdf'] = self._process_pdf
        
        if self.docx_available:
            self.supported_extensions['.docx'] = self._process_docx
        
        self.logger.info(f"DocumentProcessor initialized, supports: {list(self.supported_extensions.keys())}")
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state of the processor for serialization."""
        return {
            "supported_extensions": list(self.supported_extensions.keys()),
            "pdf_available": self.pdf_available,
            "docx_available": self.docx_available,
            "config_available": self.config is not None
        }
    
    def _check_pdf_support(self) -> bool:
        """Check if PDF processing is available."""
        try:
            import pypdf
            return True
        except ImportError:
            self.logger.warning("pypdf not available, PDF processing disabled")
            return False
    
    def _check_docx_support(self) -> bool:
        """Check if DOCX processing is available."""
        try:
            import docx
            return True
        except ImportError:
            self.logger.warning("python-docx not available, DOCX processing disabled")
            return False
    
    def process_file(self, file_path: Union[str, Path]) -> Dict[str, Any]:
        """
        Process a single file and extract content.
        
        Args:
            file_path: Path to the file
            
        Returns:
            Dictionary with extracted content and metadata
        """
        file_path = Path(file_path)
        
        if not file_path.exists():
            raise ValidationError(f"File not found: {file_path}")
        
        extension = file_path.suffix.lower()
        
        if extension not in self.supported_extensions:
            raise ValidationError(f"Unsupported file type: {extension}")
        
        try:
            processor = self.supported_extensions[extension]
            content = processor(file_path)
            
            result = {
                "file_path": str(file_path),
                "file_name": file_path.name,
                "extension": extension,
                "content": content,
                "metadata": {
                    "file_size": file_path.stat().st_size,
                    "processed_at": pd.Timestamp.now().isoformat()
                }
            }
            
            self.logger.debug(f"Processed {file_path.name}: {len(str(content))} characters")
            return result
            
        except Exception as e:
            self.logger.error(f"Failed to process {file_path}: {str(e)}")
            raise ValidationError(f"Processing failed: {str(e)}")
    
    def process_directory(
        self,
        directory_path: Union[str, Path],
        recursive: bool = True,
        pattern: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Process all supported files in a directory.
        
        Args:
            directory_path: Path to directory
            recursive: Whether to search subdirectories
            pattern: Optional glob pattern to filter files
            
        Returns:
            List of processed file results
        """
        directory_path = Path(directory_path)
        
        if not directory_path.exists():
            raise ValidationError(f"Directory not found: {directory_path}")
        
        if not directory_path.is_dir():
            raise ValidationError(f"Path is not a directory: {directory_path}")
        
        # Find files to process
        if recursive:
            search_pattern = "**/*" if not pattern else f"**/{pattern}"
        else:
            search_pattern = "*" if not pattern else pattern
        
        files_to_process = []
        for file_path in directory_path.glob(search_pattern):
            if file_path.is_file() and file_path.suffix.lower() in self.supported_extensions:
                files_to_process.append(file_path)
        
        # Process files
        results = []
        for file_path in files_to_process:
            try:
                result = self.process_file(file_path)
                results.append(result)
            except Exception as e:
                self.logger.warning(f"Skipping {file_path}: {str(e)}")
        
        self.logger.info(f"Processed {len(results)} files from {directory_path}")
        return results
    
    def _process_text(self, file_path: Path) -> str:
        """Process plain text files."""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            # Try with different encoding
            with open(file_path, 'r', encoding='latin-1') as f:
                return f.read()
    
    def _process_json(self, file_path: Path) -> Dict[str, Any]:
        """Process JSON files."""
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def _process_jsonl(self, file_path: Path) -> List[Dict[str, Any]]:
        """Process JSONL files."""
        data = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    data.append(json.loads(line))
        return data
    
    def _process_csv(self, file_path: Path) -> List[Dict[str, Any]]:
        """Process CSV files."""
        try:
            df = pd.read_csv(file_path)
            return df.to_dict('records')
        except Exception as e:
            # Fallback to basic CSV reader
            data = []
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    data.append(dict(row))
            return data
    
    def _process_excel(self, file_path: Path) -> Dict[str, List[Dict[str, Any]]]:
        """Process Excel files."""
        try:
            # Read all sheets
            xl_file = pd.ExcelFile(file_path)
            result = {}
            
            for sheet_name in xl_file.sheet_names:
                df = pd.read_excel(file_path, sheet_name=sheet_name)
                result[sheet_name] = df.to_dict('records')
            
            return result
            
        except Exception as e:
            self.logger.error(f"Excel processing failed: {str(e)}")
            raise ValidationError(f"Excel processing failed: {str(e)}")
    
    def _process_pdf(self, file_path: Path) -> str:
        """Process PDF files."""
        if not self.pdf_available:
            raise ValidationError("PDF processing not available - install pypdf")
        
        import pypdf
        
        text = ""
        try:
            with open(file_path, 'rb') as f:
                pdf_reader = pypdf.PdfReader(f)
                
                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    text += page.extract_text() + "\n"
            
            return text.strip()
            
        except Exception as e:
            self.logger.error(f"PDF processing failed: {str(e)}")
            raise ValidationError(f"PDF processing failed: {str(e)}")
    
    def _process_docx(self, file_path: Path) -> str:
        """Process DOCX files."""
        if not self.docx_available:
            raise ValidationError("DOCX processing not available - install python-docx")
        
        import docx
        
        try:
            doc = docx.Document(file_path)
            text = ""
            
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            
            return text.strip()
            
        except Exception as e:
            self.logger.error(f"DOCX processing failed: {str(e)}")
            raise ValidationError(f"DOCX processing failed: {str(e)}")
    
    def extract_training_data(
        self,
        processed_files: List[Dict[str, Any]],
        input_field: str = "input",
        output_field: str = "output"
    ) -> List[Dict[str, str]]:
        """
        Extract training data from processed files.
        
        Args:
            processed_files: List of processed file results
            input_field: Field name for input text
            output_field: Field name for output text
            
        Returns:
            List of training examples
        """
        training_data = []
        
        for file_result in processed_files:
            content = file_result["content"]
            file_name = file_result["file_name"]
            
            # Handle different content types
            if isinstance(content, str):
                # For text files, create a single example
                training_data.append({
                    input_field: f"Content from {file_name}",
                    output_field: content[:1000],  # Truncate long content
                    "source": file_name
                })
                
            elif isinstance(content, list):
                # For JSONL, CSV data
                for item in content:
                    if isinstance(item, dict):
                        # Try to find input/output fields
                        input_text = self._extract_field(item, input_field)
                        output_text = self._extract_field(item, output_field)
                        
                        if input_text and output_text:
                            training_data.append({
                                input_field: input_text,
                                output_field: output_text,
                                "source": file_name
                            })
                            
            elif isinstance(content, dict):
                # For JSON or Excel with multiple sheets
                for key, value in content.items():
                    if isinstance(value, list):
                        for item in value:
                            if isinstance(item, dict):
                                input_text = self._extract_field(item, input_field)
                                output_text = self._extract_field(item, output_field)
                                
                                if input_text and output_text:
                                    training_data.append({
                                        input_field: input_text,
                                        output_field: output_text,
                                        "source": f"{file_name}:{key}"
                                    })
        
        self.logger.info(f"Extracted {len(training_data)} training examples")
        return training_data
    
    def _extract_field(self, item: Dict[str, Any], field_name: str) -> Optional[str]:
        """Extract field value from item, trying various field name variations."""
        # Try exact match first
        if field_name in item:
            return str(item[field_name])
        
        # Try case-insensitive match
        field_lower = field_name.lower()
        for key, value in item.items():
            if key.lower() == field_lower:
                return str(value)
        
        # Try common variations
        variations = {
            "input": ["prompt", "question", "query", "text"],
            "output": ["response", "answer", "target", "label", "result"]
        }
        
        if field_name in variations:
            for variant in variations[field_name]:
                if variant in item:
                    return str(item[variant])
                
                # Case-insensitive variant match
                for key, value in item.items():
                    if key.lower() == variant.lower():
                        return str(value)
        
        return None
    
    def get_supported_extensions(self) -> List[str]:
        """Get list of supported file extensions."""
        return list(self.supported_extensions.keys())
    
    def is_supported(self, file_path: Union[str, Path]) -> bool:
        """Check if a file type is supported."""
        extension = Path(file_path).suffix.lower()
        return extension in self.supported_extensions 