from mcp.server.fastmcp import FastMCP
from typing import Dict, List, Any, Optional
import logging
import json
from pathlib import Path
import os

logger = logging.getLogger(__name__)

# Create FastMCP server instance
mcp = FastMCP("PDF Parser")

# Global document storage
documents: Dict[str, Dict] = {}

# Storage file path - use absolute path to ensure it works from any directory
STORAGE_FILE = Path(__file__).parent / "storage.json"

def load_documents():
    """Load documents from storage.json file"""
    global documents
    try:
        if STORAGE_FILE.exists():
            with open(STORAGE_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
                documents.clear()
                documents.update(data)
            logger.info(f"Loaded {len(documents)} documents from storage at {STORAGE_FILE}")
        else:
            logger.info(f"No storage file found at {STORAGE_FILE}, starting with empty documents")
            documents.clear()
    except Exception as e:
        logger.error(f"Error loading documents from storage: {str(e)}")
        documents.clear()

def reload_documents():
    """Reload documents from storage - called by tools to ensure fresh data"""
    load_documents()

# Load documents on startup
load_documents()

@mcp.tool()
def list_documents() -> str:
    """List all processed PDF documents"""
    # Always reload documents to get latest data
    reload_documents()
    
    if not documents:
        return "No documents available."
    
    doc_list = []
    for file_id, doc in documents.items():
        doc_list.append(f"- {file_id}: {doc['filename']} (Status: {doc['status']})")
    
    return "Available documents:\n" + "\n".join(doc_list)

@mcp.tool()
def get_document_summary(file_id: str) -> str:
    """Get AI-generated summary of a specific document"""
    # Always reload documents to get latest data
    reload_documents()
    
    if file_id not in documents:
        return f"Document {file_id} not found."
    
    doc = documents[file_id]
    
    if "summary" not in doc:
        return f"No summary available for document {file_id}."
    
    summary = doc["summary"]
    
    if summary["type"] == "single_chunk":
        return f"Document: {doc['filename']}\n\nSummary:\n{summary['summary']}"
    else:
        # Multi-chunk summary
        result = f"Document: {doc['filename']}\n\nOverall Summary:\n{summary['overall_summary']}\n\n"
        result += "Chunk Summaries:\n"
        for i, chunk_summary in enumerate(summary['chunk_summaries']):
            result += f"\nChunk {i+1}:\n{chunk_summary}\n"
        return result

@mcp.tool()
def get_document_content(file_id: str, chunk_id: Optional[int] = None) -> str:
    """Get full content of a specific document or a specific chunk"""
    # Always reload documents to get latest data
    reload_documents()
    
    if file_id not in documents:
        return f"Document {file_id} not found."
    
    doc = documents[file_id]
    
    if chunk_id is not None:
        # Get specific chunk
        if "chunks" in doc and chunk_id < len(doc["chunks"]):
            chunk = doc["chunks"][chunk_id]
            return f"Document: {doc['filename']}\nChunk {chunk_id} (Pages {chunk['page_range']}):\n\n{chunk['text']}"
        else:
            return f"Chunk {chunk_id} not found in document {file_id}."
    else:
        # Get full content
        if "extracted_text" in doc:
            return f"Document: {doc['filename']}\n\nFull Content:\n{doc['extracted_text']}"
        else:
            return f"Content not available for document {file_id}."

@mcp.tool()
def search_documents(query: str, file_id: Optional[str] = None) -> str:
    """Search for specific content across all documents or within a specific document"""
    # Always reload documents to get latest data
    reload_documents()
    
    results = []
    
    documents_to_search = {file_id: documents[file_id]} if file_id else documents
    
    for doc_id, doc in documents_to_search.items():
        if "extracted_text" not in doc:
            continue
        
        text = doc["extracted_text"].lower()
        query_lower = query.lower()
        
        if query_lower in text:
            # Find context around matches
            lines = doc["extracted_text"].split('\n')
            matching_lines = []
            
            for i, line in enumerate(lines):
                if query_lower in line.lower():
                    # Include context (previous and next lines)
                    start = max(0, i-2)
                    end = min(len(lines), i+3)
                    context = lines[start:end]
                    matching_lines.append(f"Line {i+1}: " + "\n".join(context))
            
            if matching_lines:
                results.append(f"Document: {doc['filename']} (ID: {doc_id})")
                results.extend(matching_lines[:5])  # Limit results
                results.append("")
    
    if not results:
        return f"No results found for query: {query}"
    
    return f"Search results for '{query}':\n\n" + "\n".join(results)

@mcp.tool()
def get_document_metadata(file_id: str) -> str:
    """Get metadata about a document (pages, size, processing time, etc.)"""
    # Always reload documents to get latest data
    reload_documents()
    
    if file_id not in documents:
        return f"Document {file_id} not found."
    
    doc = documents[file_id]
    
    metadata = {
        "File ID": file_id,
        "Filename": doc.get("filename", "N/A"),
        "Status": doc.get("status", "N/A"),
        "Processed At": doc.get("processed_at", "N/A"),
        "Chunks": len(doc.get("chunks", [])),
        "Total Tokens": sum(chunk.get("tokens", 0) for chunk in doc.get("chunks", [])),
        "Has Summary": "summary" in doc
    }
    
    if "summary" in doc:
        metadata["Summary Type"] = doc["summary"].get("type", "N/A")
    
    text = f"Metadata for {doc.get('filename', file_id)}:\n\n"
    text += "\n".join([f"{key}: {value}" for key, value in metadata.items()])
    
    return text

@mcp.tool()
def answer_question(question: str, file_id: str) -> str:
    """Answer a question about a specific document using its content and summary"""
    # Always reload documents to get latest data
    reload_documents()
    
    if file_id not in documents:
        return f"Document {file_id} not found."
    
    doc = documents[file_id]
    
    # Get document summary for context
    summary_text = "No summary available"
    if "summary" in doc:
        summary = doc["summary"]
        if summary["type"] == "single_chunk":
            summary_text = summary["summary"]
        else:
            summary_text = summary["overall_summary"]
    
    # Get document content (truncated for context)
    content = doc.get("extracted_text", "")
    content_preview = content[:5000] + "..." if len(content) > 5000 else content
    
    # Return context for the question
    return f"""Question: {question}

Document: {doc['filename']}

Summary: {summary_text}

Content Preview: {content_preview}

To get a detailed AI-powered answer, please use the Anthropic API integration in the main application."""

# Helper functions for document management
def update_document_data(file_id: str, document_data: Dict[str, Any]):
    """Update document data in MCP server"""
    documents[file_id] = document_data
    logger.info(f"Updated document data for {file_id}")

def remove_document(file_id: str):
    """Remove document from MCP server"""
    if file_id in documents:
        del documents[file_id]
        logger.info(f"Removed document {file_id} from MCP server")

def get_server():
    """Get the FastMCP server instance"""
    return mcp