# src/main.py
from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from typing import List
import uuid
from datetime import datetime
from pathlib import Path
import os

# Import custom modules1
from src.agents.rag_agent import RAGAgent
from src.models.document import AllDocumentsResponse, StoredDocument
from src.utils.document_processor import DocumentProcessor
from src.utils.conversation_summarizer import ConversationSummarizer
from src.utils.logger import logger
from src.utils.llm_utils import get_llm_instance, get_vector_store
from src.db.mongodb_store import MongoDBStore
from src.implementations.document_service import DocumentService
from src.models import (
    ChatRequest, 
    ChatResponse, 
    DocumentResponse,
    BatchUploadResponse,
    SummarizeRequest,
    SummaryResponse,
    FeedbackRequest
)
from config.config import settings

app = FastAPI(title="Chatbot API")

# Initialize MongoDB
mongodb = MongoDBStore(settings.MONGODB_URI)

# Initialize core components
doc_processor = DocumentProcessor()
summarizer = ConversationSummarizer()
document_service = DocumentService(doc_processor, mongodb)

# Create uploads directory if it doesn't exist
UPLOADS_DIR = Path("uploads")
UPLOADS_DIR.mkdir(exist_ok=True)

# Mount the uploads directory for static file serving
app.mount("/docs", StaticFiles(directory=str(UPLOADS_DIR)), name="documents")

@app.get("/documents")
async def get_all_documents():
    """Get all documents from MongoDB"""
    try:
        documents = await mongodb.get_all_documents()
        
        formatted_documents = []
        for doc in documents:
            try:
                formatted_doc = {
                    "document_id": doc.get("document_id"),
                    "filename": doc.get("filename"),
                    "content_type": doc.get("content_type"),
                    "file_size": doc.get("file_size"),
                    "url_path": doc.get("url_path"),
                    "upload_timestamp": doc.get("upload_timestamp")
                }
                formatted_documents.append(formatted_doc)
            except Exception as e:
                logger.error(f"Error formatting document {doc.get('document_id', 'unknown')}: {str(e)}")
                continue
        
        return {
            "total_documents": len(formatted_documents),
            "documents": formatted_documents
        }
    except Exception as e:
        logger.error(f"Error retrieving documents: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/documents/{document_id}/download")
async def get_document_file(document_id: str):
    """Serve a document file by its ID"""
    try:
        # Get document info from MongoDB
        doc = await mongodb.get_document(document_id)
        if not doc:
            raise HTTPException(status_code=404, detail="Document not found")
        
        # Extract filename from url_path
        filename = doc["url_path"].split("/")[-1]
        file_path = UPLOADS_DIR / filename
        
        if not file_path.exists():
            raise HTTPException(
                status_code=404, 
                detail=f"File not found on server: {filename}"
            )
            
        return FileResponse(
            path=str(file_path),
            filename=doc["filename"],
            media_type=doc["content_type"]
        )
            
    except Exception as e:
        logger.error(f"Error serving document file: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/documents/upload", response_model=BatchUploadResponse)
async def upload_documents(
    files: List[UploadFile] = File(...),
    background_tasks: BackgroundTasks = BackgroundTasks()
):
    """Upload and process multiple documents"""
    try:
        vector_store, _ = await get_vector_store()
        response = await document_service.process_documents(
            files, 
            vector_store, 
            background_tasks
        )
        return response
    except Exception as e:
        logger.error(f"Error in document upload: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))    

@app.get("/documentchunks/{document_id}")
async def get_document_chunks(document_id: str):
    """Get all chunks for a specific document"""
    try:
        vector_store, _ = await get_vector_store()
        chunks = vector_store.get_document_chunks(document_id)
        
        if not chunks:
            raise HTTPException(status_code=404, detail="Document not found")
            
        return {
            "document_id": document_id,
            "total_chunks": len(chunks),
            "chunks": chunks
        }
    except Exception as e:
        logger.error(f"Error retrieving document chunks: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/documents/{document_id}")
async def delete_document(document_id: str):
    """Delete document from MongoDB, ChromaDB, and physical storage"""
    try:
        # First get document details from MongoDB to get file path
        document = await mongodb.get_document(document_id)
        if not document:
            raise HTTPException(status_code=404, detail="Document not found")
        
        # Get vector store instance
        vector_store, _ = await get_vector_store()
        
        # Delete physical file using document service
        deletion_success = await document_service.delete_document(document_id)
        if not deletion_success:
            logger.warning(f"Failed to delete physical file for document {document_id}")
            
        # Delete from vector store
        try:
            vector_store.delete_document(document_id)
        except Exception as e:
            logger.error(f"Error deleting document from vector store: {str(e)}")
            raise HTTPException(
                status_code=500, 
                detail=f"Failed to delete document from vector store: {str(e)}"
            )
            
        # Delete from MongoDB - don't check return value since document might already be deleted
        await mongodb.delete_document(document_id)
        
        return {
            "status": "success",
            "message": f"Document {document_id} successfully deleted from all stores"
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in delete_document endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))
    
