import asyncio
import os
import uuid
from typing import List, Dict, Any, Optional
from mcp.server.fastmcp import FastMCP
from rag_store import RAGDocumentStore
from document_processors import ContentExtractor
from config import Config
import logging

import sys
import signal
import traceback

# Setup signal handlers for graceful shutdown
def signal_handler(signum, frame):
    logger.info(f"Received signal {signum}, shutting down gracefully...")
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize MCP server
mcp = FastMCP("RAGLargeFileProcessor")

# Global RAG store (lazy initialized)
rag_store = None

def get_rag_store():
    """Get or create RAG store"""
    global rag_store
    if rag_store is None:
        from rag_store import RAGDocumentStore
        rag_store = RAGDocumentStore()
    return rag_store
# @mcp.tool()
# async def upload_documents_with_rag(file_paths: List[str]) -> Dict[str, Any]:
#     """
#     Upload documents and create vector embeddings for RAG with large file support.
    
#     Args:
#         file_paths: List of file paths to process
    
#     Returns:
#         Upload results with RAG processing status
#     """
#     results = {
#         "processed_documents": [],
#         "failed_documents": [],
#         "rag_summary": {
#             "total_chunks_created": 0,
#             "total_documents": 0,
#             "processing_time_ms": 0
#         }
#     }
    
#     import time
#     start_time = time.time()
    
#     for file_path in file_paths:
#         try:
#             logger.info(f"Processing file: {file_path}")
            
#             if not os.path.exists(file_path):
#                 results["failed_documents"].append({
#                     "file_path": file_path,
#                     "error": "File not found"
#                 })
#                 continue
            
#             # Check file size
#             file_info = await ContentExtractor.get_file_info(file_path)
#             if file_info.get("file_size_mb", 0) > Config.MAX_FILE_SIZE_MB:
#                 results["failed_documents"].append({
#                     "file_path": file_path,
#                     "error": f"File too large. Maximum size: {Config.MAX_FILE_SIZE_MB}MB"
#                 })
#                 continue
            
#             # Generate document ID
#             doc_id = str(uuid.uuid4())
#             filename = os.path.basename(file_path)
            
#             # Extract content based on file type
#             try:
#                 content, file_type = await ContentExtractor.extract_content(file_path)
#             except ValueError as e:
#                 results["failed_documents"].append({
#                     "file_path": file_path,
#                     "error": str(e)
#                 })
#                 continue
            
#             if not content or content.startswith("Error"):
#                 results["failed_documents"].append({
#                     "file_path": file_path,
#                     "error": "Failed to extract content or content is empty"
#                 })
#                 continue
            
#             # Prepare metadata
#             metadata = {
#                 "file_size_bytes": file_info["file_size_bytes"],
#                 "file_size_mb": file_info["file_size_mb"],
#                 "file_extension": os.path.splitext(file_path)[1].lower(),
#                 "processing_strategy": file_info["processing_strategy"]
#             }
            
#             # Add to RAG store
#             rag_result = await rag_store.add_document(doc_id, filename, file_type, content, metadata)
            
#             if "error" in rag_result:
#                 results["failed_documents"].append({
#                     "file_path": file_path,
#                     "error": rag_result["error"]
#                 })
#                 continue
            
#             results["processed_documents"].append({
#                 "doc_id": doc_id,
#                 "filename": filename,
#                 "file_type": file_type,
#                 "chunks_created": rag_result["chunks_created"],
#                 "content_length": len(content),
#                 "file_size_mb": file_info["file_size_mb"],
#                 "processing_strategy": file_info["processing_strategy"],
#                 "status": "success"
#             })
            
#             results["rag_summary"]["total_chunks_created"] += rag_result["chunks_created"]
#             results["rag_summary"]["total_documents"] += 1
            
#             logger.info(f"Successfully processed {filename}: {rag_result['chunks_created']} chunks")
            
