from flask import Flask, Response,request, jsonify
from werkzeug.utils import secure_filename
import os 
from rag_engine import RagEngine
from flask_cors import CORS, cross_origin

app = Flask(__name__)
cors = CORS(app)
app.config["CORS_HEADERS"]= 'Content-Type'

app.config["UPLOAD_FOLDER"]= "uploads"
os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)

rag = RagEngine()

@app.route("/upload", methods=["POST"])
@cross_origin()
def upload_pdf():
    file = request.files.get('file')
    if not file.filename.endswith(".pdf"):
        return jsonify({"error":"Only pdf files are support"}),400
    filename = secure_filename(file.filename)
    filepath = os.path.join(app.config["UPLOAD_FOLDER"], filename)
    file.save(filepath)
    
    try :
        rag.index_pdf(filepath)
    except ValueError as ve:
        return jsonify({"error": str(ve)}), 400
    
    return jsonify({"message":f"file {filename} uploaded and indexed successfully"})


@app.route ("/stream", methods=["POST"])
@cross_origin()
def stream_answer():
    
    question = request.json.get("question", "")
    print (question)
    if not question.strip():
        return jsonify({"error": "Empty question"}), 400

    def generate():
        for token in rag.ask_question_stream(question):
            yield token

    return Response(generate(), mimetype='text/plain')
    
    
@app.route("/ask", methods=["POST"])
@cross_origin()
def ask():
    question = request.json.get("question", "")
    if not question.strip():
        return jsonify({"error": "Empty question"}), 400
    try :
        answer = rag.ask_question(question)
    except Exception as e:
        return jsonify({"error": str(e)}),500
    return jsonify({"message": answer})

@app.route("/stream_answer",methods=["POST"])
@cross_origin()
def stream_question():
    data = request.get_json()
    question = data.get("question","")
    if not question:
        return jsonify({"error": "No question provided"}),400
    def event_stream():
        for token in rag.stream_answer(question=question):
            yield token
    return Response(event_stream(), content_type ="text/event-stream")


if __name__ == "__main__":
      app.run(host="0.0.0.0", port=7860)