from flask import Flask, request, jsonify
from flask_cors import CORS
import traceback
import os
from queryrun import FAISSQuerySystem
from huggingface_hub import HfFolder

# --- Flask App Setup ---
app = Flask(__name__)
CORS(app)

# --- Initialize the RAG system ---
query_system = None
try:
    print("Initializing RAG system...")
    query_system = FAISSQuerySystem()
    print("RAG system ready for queries.")
except Exception as e:
    print("--- APPLICATION FAILED TO START: RAG SYSTEM INITIALIZATION ERROR ---")
    print(f"Error: {str(e)}")

@app.route('/')
def home():
    return jsonify({
        "status": "ok",
        "message": "RAG API is running",
        "endpoints": {
            "/api/query": "POST - Send queries to the RAG system"
        }
    })

@app.route('/api/query', methods=['POST'])
def handle_query():
    if query_system is None:
        return jsonify({"error": "RAG system is not initialized"}), 500
    
    try:
        data = request.get_json()
        if not data or 'query' not in data:
            return jsonify({"error": "No query provided"}), 400
        
        query = str(data['query'])
        print(f"Received query: {query}")

        search_results = query_system.search(query, k=5)
        response_text = query_system.generate_response(query, search_results)

        sources_for_response = [
            {
                "id": i + 1,
                "score": round(doc['score'], 4),
                "metadata": doc.get('metadata', {})
            }
            for i, doc in enumerate(search_results)
        ]

        return jsonify({
            "response": response_text,
            "sources": sources_for_response
        })

    except Exception as e:
        print(f"Error processing query: {str(e)}")
        traceback.print_exc()
        return jsonify({"error": "An internal error occurred"}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)  # Hugging Face Spaces uses port 7860