import os
import time
import json
import keyboard  # pip install keyboard==0.13.5
import numpy as np
import faiss
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed

# Import the AudioProcessor class from your STT module.
from stt import AudioProcessor
from llm import LLMProcessor
# Import vectorstore functions from vectorstore.py (located in app/)
from vectorstore import (
    update_vectorstore_from_pdf,
    save_vectorstore_with_timestamp,
    load_vectorstore,
    search_vectorstore
)
# Import TTS functions from the tts folder.
from tts import chunk_text_for_tts, synthesize_speech_guideline, play_audio


def combine_vectorstores(index1, chunks1, index2, chunks2):
    """
    Combine two FAISS indexes (both IndexFlatL2) by retrieving all embeddings,
    concatenating them, and building a new FAISS index.
    Also combines the two lists of text chunks.
    """
    embeddings1 = np.array([index1.reconstruct(i) for i in range(index1.ntotal)])
    embeddings2 = np.array([index2.reconstruct(i) for i in range(index2.ntotal)])
    combined_embeddings = np.concatenate([embeddings1, embeddings2], axis=0)
    combined_chunks = chunks1 + chunks2
    dimension = combined_embeddings.shape[1]
    combined_index = faiss.IndexFlatL2(dimension)
    combined_index.add(combined_embeddings)
    return combined_index, combined_chunks


def build_and_save_vectorstore():
    """
    Process two PDFs, combine their embeddings into one vectorstore,
    save the combined index with a timestamp, and return the index,
    chunks, and embedder used.
    """
    base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    pdf_path1 = os.path.join(base_dir, "data", "PDF", "cancer_dictionary.pdf")
    pdf_path2 = os.path.join(base_dir, "data", "PDF", "Medical Dictionary.pdf")
    
    index1, chunks1, embedder = update_vectorstore_from_pdf(pdf_path1, chunk_size=600)
    index2, chunks2, _ = update_vectorstore_from_pdf(pdf_path2, chunk_size=600)
    
    combined_index, combined_chunks = combine_vectorstores(index1, chunks1, index2, chunks2)
    
    saved_index_path = save_vectorstore_with_timestamp(combined_index)
    print(f"Combined vectorstore saved at: {saved_index_path}")
    
    loaded_index = load_vectorstore(saved_index_path)
    return loaded_index, combined_chunks, embedder


def synthesize_audio_chunks(text, max_chars=300):
    """
    Breaks the given text into chunks and asynchronously synthesizes them.
    Yields the output file for each synthesized audio chunk.
    """
    if not text:
        raise ValueError("No text provided for TTS synthesis.")

    # Break the text into manageable chunks.
    chunks = chunk_text_for_tts(text, max_chars=max_chars)
    print("TTS Text chunks:")
    for i, chunk in enumerate(chunks, start=1):
        print(f"Chunk {i}: {chunk}\n")
    
    # Use a ThreadPoolExecutor to process the chunks concurrently.
    with ThreadPoolExecutor(max_workers=4) as executor:
        # Submit synthesis tasks and keep a mapping of future to its chunk index.
        future_to_index = {
            executor.submit(synthesize_speech_guideline, chunk, output_file=f"audio_chunk_{i}.wav"): i
            for i, chunk in enumerate(chunks, start=1)
        }
        
        # Yield audio files as soon as each future completes.
        for future in as_completed(future_to_index):
            index = future_to_index[future]
            try:
                audio_file = future.result()
                if audio_file is None:
                    print(f"Warning: Synthesis for chunk {index} returned None.")
                else:
                    print(f"Chunk {index} synthesized successfully: {audio_file}")
                    yield audio_file
            except Exception as exc:
                print(f"Error synthesizing chunk {index}: {exc}")

def synthesize_audio_chunks(text, max_chars=300):
    """
    Breaks the given text into chunks and asynchronously synthesizes them.
    Yields the output file for each synthesized audio chunk.
    """
    if not text:
        raise ValueError("No text provided for TTS synthesis.")
    chunks = chunk_text_for_tts(text, max_chars=max_chars)
    print("TTS Text chunks:")
    for i, chunk in enumerate(chunks, 1):
        print(f"Chunk {i}: {chunk}\n")
    
    with ThreadPoolExecutor(max_workers=4) as executor:
        # Submit synthesis tasks for each chunk
        futures = [executor.submit(synthesize_speech_guideline, chunk, output_file=f"audio_chunk_{i}.wav")
                    for i, chunk in enumerate(chunks)]
        # Yield audio files as soon as each synthesis task completes
        for future in futures:
            audio_file = future.result()
            yield audio_file


def main():
    print("Building vectorstore from PDFs...")
    vector_index, vector_chunks, embedder = build_and_save_vectorstore()
    
    processor = AudioProcessor()
    llm_processor = LLMProcessor()
    
    print("\nLooping process: Use push-to-talk to record and transcribe.")
    print("Press 's' at any time to stop the program.")
    Saved_response = ""
    
    while True:
        if keyboard.is_pressed('s'):
            print("Stop key pressed. Exiting program.")
            break
        
        print("\nReady to record. Press 'q' to start and 'q' again to stop recording.\n----------------\n")
        denoised_audio, fs = processor.record_and_denoise()
        
        if denoised_audio is None:
            print("No audio was recorded. Skipping transcription.")
        else:
            print("Transcribing audio...")
            try:
                transcription = processor.transcribe_audio(processor.default_output_file, language="en")
                print("Transcribed Text:\n", transcription)
                
                search_results = search_vectorstore(transcription, embedder, vector_index, vector_chunks, top_k=3)
                query = Saved_response + transcription
                for i, result in enumerate(search_results, 1):
                    query += f"\n{i}. {result}"
                
                try:
                    gemini_response = llm_processor.call_gemini_llm("gemini-2.0-flash", query)
                    print("Gemini response:", gemini_response)
                except Exception as e:
                    print("Gemini error:", e)
                    gemini_response = f"Error: {str(e)}"
                
                Saved_response += "transcription: " + transcription + "\n" + "gemini_response: " + gemini_response + "\n"
                processor.log_conversation(transcription, bot_text=gemini_response)
                
                # Synthesize and play TTS output in a pipelined fashion.
                print("Synthesizing and playing LLM response via TTS...")
                print()
                audio_generator = synthesize_audio_chunks(gemini_response, max_chars=300)
                # Immediately get the first audio chunk and play it
                try:
                    first_audio = next(audio_generator)
                    play_audio(first_audio)
                except StopIteration:
                    print("No audio chunks generated.")
                
                # Now play the rest of the chunks as soon as they are ready.
                for audio_file in audio_generator:
                    play_audio(audio_file)
                
            except Exception as e:
                print("Error during transcription:", e)
        
        print("Iteration complete. Waiting 5 seconds before next recording...")
        time.sleep(5)
    
    print("Program terminated.")



if __name__ == "__main__":
    main()