#         except Exception as e:
#             logger.error(f"Error processing {file_path}: {e}")
#             results["failed_documents"].append({
#                 "file_path": file_path,
#                 "error": str(e)
#             })
    
#     results["rag_summary"]["processing_time_ms"] = int((time.time() - start_time) * 1000)
    
#     return results

@mcp.tool()
async def upload_documents_with_rag(file_paths: List[str]) -> Dict[str, Any]:
    """Upload documents with comprehensive error handling"""
    try:
        logger.info(f"Starting document upload for {len(file_paths)} files")
        
        results = {
            "processed_documents": [],
            "failed_documents": [],
            "rag_summary": {
                "total_chunks_created": 0,
                "total_documents": 0,
                "processing_time_ms": 0
            }
        }
        
        import time
        start_time = time.time()
        
        for file_path in file_paths:
            try:
                logger.info(f"Processing file: {file_path}")
                
                if not os.path.exists(file_path):
                    error_msg = f"File not found: {file_path}"
                    logger.error(error_msg)
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": error_msg
                    })
                    continue
                
                # Check file size
                try:
                    file_info = await ContentExtractor.get_file_info(file_path)
                    if "error" in file_info:
                        results["failed_documents"].append({
                            "file_path": file_path,
                            "error": file_info["error"]
                        })
                        continue
                        
                except Exception as e:
                    error_msg = f"Failed to get file info: {str(e)}"
                    logger.error(error_msg)
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": error_msg
                    })
                    continue
                
                if file_info.get("file_size_mb", 0) > Config.MAX_FILE_SIZE_MB:
                    error_msg = f"File too large. Maximum size: {Config.MAX_FILE_SIZE_MB}MB"
                    logger.error(error_msg)
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": error_msg
                    })
                    continue
                
                # Generate document ID
                doc_id = str(uuid.uuid4())
                filename = os.path.basename(file_path)
                
                # Extract content based on file type
                try:
                    content, file_type = await ContentExtractor.extract_content(file_path)
                except Exception as e:
                    error_msg = f"Content extraction failed: {str(e)}"
                    logger.error(error_msg)
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": error_msg
                    })
                    continue
                
                if not content or content.startswith("Error"):
                    error_msg = "Failed to extract content or content is empty"
                    logger.error(error_msg)
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": error_msg
                    })
                    continue
                
                # Prepare metadata
                metadata = {
                    "file_size_bytes": file_info["file_size_bytes"],
                    "file_size_mb": file_info["file_size_mb"],
                    "file_extension": os.path.splitext(file_path)[1].lower(),
                    "processing_strategy": file_info["processing_strategy"]
                }
                
                # Add to RAG store
                try:
                    rag_result = await rag_store.add_document(doc_id, filename, file_type, content, metadata)
                except Exception as e:
                    error_msg = f"RAG store addition failed: {str(e)}"
                    logger.error(error_msg)
                    logger.error(traceback.format_exc())
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": error_msg
                    })
                    continue
                
                if "error" in rag_result:
                    results["failed_documents"].append({
                        "file_path": file_path,
                        "error": rag_result["error"]
                    })
                    continue
                
                results["processed_documents"].append({
                    "doc_id": doc_id,
                    "filename": filename,
                    "file_type": file_type,
                    "chunks_created": rag_result["chunks_created"],
                    "content_length": len(content),
                    "file_size_mb": file_info["file_size_mb"],
                    "processing_strategy": file_info["processing_strategy"],
                    "status": "success"
                })
                
                results["rag_summary"]["total_chunks_created"] += rag_result["chunks_created"]
                results["rag_summary"]["total_documents"] += 1
                
                logger.info(f"Successfully processed {filename}: {rag_result['chunks_created']} chunks")
                
            except Exception as e:
                error_msg = f"Unexpected error processing {file_path}: {str(e)}"
                logger.error(error_msg)
                logger.error(traceback.format_exc())
                results["failed_documents"].append({
                    "file_path": file_path,
                    "error": error_msg
                })
        
        results["rag_summary"]["processing_time_ms"] = int((time.time() - start_time) * 1000)
        
        logger.info(f"Upload completed: {results['rag_summary']['total_documents']} successful, {len(results['failed_documents'])} failed")
        return results
        
    except Exception as e:
        error_msg = f"Critical error in upload_documents_with_rag: {str(e)}"
        logger.error(error_msg)
        logger.error(traceback.format_exc())
        return {"error": error_msg}

