import gradio as gr
import spaces
import os
from modules.asr import transcribe
from modules.embed import embed_text
from modules.retriever import load_faiss_index, retrieve_chunks
from modules.llm import ask_llm
from modules.tts import synthesize_speech
from modules.offline_indexer import build_faiss_index
import numpy as np
from typing import List, Tuple, Optional
import json

# Constants for memory management
MAX_MEMORY_TURNS = 5  # Maximum number of conversation turns to remember
MEMORY_FILE = "conversation_memory.json"

class ConversationMemory:
    def __init__(self, max_turns: int = MAX_MEMORY_TURNS):
        self.max_turns = max_turns
        self.conversations: List[Tuple[str, str]] = []
        self.load_memory()
    
    def add_interaction(self, query: str, response: str):
        """Add a new interaction to memory"""
        self.conversations.append((query, response))
        if len(self.conversations) > self.max_turns:
            self.conversations.pop(0)  # Remove oldest interaction
        self.save_memory()
    
    def get_context(self) -> str:
        """Get formatted conversation history for context"""
        if not self.conversations:
            return ""
        
        context = "Previous relevant conversations:\n"
        for i, (q, a) in enumerate(self.conversations, 1):
            context += f"Q{i}: {q}\nA{i}: {a}\n"
        return context
    
    def clear(self):
        """Clear conversation memory"""
        self.conversations = []
        self.save_memory()
    
    def save_memory(self):
        """Save conversations to file"""
        try:
            with open(MEMORY_FILE, "w") as f:
                json.dump(self.conversations, f)
        except Exception as e:
            print(f"Error saving memory: {e}")
    
    def load_memory(self):
        """Load conversations from file"""
        try:
            if os.path.exists(MEMORY_FILE):
                with open(MEMORY_FILE, "r") as f:
                    self.conversations = json.load(f)
        except Exception as e:
            print(f"Error loading memory: {e}")
            self.conversations = []

# Initialize conversation memory
memory = ConversationMemory()

# Try to load FAISS index at startup, but don't crash if it's not available
index, metadata = load_faiss_index()

# Check if we have a knowledge base file but no index
if os.path.exists("knowledge_base.txt") and (index is None or metadata is None):
    try:
        print("Found knowledge base but no index. Building index...")
        build_faiss_index()
        index, metadata = load_faiss_index()
    except Exception as e:
        print(f"Error building index: {e}")

def stream_voice_qa(audio_chunk, history):
    """Process streaming audio in real-time"""
    # Check if index is loaded
    if index is None or metadata is None:
        history.append((None, "No knowledge base loaded. Please upload a knowledge base first."))
        return history

    if audio_chunk is None:
        return history
        
    # Step 1: Transcribe audio chunk
    query = transcribe(audio_chunk)
    if not query.strip():
        return history
    
    # Add user message
    history.append((query, None))
    
    try:
        # Step 2: Generate embeddings for the query
        query_embedding = embed_text(query)
        
        # Step 3: Retrieve relevant context from knowledge base
        context = retrieve_chunks(query_embedding, index, metadata)
        
        # Get conversation memory context
        memory_context = memory.get_context()
        
        # Step 4: Generate answer using LLM with memory context
        prompt = f"""Answer the question based on the following context and previous conversations:

Previous conversations:
{memory_context}

Knowledge base context:
{context}

Question: {query}"""
        
        answer = ask_llm(prompt)
        
        # Step 5: Convert answer to speech
        audio_output = synthesize_speech(answer)
        
        # Update conversation memory
        memory.add_interaction(query, answer)
        
        # Update last message with AI response
        history[-1] = (query, (answer, audio_output))
    except Exception as e:
        # Handle any errors gracefully
        error_msg = f"Error processing query: {str(e)}"
        history[-1] = (query, (error_msg, None))
    
    return history

# Function to update knowledge base
def update_kb(kb_file):
    if kb_file is None:
        return "Error: No file uploaded. Please select a file."
        
    try:
        with open("knowledge_base.txt", "wb") as f:
            f.write(kb_file.read())
        build_faiss_index()
        global index, metadata
        index, metadata = load_faiss_index()
        return "Knowledge base and index updated successfully."
    except Exception as e:
        return f"Error updating knowledge base: {e}"

# Create Gradio interface
with gr.Blocks(title="Real-time Voice Chat with ZeroGPU") as demo:
    gr.Markdown("""
    # 🎙️ Real-time Voice Chat using ZeroGPU
    
    ### Instructions:
    1. Click and hold the microphone button to speak
    2. Release when you're done speaking
    3. Wait for the AI to respond with text and voice
    """)
    
    with gr.Tab("Voice Chat"):
        with gr.Row():
            # Chat history to maintain context
            chatbot = gr.Chatbot(
                [],
                elem_id="chatbot",
                bubble_full_width=False,
                avatar_images=(None, "🤖"),
                height=400
            )
        
        with gr.Row():
            # Audio input with streaming
            audio_input = gr.Audio(
                sources=["microphone"],
                type="filepath",
                streaming=True,
                label="Click and hold to speak",
                elem_id="audio-input"
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                clear_btn = gr.Button("🗑️ Clear Chat", elem_id="clear-btn")
            with gr.Column(scale=2):
                status = gr.Textbox(
                    label="Status",
                    value="Ready! Click and hold the microphone button to speak.",
                    interactive=False
                )
    
    with gr.Tab("Upload Knowledge Base"):
        kb_file = gr.File(label="Upload Knowledge Base (.txt)", type="binary")
        kb_submit = gr.Button("Update Knowledge Base")
        kb_output = gr.Textbox(label="Status")
        
        kb_submit.click(
            fn=update_kb,
            inputs=kb_file,
            outputs=kb_output
        )
        
        # Clear both chat history and conversation memory
        def clear_all():
            memory.clear()
            return []
        
        clear_btn.click(
            fn=clear_all,
            inputs=None,
            outputs=chatbot
        )
        
        # Handle real-time audio streaming
        audio_input.stream(
            fn=stream_voice_qa,
            inputs=[audio_input, chatbot],
            outputs=[chatbot],
            show_progress=False
        )

# Launch the app with queue enabled for better handling of concurrent requests
if __name__ == "__main__":
    demo.queue().launch()