from flask import Blueprint, request, jsonify
from src.rag.gemeni import GeminiClient
from src.gemini import ask_gemini_with_rag
from src.chroma.vector_store import search_documents
from src.rag.embeddings import get_embedding

gemini_bp = Blueprint('gemini', __name__)
brain = None

def get_brain():
    global brain
    if brain is None:
        brain = GeminiClient()
    return brain

@gemini_bp.route('/ask', methods=['POST'])
def ask():
    data = request.json
    prompt = data.get("prompt")
    context = data.get("context", "")

    if not prompt:
        return jsonify({"error": "No prompt provided"}), 400

    try:
        # Step 1: Retrieve relevant context from ChromaDB (RAG)
        print(f"Generating embedding for prompt: {prompt}")
        query_embedding = get_embedding(prompt)
        
        print("Searching Vector Database for context...")
        search_results = search_documents(query_embedding, num_results=50)
        
        # Special handling for Georgetown University queries to ensure those docs are included
        # even if generic corporate reports outrank them in vector search.
        if "georgetown" in prompt.lower():
            try:
                from src.mongo.connection import get_mongo_client
                client = get_mongo_client()
                # Use hardcoded DB name to match vector_store.py
                db = client["ethix_vectors"]
                collection = db["rag_documents"]
                
                # Fetch docs with Georgetown or MAP_INFO in the filename/source
                gt_docs = list(collection.find({"source": {"$regex": "Georgetown|MAP_INFO", "$options": "i"}}).limit(30))
                
                if gt_docs:
                    print(f"Direct Match: Found {len(gt_docs)} Georgetown specific documents.")
                    for doc in gt_docs:
                        # Normalize to match search_results format
                        search_results.append({
                            "text": doc.get("text", ""),
                            "metadata": doc.get("metadata", {"source": doc.get("source", "Georgetown File")}),
                            "score": 1.0 # High priority
                        })
            except Exception as e:
                print(f"Error checking Georgetown docs: {e}")

        retrieved_context = ""
        if search_results:
            print(f"Found {len(search_results)} documents (total).")
            print("Sources found:")
            seen_sources = set()
            
            for res in search_results:
                # Include metadata if useful, e.g. brand name or date
                meta = res.get('metadata', {})
                source = meta.get('source', 'unknown')
                
                # Deduplication of printing and adding context if exact text overlap?
                # For now just append. LLM can handle duplication.
                
                if source not in seen_sources:
                     print(f" - {source}")
                     seen_sources.add(source)
                
                source_info = f"[Source: {meta.get('type', 'doc')} - {source}]"
                retrieved_context += f"{source_info}\n{res['text']}\n\n"
        else:
            print("No relevant documents found.")

        # Step 2: Combine manual context (if any) with retrieved context
        full_context = context
        if retrieved_context:
            if full_context:
                full_context += "\n\n" + retrieved_context
            else:
                full_context = retrieved_context

        # Step 3: Ask Gemini
        client = get_brain()
        response = client.ask(prompt, full_context)
        
        return jsonify({
            "status": "success",
            "reply": response
        })
    except Exception as e:
        print(f"Error in RAG flow: {e}")
        return jsonify({
            "status": "error",
            "message": str(e)
        }), 500 

@gemini_bp .route ('/rag',methods =['POST'])
def rag ():
    data =request .json 
    prompt =data .get ("prompt")
    category =data .get ("category")

    if not prompt :
        return jsonify ({"error":"No prompt provided"}),400 

    try :
        response =ask_gemini_with_rag (prompt ,category =category )
        return jsonify ({
        "status":"success",
        "reply":response 
        })
    except Exception as e :
        return jsonify ({
        "status":"error",
        "message":str (e )
        }),500 

@gemini_bp .route ('/vision',methods =['POST'])
def vision ():
    return jsonify ({
    "status":"error",
    "message":"Vision endpoint not yet implemented"
    }),501 