@mcp.tool()
async def semantic_search_documents(query: str, n_results: int = 10, 
                                  document_ids: Optional[List[str]] = None) -> Dict[str, Any]:
    """
    Perform semantic search across uploaded documents using RAG.
    
    Args:
        query: Search query
        n_results: Number of results to return
        document_ids: Optional list of specific document IDs to search
    
    Returns:
        Semantically similar content chunks
    """
    try:
        results = await rag_store.semantic_search(query, n_results, document_ids)
        
        return {
            "query": query,
            "total_results": len(results),
            "results": results,
            "search_type": "semantic_similarity"
        }
    except Exception as e:
        logger.error(f"Semantic search failed: {e}")
        return {"error": f"Semantic search failed: {str(e)}"}

@mcp.tool()
async def retrieve_context_for_query(query: str, max_context_length: int = 4000) -> Dict[str, Any]:
    """
    Retrieve relevant context for a query using RAG (optimized for LLM consumption).
    
    Args:
        query: Query to find relevant context for
        max_context_length: Maximum length of combined context
    
    Returns:
        Retrieved context optimized for LLM consumption
    """
    try:
        # Get semantic search results
        search_results = await rag_store.semantic_search(query, n_results=8)
        
        # Combine most relevant chunks
        context_parts = []
        current_length = 0
        sources_used = set()
        
        for result in search_results:
            content = result["content"]
            metadata = result["metadata"]
            similarity = result["similarity_score"]
            
            # Skip if similarity is too low
            if similarity < 0.3:
                continue
            
            # Add source information
            source_info = f"\n[Source: {metadata['filename']} ({metadata['file_type']}) - Similarity: {similarity:.2f}]\n"
            full_content = source_info + content.strip()
            
            if current_length + len(full_content) <= max_context_length:
                context_parts.append(full_content)
                current_length += len(full_content)
                sources_used.add(metadata['filename'])
            else:
                break
        
        combined_context = "\n\n".join(context_parts)
        
        return {
            "query": query,
            "retrieved_context": combined_context,
            "context_length": len(combined_context),
            "sources_used": list(sources_used),
            "chunks_retrieved": len(context_parts),
            "retrieval_method": "semantic_similarity"
        }
    except Exception as e:
        logger.error(f"Context retrieval failed: {e}")
        return {"error": f"Context retrieval failed: {str(e)}"}

@mcp.tool()
async def get_rag_document_summary() -> Dict[str, Any]:
    """
    Get summary of RAG-enabled document store.
    
    Returns:
        Summary of documents and vector database status
    """
    return await rag_store.get_document_summary()

@mcp.tool()
async def analyze_document_collection() -> Dict[str, Any]:
    """
    Analyze the collection of uploaded documents.
    
    Returns:
        Analysis of document collection
    """
    try:
        summary = await rag_store.get_document_summary()
        
        if summary.get("total_documents", 0) == 0:
            return {"error": "No documents in collection"}
        
        # Analyze by file type
        file_types = {}
        total_size_mb = 0
        large_files = 0
        
        for doc_id, doc_info in summary.get("documents", {}).items():
            file_type = doc_info["file_type"]
            file_types[file_type] = file_types.get(file_type, 0) + 1
            
            doc_size = doc_info.get("metadata", {}).get("file_size_mb", 0)
            total_size_mb += doc_size
            
            if doc_size > Config.MEMORY_THRESHOLD_MB:
                large_files += 1
        
        return {
            "collection_analysis": {
                "total_documents": summary["total_documents"],
                "total_chunks": summary["total_chunks"],
                "file_type_distribution": file_types,
                "total_size_mb": round(total_size_mb, 2),
                "large_files_count": large_files,
                "average_chunks_per_doc": round(summary["total_chunks"] / summary["total_documents"], 1) if summary["total_documents"] > 0 else 0
            },
            "embedding_info": {
                "model": summary.get("embedding_model", "unknown"),
                "vector_db_status": summary.get("vector_db_status", "unknown")
            }
        }
    except Exception as e:
        logger.error(f"Document collection analysis failed: {e}")
        return {"error": str(e)}
    
