import spaces
import gradio as gr

from pathlib import Path

from src.vectorstore.faiss_index import load_index
from src.vectorstore.metadata_store import load_metadata

from src.retrieval.retriever import retrieve
from src.retrieval.reranker import rerank

from src.llm.generator import generate_answer
from src.llm.postprocessing import clean_answer

from src.data_processing.txt_embeddings import compute_embeddings

from src.voice.assistant import text_to_audio
from src.voice.speech_to_text import transcribe_audio


# ============================================================
# Chargement des ressources
# ============================================================

BASE_DIR = Path(__file__).parent

INDEX_PATH = str(BASE_DIR / "data/vectorstore/faiss.index")
METADATA_PATH = str(BASE_DIR / "data/metadata/metadata.json")


index = load_index(INDEX_PATH)
metadata = load_metadata(METADATA_PATH)


# ============================================================
# Fonction RAG + TTS
# ============================================================
@spaces.GPU

#
def chat_fn(message, history):

    # 1. Embedding question
    query_emb = compute_embeddings([message])


    # 2. Retrieval FAISS
    docs = retrieve(
        query_emb,
        index,
        metadata,
        top_k=10
    )


    # 3. Reranking
    docs = rerank(
        message,
        docs,
        top_k=4
    )


    # 4. Construction contexte
    context = "\n\n".join(
        [
            doc["text"]
            for doc, score in docs
        ]
    )


    # 5. Prompt LLM
    prompt = f"""
Context:
{context}

Question:
{message}

Answer:
"""
    # 6. Génération réponse
    answer = generate_answer(prompt)

    # 7. Nettoyage
    answer = clean_answer(answer)

    # 8. Génération audio Kokoro

    try:
        audio = text_to_audio(answer)
    except Exception:
        audio = None

    # Debug sources
    debug = "\n\n".join(
        [
            f"📄 Chunk {i+1} (Score : {score:.3f})\n{doc['text']}"
            for i, (doc, score) in enumerate(docs)
        ]
    )


    final_output = (
        f"{answer}\n\n"
        "---\n\n"
        f"🔍 Retrieved chunks:\n{debug}"
    )

    return final_output, audio

#
def respond_audio(audio, history):

    print("respond_audio appelée")
    print(audio)

    if audio is None:
        return history, None

    if history is None:
        history = []

    message = transcribe_audio(audio)

    answer, output_audio = chat_fn(
        message,
        history
    )

    history.append(
        {
            "role": "user",
            "content": message
        }
    )

    history.append(
        {
            "role": "assistant",
            "content": answer
        }
    )

    return history, output_audio

# ============================================================
# Interface Gradio
# ============================================================

with gr.Blocks() as demo:

    gr.Markdown(
        "# 💬 RAG Assistant vocal"
    )


    # ========================================================
    # Ligne 1 : Chatbot
    # ========================================================

    chatbot = gr.Chatbot(
        label="Conversation",
        height=450
    )


    # ========================================================
    # Ligne 2 : Question texte
    # ========================================================

    msg = gr.Textbox(
        placeholder="Pose ta question...",
        label="✍️ Question texte"
    )


    # ========================================================
    # Ligne 3 : Audio entrée / sortie
    # ========================================================

    with gr.Row():

        with gr.Column():

            audio_input = gr.Audio(
                sources=[
                    "microphone",
                    "upload"
                ],
                type="filepath",
                waveform_options=gr.WaveformOptions(
                    show_recording_waveform=True
                ),
                label="🎤 Parlez ou déposez un fichier audio"
            )


        with gr.Column():

            audio_output = gr.Audio(
                label="🔊 Réponse audio"
            )


    # ========================================================
    # Réponse texte
    # ========================================================

    def respond(message, history):

        if history is None:
            history = []

        answer, audio = chat_fn(
            message,
            history
        )


        history.append(
            {
                "role": "user",
                "content": message
            }
        )


        history.append(
            {
                "role": "assistant",
                "content": answer
            }
        )


        return (
            "",
            history,
            audio
        )

    # ========================================================
    # Events
    # ========================================================

    msg.submit(
        respond,
        inputs=[
            msg,
            chatbot
        ],
        outputs=[
            msg,
            chatbot,
            audio_output
        ]
    )


    audio_input.change(
        respond_audio,
        inputs=[
            audio_input,
            chatbot
        ],
        outputs=[
            chatbot,
            audio_output
        ]
    )

# ============================================================
# Run
# ============================================================

if __name__ == "__main__":
    demo.launch()