@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(
    request: ChatRequest,
    background_tasks: BackgroundTasks
):
    """Chat endpoint with RAG support"""
    try:
        vector_store, embedding_model = await get_vector_store()
        llm = get_llm_instance(request.llm_provider)
        
        rag_agent = RAGAgent(
            llm=llm,
            embedding=embedding_model,
            vector_store=vector_store
        )
        
        if request.stream:
            return StreamingResponse(
                rag_agent.generate_streaming_response(request.query),
                media_type="text/event-stream"
            )
        
        response = await rag_agent.generate_response(
            query=request.query,
            temperature=request.temperature
        )
        
        conversation_id = request.conversation_id or str(uuid.uuid4())
        
        # Store chat history in MongoDB
        await mongodb.store_message(
            conversation_id=conversation_id,
            query=request.query,
            response=response.response,
            context=response.context_docs,
            sources=response.sources,
            llm_provider=request.llm_provider
        )
        
        return ChatResponse(
            response=response.response,
            context=response.context_docs,
            sources=response.sources,
            conversation_id=conversation_id,
            timestamp=datetime.now(),
            relevant_doc_scores=response.scores if hasattr(response, 'scores') else None
        )
        
    except Exception as e:
        logger.error(f"Error in chat endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/chat/history/{conversation_id}")
async def get_conversation_history(conversation_id: str):
    """Get complete conversation history"""
    history = await mongodb.get_conversation_history(conversation_id)
    
    if not history:
        raise HTTPException(status_code=404, detail="Conversation not found")
        
    return {
        "conversation_id": conversation_id,
        "messages": history
    }

@app.post("/chat/summarize", response_model=SummaryResponse)
async def summarize_conversation(request: SummarizeRequest):
    """Generate a summary of a conversation"""
    try:
        messages = await mongodb.get_messages_for_summary(request.conversation_id)
        
        if not messages:
            raise HTTPException(status_code=404, detail="Conversation not found")
        
        summary = await summarizer.summarize_conversation(
            messages,
            include_metadata=request.include_metadata
        )
        
        return SummaryResponse(**summary)
        
    except Exception as e:
        logger.error(f"Error generating summary: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat/feedback/{conversation_id}")
async def submit_feedback(
    conversation_id: str,
    feedback_request: FeedbackRequest
):
    """Submit feedback for a conversation"""
    try:
        success = await mongodb.update_feedback(
            conversation_id=conversation_id,
            feedback=feedback_request.feedback,
            rating=feedback_request.rating
        )
        
        if not success:
            raise HTTPException(status_code=404, detail="Conversation not found")
            
        return {"status": "Feedback submitted successfully"}
        
    except Exception as e:
        logger.error(f"Error submitting feedback: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/debug/config")
async def debug_config():
    """Debug endpoint to check configuration"""
    import os
    from config.config import settings
    from pathlib import Path
    
    debug_info = {
        "environment_variables": {
            "OPENAI_API_KEY": "[SET]" if os.getenv('OPENAI_API_KEY') else "[NOT SET]",
            "OPENAI_MODEL": os.getenv('OPENAI_MODEL', '[NOT SET]')
        },
        "settings": {
            "OPENAI_API_KEY": "[SET]" if settings.OPENAI_API_KEY else "[NOT SET]",
            "OPENAI_MODEL": settings.OPENAI_MODEL,
        },
        "files": {
            "env_file_exists": Path('.env').exists(),
            "openai_config_exists": (Path.home() / '.openai' / 'api_key').exists()
        }
    }
    
    if settings.OPENAI_API_KEY:
        key = settings.OPENAI_API_KEY
        debug_info["api_key_info"] = {
            "length": len(key),
            "preview": f"{key[:4]}...{key[-4:]}" if len(key) > 8 else "[INVALID LENGTH]"
        }
        
    return debug_info

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

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