from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import os
import uuid
from pathlib import Path
from typing import Dict, List, Optional
import logging
from datetime import datetime
import json

from pdf_processor import PDFProcessor
from mcp_server import update_document_data, remove_document
from anthropic_client import AnthropicClient

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

app = FastAPI(title="PDF Parser MCP Server", version="1.0.0")

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize components
pdf_processor = PDFProcessor()
anthropic_client = AnthropicClient()

# Storage for processed documents
processed_documents: Dict[str, Dict] = {}

# Ensure upload directory exists
UPLOAD_DIR = Path(os.getenv("UPLOAD_DIR", "uploads"))
UPLOAD_DIR.mkdir(exist_ok=True)

STORAGE_FILE = Path("storage.json")

def save_documents():
    with open(STORAGE_FILE, "w", encoding="utf-8") as f:
        json.dump(processed_documents, f, ensure_ascii=False, indent=2)

def load_documents():
    global processed_documents
    if STORAGE_FILE.exists():
        with open(STORAGE_FILE, "r", encoding="utf-8") as f:
            processed_documents.clear()
            processed_documents.update(json.load(f))

# Load documents on startup
load_documents()

@app.get("/")
async def root():
    """Welcome message and API information"""
    return {
        "message": "PDF Parser MCP Server API",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/health",
        "endpoints": {
            "upload": "POST /upload-pdf/",
            "list_documents": "GET /documents/",
            "document_info": "GET /documents/{file_id}",
            "document_content": "GET /documents/{file_id}/content",
            "document_summary": "GET /documents/{file_id}/summary", 
            "document_chunks": "GET /documents/{file_id}/chunks",
            "full_document": "GET /documents/{file_id}/full"
        }
    }

@app.post("/upload-pdf/")
async def upload_pdf(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...)
):
    """Upload and process a PDF file"""
    
    # Validate file type
    if not file.filename.endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed")
    
    # Generate unique file ID
    file_id = str(uuid.uuid4())
    file_path = UPLOAD_DIR / f"{file_id}.pdf"
    
    try:
        # Save uploaded file
        with open(file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)
        
        # Process PDF in background
        background_tasks.add_task(process_pdf_background, file_id, file_path, file.filename)
        
        return JSONResponse({
            "file_id": file_id,
            "filename": file.filename,
            "status": "processing",
            "message": "PDF uploaded successfully. Processing in background."
        })
        
    except Exception as e:
        logger.error(f"Error uploading file: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error uploading file: {str(e)}")

async def process_pdf_background(file_id: str, file_path: Path, filename: str):
    """Background task to process PDF"""
    try:
        logger.info(f"Processing PDF: {filename}")
        
        # Extract text from PDF
        extracted_text = pdf_processor.extract_text(file_path)
        
        # Chunk text if needed
        chunks = pdf_processor.chunk_text(extracted_text)
        
        # Generate summary using Anthropic
        summary = await anthropic_client.generate_summary(chunks)
        
        # Store processed document
        processed_documents[file_id] = {
            "filename": filename,
            "extracted_text": extracted_text,
            "chunks": chunks,
            "summary": summary,
            "processed_at": datetime.now().isoformat(),
            "status": "completed"
        }
        
        # Update MCP server data
        update_document_data(file_id, processed_documents[file_id])
        
        logger.info(f"Successfully processed PDF: {filename}")
        
    except Exception as e:
        logger.error(f"Error processing PDF {filename}: {str(e)}")
        processed_documents[file_id] = {
            "filename": filename,
            "status": "error",
            "error": str(e),
            "processed_at": datetime.now().isoformat()
        }
    finally:
        save_documents()

@app.get("/status/{file_id}")
async def get_processing_status(file_id: str):
    """Get processing status of a PDF"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="File not found")
    
    doc = processed_documents[file_id]
    return {
        "file_id": file_id,
        "filename": doc.get("filename"),
        "status": doc.get("status"),
        "processed_at": doc.get("processed_at"),
        "has_summary": "summary" in doc,
        "error": doc.get("error")
    }

@app.get("/documents/")
async def list_documents():
    """List all processed documents"""
    return {
        "documents": [
            {
                "file_id": file_id,
                "filename": doc.get("filename"),
                "status": doc.get("status"),
                "processed_at": doc.get("processed_at"),
                "has_text": "extracted_text" in doc,
                "has_summary": "summary" in doc,
                "has_chunks": "chunks" in doc
            }
            for file_id, doc in processed_documents.items()
        ]
    }

@app.get("/documents/{file_id}")
async def get_document_info(file_id: str):
    """Get document metadata and info"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="Document not found")
    
    doc = processed_documents[file_id]
    return {
        "file_id": file_id,
        "filename": doc.get("filename"),
        "status": doc.get("status"),
        "processed_at": doc.get("processed_at"),
        "has_text": "extracted_text" in doc,
        "has_summary": "summary" in doc,
        "has_chunks": "chunks" in doc,
        "text_length": len(doc.get("extracted_text", "")),
        "num_chunks": len(doc.get("chunks", [])),
        "error": doc.get("error")
    }

@app.get("/documents/{file_id}/content")
async def get_document_content(file_id: str):
    """Get extracted text content of a document"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="Document not found")
    
    doc = processed_documents[file_id]
    if "extracted_text" not in doc:
        raise HTTPException(status_code=404, detail="Extracted text not available")
    
    return {
        "file_id": file_id,
        "filename": doc.get("filename"),
        "extracted_text": doc["extracted_text"],
        "text_length": len(doc["extracted_text"])
    }

@app.get("/documents/{file_id}/summary")
async def get_document_summary(file_id: str):
    """Get AI-generated summary of a document"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="Document not found")
    
    doc = processed_documents[file_id]
    if "summary" not in doc:
        raise HTTPException(status_code=404, detail="Summary not available")
    
    return {
        "file_id": file_id,
        "filename": doc.get("filename"),
        "summary": doc["summary"]
    }

@app.get("/documents/{file_id}/chunks")
async def get_document_chunks(file_id: str):
    """Get text chunks of a document"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="Document not found")
    
    doc = processed_documents[file_id]
    if "chunks" not in doc:
        raise HTTPException(status_code=404, detail="Chunks not available")
    
    return {
        "file_id": file_id,
        "filename": doc.get("filename"),
        "chunks": doc["chunks"],
        "num_chunks": len(doc["chunks"])
    }

@app.get("/documents/{file_id}/full")
async def get_full_document(file_id: str):
    """Get complete document data (text, summary, chunks)"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="Document not found")
    
    return processed_documents[file_id]

@app.delete("/documents/{file_id}")
async def delete_document(file_id: str):
    """Delete a processed document"""
    if file_id not in processed_documents:
        raise HTTPException(status_code=404, detail="Document not found")
    
    # Remove from storage
    del processed_documents[file_id]
    
    # Remove from MCP server
    remove_document(file_id)
    
    # Remove uploaded file
    file_path = UPLOAD_DIR / f"{file_id}.pdf"
    if file_path.exists():
        file_path.unlink()
    
    save_documents()
    
    return {"message": "Document deleted successfully"}

# Keep the old endpoints for backwards compatibility
@app.get("/document/{file_id}")
async def get_document(file_id: str):
    """Get full document data (deprecated - use /documents/{file_id}/full)"""
    return await get_full_document(file_id)

@app.delete("/document/{file_id}")
async def delete_document_old(file_id: str):
    """Delete a processed document (deprecated - use /documents/{file_id})"""
    return await delete_document(file_id)

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)