main.py

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List, Dict
from PIL import Image
import io
import numpy as np
import os
from datetime import datetime
from pymongo import MongoClient
from huggingface_hub import InferenceClient

from embedding_service import JinaClipEmbeddingService
from qdrant_service import QdrantVectorService

# Initialize FastAPI app
app = FastAPI(
    title="Event Social Media Embeddings & ChatbotRAG API",
    description="API để embeddings, search và ChatbotRAG với Jina CLIP v2 + Qdrant + MongoDB + LLM",
    version="2.0.0"
)

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

# Initialize services
print("Initializing services...")
embedding_service = JinaClipEmbeddingService(model_path="jinaai/jina-clip-v2")

collection_name = os.getenv("COLLECTION_NAME", "event_social_media")
qdrant_service = QdrantVectorService(
    collection_name=collection_name,
    vector_size=embedding_service.get_embedding_dimension()
)
print(f"✓ Qdrant collection: {collection_name}")

# MongoDB connection
mongodb_uri = os.getenv("MONGODB_URI", "mongodb+srv://truongtn7122003:7KaI9OT5KTUxWjVI@truongtn7122003.xogin4q.mongodb.net/")
mongo_client = MongoClient(mongodb_uri)
db = mongo_client[os.getenv("MONGODB_DB_NAME", "chatbot_rag")]
documents_collection = db["documents"]
chat_history_collection = db["chat_history"]
print("✓ MongoDB connected")

# Hugging Face token
hf_token = os.getenv("HUGGINGFACE_TOKEN")
if hf_token:
    print("✓ Hugging Face token configured")

print("✓ Services initialized successfully")


# Pydantic models for embeddings
class SearchRequest(BaseModel):
    text: Optional[str] = None
    limit: int = 10
    score_threshold: Optional[float] = None
    text_weight: float = 0.5
    image_weight: float = 0.5


class SearchResponse(BaseModel):
    id: str
    confidence: float
    metadata: dict


class IndexResponse(BaseModel):
    success: bool
    id: str
    message: str


# Pydantic models for ChatbotRAG
class ChatRequest(BaseModel):
    message: str
    use_rag: bool = True
    top_k: int = 3
    system_message: Optional[str] = "You are a helpful AI assistant."
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.95
    hf_token: Optional[str] = None


class ChatResponse(BaseModel):
    response: str
    context_used: List[Dict]
    timestamp: str


class AddDocumentRequest(BaseModel):
    text: str
    metadata: Optional[Dict] = None


class AddDocumentResponse(BaseModel):
    success: bool
    doc_id: str
    message: str


@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "running",
        "service": "Event Social Media Embeddings & ChatbotRAG API",
        "embedding_model": "Jina CLIP v2",
        "vector_db": "Qdrant",
        "language_support": "Vietnamese + 88 other languages",
        "endpoints": {
            "embeddings": {
                "POST /index": "Index data với text/image",
                "POST /search": "Hybrid search",
                "POST /search/text": "Text search",
                "POST /search/image": "Image search",
                "DELETE /delete/{doc_id}": "Delete document",
                "GET /document/{doc_id}": "Get document",
                "GET /stats": "Collection statistics"
            },
            "chatbot_rag": {
                "POST /chat": "Chat với RAG",
                "POST /documents": "Add document to knowledge base",
                "POST /rag/search": "Search in knowledge base",
                "GET /history": "Get chat history",
                "DELETE /documents/{doc_id}": "Delete document from knowledge base"
            }
        }
    }


