from flask import Flask, render_template, request, jsonify
from rag_model import get_qa_chain, build_rag_system, PERSIST_DIRECTORY, HuggingFaceEmbeddings, Chroma
from dotenv import load_dotenv
import os
load_dotenv(override=True)
print(f"App: API Key loaded (starting with {os.getenv('GOOGLE_API_KEY', 'None')[:5]}...)")

app = Flask(__name__)

# Load or Build Vector Store
if not os.path.exists(PERSIST_DIRECTORY):
    print("Vector store not found. Building...")
    vectorstore = build_rag_system()
else:
    print("Loading existing vector store...")
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    vectorstore = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)

qa_chain = get_qa_chain(vectorstore)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/ask', methods=['POST'])
def ask():
    data = request.json
    query = data.get('query')
    print(f"Processing query: {query}")
    
    if not query:
        return jsonify({"error": "No query provided"}), 400
    
    if not qa_chain:
        return jsonify({"error": "QA chain not initialized. Check GOOGLE_API_KEY."}), 500

    try:
        result = qa_chain({"query": query})
        answer = result['result']
        sources = []
        seen = set()
        for doc in result['source_documents']:
            source_name = doc.metadata.get('name', 'Unknown')
            if source_name not in seen:
                sources.append(source_name)
                seen.add(source_name)
        
        return jsonify({
            "answer": answer,
            "sources": sources[:3]  # Return top 3 unique names
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)