from flask import Flask, request, render_template, send_file
from werkzeug.utils import secure_filename
from src.ollama_chain import OllamaChain, OllamaRAGChain
from src.llama_cpp_chains import LlamaChain
from src.pdf_handler import extract_pdf
from src.vqa import answer_visual_question
from src.audio_processor import AudioProcessor
from dotenv import load_dotenv
import os

app = Flask(__name__)
load_dotenv()

UPLOAD_FOLDER = './.cache/temp_files'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

audio_processor = AudioProcessor()
chat_memory = []  # Replace with appropriate chat memory class if needed
llm_chain = None
pdf_mode = False


def get_chain():
    global llm_chain
    if llm_chain is None:
        llm_chain = OllamaRAGChain(chat_memory) if pdf_mode else OllamaChain(chat_memory)
    return llm_chain


@app.route("/", methods=["GET", "POST"])
def index():
    global llm_chain, pdf_mode

    response_text = ""
    audio_file_path = ""

    if request.method == "POST":
        user_input = request.form.get("user_input")
        uploaded_pdf = request.files.getlist("uploaded_pdf")
        uploaded_image = request.files.get("uploaded_image")
        uploaded_audio = request.files.get("uploaded_audio")
        pdf_mode = True if request.form.get("pdf_chat") == "on" else False

        chain = get_chain()

        # PDF processing
        if pdf_mode and uploaded_pdf:
            try:
                chain.update_chain(uploaded_pdf)
            except Exception as e:
                return f"Error updating knowledge base: {e}"

        try:
            if uploaded_image and uploaded_image.filename != '':
                image_path = os.path.join(UPLOAD_FOLDER, secure_filename(uploaded_image.filename))
                uploaded_image.save(image_path)
                response_text = answer_visual_question(image_path, user_input)

            elif uploaded_audio and uploaded_audio.filename != '':
                audio_path = os.path.join(UPLOAD_FOLDER, secure_filename(uploaded_audio.filename))
                uploaded_audio.save(audio_path)
                question = audio_processor.audio_to_text(audio_path)
                response_text = chain.run(question)

            else:
                response_text = chain.run(user_input)

            audio_file_path = audio_processor.text_to_speech(response_text)

        except Exception as e:
            return f"Error during chat: {e}"

    return render_template("index.html", response=response_text, audio_path=audio_file_path)


@app.route("/audio")
def audio():
    path = request.args.get("path")
    return send_file(path, mimetype="audio/mp3")


if __name__ == "__main__":
    app.run(debug=True)