@app.post("/index", response_model=IndexResponse)
async def index_data(
    id: str = Form(...),
    text: str = Form(...),
    image: Optional[UploadFile] = File(None)
):
    """
    Index data vào vector database

    Body:
    - id: Document ID (event ID, post ID, etc.)
    - text: Text content (tiếng Việt supported)
    - image: Image file (optional)

    Returns:
    - success: True/False
    - id: Document ID
    - message: Status message
    """
    try:
        # Prepare embeddings
        text_embedding = None
        image_embedding = None

        # Encode text (tiếng Việt)
        if text and text.strip():
            text_embedding = embedding_service.encode_text(text)

        # Encode image nếu có
        if image:
            image_bytes = await image.read()
            pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
            image_embedding = embedding_service.encode_image(pil_image)

        # Combine embeddings
        if text_embedding is not None and image_embedding is not None:
            # Average của text và image embeddings
            combined_embedding = np.mean([text_embedding, image_embedding], axis=0)
        elif text_embedding is not None:
            combined_embedding = text_embedding
        elif image_embedding is not None:
            combined_embedding = image_embedding
        else:
            raise HTTPException(status_code=400, detail="Phải cung cấp ít nhất text hoặc image")

        # Normalize
        combined_embedding = combined_embedding / np.linalg.norm(combined_embedding, axis=1, keepdims=True)

        # Index vào Qdrant
        metadata = {
            "text": text,
            "has_image": image is not None,
            "image_filename": image.filename if image else None
        }

        result = qdrant_service.index_data(
            doc_id=id,
            embedding=combined_embedding,
            metadata=metadata
        )

        return IndexResponse(
            success=True,
            id=result["original_id"],  # Trả về MongoDB ObjectId
            message=f"Đã index thành công document {result['original_id']} (Qdrant UUID: {result['qdrant_id']})"
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi index: {str(e)}")


@app.post("/search", response_model=List[SearchResponse])
async def search(
    text: Optional[str] = Form(None),
    image: Optional[UploadFile] = File(None),
    limit: int = Form(10),
    score_threshold: Optional[float] = Form(None),
    text_weight: float = Form(0.5),
    image_weight: float = Form(0.5)
):
    """
    Search similar documents bằng text và/hoặc image

    Body:
    - text: Query text (tiếng Việt supported)
    - image: Query image (optional)
    - limit: Số lượng kết quả (default: 10)
    - score_threshold: Minimum confidence score (0-1)
    - text_weight: Weight cho text search (default: 0.5)
    - image_weight: Weight cho image search (default: 0.5)

    Returns:
    - List of results với id, confidence, và metadata
    """
    try:
        # Prepare query embeddings
        text_embedding = None
        image_embedding = None

        # Encode text query
        if text and text.strip():
            text_embedding = embedding_service.encode_text(text)

        # Encode image query
        if image:
            image_bytes = await image.read()
            pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
            image_embedding = embedding_service.encode_image(pil_image)

        # Validate input
        if text_embedding is None and image_embedding is None:
            raise HTTPException(status_code=400, detail="Phải cung cấp ít nhất text hoặc image để search")

        # Hybrid search với Qdrant
        results = qdrant_service.hybrid_search(
            text_embedding=text_embedding,
            image_embedding=image_embedding,
            text_weight=text_weight,
            image_weight=image_weight,
            limit=limit,
            score_threshold=score_threshold,
            ef=256  # High accuracy search
        )

        # Format response
        return [
            SearchResponse(
                id=result["id"],
                confidence=result["confidence"],
                metadata=result["metadata"]
            )
            for result in results
        ]

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi search: {str(e)}")


@app.post("/search/text", response_model=List[SearchResponse])
async def search_by_text(
    text: str = Form(...),
    limit: int = Form(10),
    score_threshold: Optional[float] = Form(None)
):
    """
    Search chỉ bằng text (tiếng Việt)

    Body:
    - text: Query text (tiếng Việt)
    - limit: Số lượng kết quả
    - score_threshold: Minimum confidence score

    Returns:
    - List of results
    """
    try:
        # Encode text
        text_embedding = embedding_service.encode_text(text)

        # Search
        results = qdrant_service.search(
            query_embedding=text_embedding,
            limit=limit,
            score_threshold=score_threshold,
            ef=256
        )

        return [
            SearchResponse(
                id=result["id"],
                confidence=result["confidence"],
                metadata=result["metadata"]
            )
            for result in results
        ]

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi search: {str(e)}")


@app.post("/search/image", response_model=List[SearchResponse])
async def search_by_image(
    image: UploadFile = File(...),
    limit: int = Form(10),
    score_threshold: Optional[float] = Form(None)
):
    """
    Search chỉ bằng image

    Body:
    - image: Query image
    - limit: Số lượng kết quả
    - score_threshold: Minimum confidence score

    Returns:
    - List of results
    """
    try:
        # Encode image
        image_bytes = await image.read()
        pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        image_embedding = embedding_service.encode_image(pil_image)

        # Search
        results = qdrant_service.search(
            query_embedding=image_embedding,
            limit=limit,
            score_threshold=score_threshold,
            ef=256
        )

        return [
            SearchResponse(
                id=result["id"],
                confidence=result["confidence"],
                metadata=result["metadata"]
            )
            for result in results
        ]

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi search: {str(e)}")


@app.delete("/delete/{doc_id}")
async def delete_document(doc_id: str):
    """
    Delete document by ID (MongoDB ObjectId hoặc UUID)

    Args:
    - doc_id: Document ID to delete

    Returns:
    - Success message
    """
    try:
        qdrant_service.delete_by_id(doc_id)
        return {"success": True, "message": f"Đã xóa document {doc_id}"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi xóa: {str(e)}")


@app.get("/document/{doc_id}")
async def get_document(doc_id: str):
    """
    Get document by ID (MongoDB ObjectId hoặc UUID)

    Args:
    - doc_id: Document ID (MongoDB ObjectId)

    Returns:
    - Document data
    """
    try:
        doc = qdrant_service.get_by_id(doc_id)
        if doc:
            return {
                "success": True,
                "data": doc
            }
        raise HTTPException(status_code=404, detail=f"Không tìm thấy document {doc_id}")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi get document: {str(e)}")


@app.get("/stats")
async def get_stats():
    """
    Lấy thông tin thống kê collection

    Returns:
    - Collection statistics
    """
    try:
        info = qdrant_service.get_collection_info()
        return info
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Lỗi khi lấy stats: {str(e)}")


# ============================================
# ChatbotRAG Endpoints
# ============================================

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """
    Chat endpoint với RAG

    Body:
    - message: User message
    - use_rag: Enable RAG retrieval (default: true)
    - top_k: Number of documents to retrieve (default: 3)
    - system_message: System prompt (optional)
    - max_tokens: Max tokens for response (default: 512)
    - temperature: Temperature for generation (default: 0.7)
    - hf_token: Hugging Face token (optional, sẽ dùng env nếu không truyền)

    Returns:
    - response: Generated response
    - context_used: Retrieved context documents
    - timestamp: Response timestamp
    """
    try:
        # Retrieve context if RAG enabled
        context_used = []
        if request.use_rag:
            # Generate query embedding
            query_embedding = embedding_service.encode_text(request.message)
            
            # Search in Qdrant
            results = qdrant_service.search(
                query_embedding=query_embedding,
                limit=request.top_k,
                score_threshold=0.5
            )
            context_used = results

        # Build context text
        context_text = ""
        if context_used:
            context_text = "\n\nRelevant Context:\n"
            for i, doc in enumerate(context_used, 1):
                doc_text = doc["metadata"].get("text", "")
                confidence = doc["confidence"]
                context_text += f"\n[{i}] (Confidence: {confidence:.2f})\n{doc_text}\n"

            # Add context to system message
            system_message = f"{request.system_message}\n{context_text}\n\nPlease use the above context to answer the user's question when relevant."
        else:
            system_message = request.system_message

        # Use token from request or fallback to env
        token = request.hf_token or hf_token

        # Generate response
        if not token:
            response = f"""[LLM Response Placeholder]

Context retrieved: {len(context_used)} documents
User question: {request.message}

To enable actual LLM generation:
1. Set HUGGINGFACE_TOKEN environment variable, OR
2. Pass hf_token in request body

Example:
{{
  "message": "Your question",
  "hf_token": "hf_xxxxxxxxxxxxx"
}}
"""
        else:
            try:
                client = InferenceClient(
                    token=token,
                    model="openai/gpt-oss-20b"
                )

                # Build messages
                messages = [
                    {"role": "system", "content": system_message},
                    {"role": "user", "content": request.message}
                ]

                # Generate response
                response = ""
                for msg in client.chat_completion(
                    messages,
                    max_tokens=request.max_tokens,
                    stream=True,
                    temperature=request.temperature,
                    top_p=request.top_p,
                ):
                    choices = msg.choices
                    if len(choices) and choices[0].delta.content:
                        response += choices[0].delta.content

            except Exception as e:
                response = f"Error generating response with LLM: {str(e)}\n\nContext was retrieved successfully, but LLM generation failed."

        # Save to history
        chat_data = {
            "user_message": request.message,
            "assistant_response": response,
            "context_used": context_used,
            "timestamp": datetime.utcnow()
        }
        chat_history_collection.insert_one(chat_data)

        return ChatResponse(
            response=response,
            context_used=context_used,
            timestamp=datetime.utcnow().isoformat()
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@app.post("/documents", response_model=AddDocumentResponse)
async def add_document(request: AddDocumentRequest):
    """
    Add document to knowledge base

    Body:
    - text: Document text
    - metadata: Additional metadata (optional)

    Returns:
    - success: True/False
    - doc_id: MongoDB document ID
    - message: Status message
    """
    try:
        # Save to MongoDB
        doc_data = {
            "text": request.text,
            "metadata": request.metadata or {},
            "created_at": datetime.utcnow()
        }
        result = documents_collection.insert_one(doc_data)
        doc_id = str(result.inserted_id)

        # Generate embedding
        embedding = embedding_service.encode_text(request.text)

        # Index to Qdrant
        qdrant_service.index_data(
            doc_id=doc_id,
            embedding=embedding,
            metadata={
                "text": request.text,
                "source": "api",
                **(request.metadata or {})
            }
        )

        return AddDocumentResponse(
            success=True,
            doc_id=doc_id,
            message=f"Document added successfully with ID: {doc_id}"
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@app.post("/rag/search", response_model=List[SearchResponse])
async def rag_search(
    query: str = Form(...),
    top_k: int = Form(5),
    score_threshold: Optional[float] = Form(0.5)
):
    """
    Search in knowledge base

    Body:
    - query: Search query
    - top_k: Number of results (default: 5)
    - score_threshold: Minimum score (default: 0.5)

    Returns:
    - results: List of matching documents
    """
    try:
        # Generate query embedding
        query_embedding = embedding_service.encode_text(query)

        # Search in Qdrant
        results = qdrant_service.search(
            query_embedding=query_embedding,
            limit=top_k,
            score_threshold=score_threshold
        )

        return [
            SearchResponse(
                id=result["id"],
                confidence=result["confidence"],
                metadata=result["metadata"]
            )
            for result in results
        ]

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@app.get("/history")
async def get_history(limit: int = 10, skip: int = 0):
    """
    Get chat history

    Query params:
    - limit: Number of messages to return (default: 10)
    - skip: Number of messages to skip (default: 0)

    Returns:
    - history: List of chat messages
    """
    try:
        history = list(
            chat_history_collection
            .find({}, {"_id": 0})
            .sort("timestamp", -1)
            .skip(skip)
            .limit(limit)
        )

        # Convert datetime to string
        for msg in history:
            if "timestamp" in msg:
                msg["timestamp"] = msg["timestamp"].isoformat()

        return {
            "history": history,
            "total": chat_history_collection.count_documents({})
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@app.delete("/documents/{doc_id}")
async def delete_document_from_kb(doc_id: str):
    """
    Delete document from knowledge base

    Args:
    - doc_id: Document ID (MongoDB ObjectId)

    Returns:
    - success: True/False
    - message: Status message
    """
    try:
        # Delete from MongoDB
        result = documents_collection.delete_one({"_id": doc_id})

        # Delete from Qdrant
        if result.deleted_count > 0:
            qdrant_service.delete_by_id(doc_id)
            return {"success": True, "message": f"Document {doc_id} deleted from knowledge base"}
        else:
            raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


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