@mcp.tool()
async def process_images_with_ocr(image_paths: List[str], enhancement_level: str = "standard") -> Dict[str, Any]:
    """
    Process multiple images with OCR using specified enhancement level.
    
    Args:
        image_paths: List of image file paths
        enhancement_level: OCR enhancement level (light, standard, aggressive)
    
    Returns:
        OCR processing results
    """
    try:
        logger.info(f"Processing {len(image_paths)} images with OCR (level: {enhancement_level})")
        
        results = {
            "processed_images": [],
            "failed_images": [],
            "batch_summary": {
                "total_images": len(image_paths),
                "successful_extractions": 0,
                "total_characters": 0,
                "processing_time_ms": 0
            }
        }
        
        import time
        start_time = time.time()
        
        for image_path in image_paths:
            try:
                if not os.path.exists(image_path):
                    results["failed_images"].append({
                        "image_path": image_path,
                        "error": "File not found"
                    })
                    continue
                
                # Check file size
                file_info = await ContentExtractor.get_file_info(image_path)
                if file_info.get("file_size_mb", 0) > Config.MAX_FILE_SIZE_MB:
                    results["failed_images"].append({
                        "image_path": image_path,
                        "error": f"File too large. Maximum size: {Config.MAX_FILE_SIZE_MB}MB"
                    })
                    continue
                
                # Extract content from image
                content = await ContentExtractor.extract_image_content(image_path)
                
                if content and not content.startswith("Error") and not content.startswith("No text"):
                    results["processed_images"].append({
                        "image_path": image_path,
                        "filename": os.path.basename(image_path),
                        "content_length": len(content),
                        "file_size_mb": file_info["file_size_mb"],
                        "status": "success",
                        "content_preview": content[:200] + "..." if len(content) > 200 else content
                    })
                    
                    results["batch_summary"]["successful_extractions"] += 1
                    results["batch_summary"]["total_characters"] += len(content)
                else:
                    results["failed_images"].append({
                        "image_path": image_path,
                        "error": content or "No text extracted"
                    })
            
            except Exception as e:
                logger.error(f"Error processing image {image_path}: {e}")
                results["failed_images"].append({
                    "image_path": image_path,
                    "error": str(e)
                })
        
        results["batch_summary"]["processing_time_ms"] = int((time.time() - start_time) * 1000)
        
        logger.info(f"Image processing complete: {results['batch_summary']['successful_extractions']}/{len(image_paths)} successful")
        return results
        
    except Exception as e:
        logger.error(f"Image batch processing failed: {e}")
        return {"error": str(e)}

@mcp.tool()
async def clear_document_store() -> Dict[str, Any]:
    """
    Clear all documents from the RAG store.
    
    Returns:
        Confirmation of clearing operation
    """
    try:
        rag_store.clear_all_documents()
        return {
            "status": "success",
            "message": "All documents cleared from RAG store",
            "timestamp": str(asyncio.get_event_loop().time())
        }
    except Exception as e:
        logger.error(f"Failed to clear document store: {e}")
        return {"error": str(e)}

if __name__ == "__main__":
    logger.info("Starting RAG Large File MCP Server...")
    mcp.run(transport="stdio")

def main():
    """Main entry point for the MCP server"""
    logger.info("Starting RAG Large File MCP Server...")
    mcp.run(transport="stdio")

if __name__ == "__main__":
    main()
