import mss
from PIL import Image
import pytesseract
import speech_recognition as sr
import pyttsx3
import asyncio
import aiohttp
from bs4 import BeautifulSoup
import ollama  # Ensure you have the Ollama SDK installed and configured
import keyboard
import threading
import time
import re
import logging
import tkinter as tk
from word2number import w2n  # Optional for number word conversion

# -------------------- Configuration --------------------

# Configure logging to capture DEBUG level logs
logging.basicConfig(
    level=logging.DEBUG,  # Set to DEBUG for detailed logs
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),  # Outputs to console
        # Uncomment the next line to enable logging to a file
        # logging.FileHandler("qa_assistant.log"),
    ]
)

# Initialize the TTS engine globally
tts_engine = pyttsx3.init()
is_listening = False  # Initialize this global variable at the start
speech_text = ""  # Initialize this global variable at the start

def transcribe_speech():
    global is_listening, speech_text
    recognizer = sr.Recognizer()
    with sr.Microphone() as source:
        logging.info("Listening... Please say the question number (e.g., 'Question 597').")
        recognizer.adjust_for_ambient_noise(source, duration=1)  # Adjust for noise for 1 second
        try:
            # Listen for speech, but stop after a timeout
            audio = recognizer.listen(source, timeout=10, phrase_time_limit=5)
            logging.info("Processing the audio with Google Speech Recognition...")
            text = recognizer.recognize_google(audio)
            # Convert number words to digits
            speech_text = convert_number_words(text)
            logging.info(f"Speech Transcribed: {speech_text}")
            
            # Provide confirmation
            speak_text(f"I heard question number {speech_text}.")
        except sr.WaitTimeoutError:
            logging.warning("Listening timed out, no speech detected.")
        except sr.UnknownValueError:
            logging.warning("Speech recognition could not understand audio.")
        except sr.RequestError as e:
            logging.error(f"Could not request results from speech recognition service; {e}")
        finally:
            with lock:
                is_listening = False  # Ensure listening flag is reset

def configure_tts_voice():
    global tts_engine
    voices = tts_engine.getProperty('voices')
    
    # List available voices (optional)
    for index, voice in enumerate(voices):
        logging.debug(f"Voice {index}: {voice.name}, {voice.gender}, {voice.languages}")
    
    # Change voice (select male/female, adjust the index as per your preference)
    if len(voices) > 1:
        tts_engine.setProperty('voice', voices[1].id)  # Change index to select a different voice
        logging.info(f"TTS voice set to: {voices[1].name}")
    else:
        logging.info("Only one voice available. Using default voice.")
    
    # Adjust speaking rate (lower is slower and more human-like)
    tts_engine.setProperty('rate', 160)  # Default is 200, lowering makes it more human-like
    
    # Adjust volume (0.0 to 1.0)
    tts_engine.setProperty('volume', 1.0)

# -------------------- Helper Functions --------------------

def normalize_option_label(label):
    """
    Normalizes OCR-misread option labels to standard A-E.
    """
    label = label.upper()
    if label in ['A', 'B', 'C', 'D', 'E']:
        return label
    # Handle common OCR misinterpretations
    ocr_misreads = {
        '©': 'C',
        'S': '5',
        'À': 'A',
        'ß': 'B',
        '(': '',  # Remove any stray characters
        ')': '',
        '-': '',
        # Add more mappings if needed
    }
    return ocr_misreads.get(label, label)  # Default to original if no mapping found

def convert_number_words(text):
    """
    Converts number words in the speech text to digits.
    Example: "five hundred ninety-seven" -> "597"
    """
    try:
        # Attempt to convert the entire text
        number = w2n.word_to_num(text)
        return str(number)
    except:
        # Attempt to extract digits directly if conversion fails
        digits = re.findall(r'\d+', text)
        return ''.join(digits) if digits else text  # Return the number if found, else original text

def perform_ocr(image):
    global ocr_text
    try:
        logging.info("Performing OCR on the selected screen region...")
        raw_text = pytesseract.image_to_string(image)
        ocr_text = raw_text
        logging.debug("Raw OCR Text:")
        logging.debug(ocr_text)  # Log the raw OCR text for debugging
        logging.info("OCR complete.")
    except Exception as e:
        logging.error(f"OCR Error: {e}")
        ocr_text = ""

def get_selected_region():
    """
    Allows the user to select a screen region by drawing a rectangle.
    Returns the coordinates of the selected region.
    """
    # Initialize the tkinter window
    root = tk.Tk()
    root.attributes('-alpha', 0.3)  # Transparency
    root.attributes("-topmost", True)
    root.overrideredirect(True)  # Remove window decorations

    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()
    root.geometry(f"{screen_width}x{screen_height}")

    canvas = tk.Canvas(root, cursor="cross", bg="grey")
    canvas.pack(fill=tk.BOTH, expand=True)

    rect = None
    start_x = start_y = end_x = end_y = 0

    def on_button_press(event):
        nonlocal start_x, start_y, rect
        start_x = event.x
        start_y = event.y
        rect = canvas.create_rectangle(start_x, start_y, start_x, start_y, outline='red')

    def on_move_press(event):
        nonlocal end_x, end_y
        end_x, end_y = event.x, event.y
        canvas.coords(rect, start_x, start_y, end_x, end_y)

    def on_button_release(event):
        root.quit()
        root.destroy()  # Destroy the window to remove the overlay


    canvas.bind("<ButtonPress-1>", on_button_press)
    canvas.bind("<B1-Motion>", on_move_press)
    canvas.bind("<ButtonRelease-1>", on_button_release)

    root.mainloop()

    # Calculate absolute coordinates
    left = min(start_x, end_x)
    top = min(start_y, end_y)
    right = max(start_x, end_x)
    bottom = max(start_y, end_y)

    logging.debug(f"Selected region coordinates: Left={left}, Top={top}, Right={right}, Bottom={bottom}")
    return (left, top, right, bottom)

def speak_text(text):
    global tts_engine
    logging.info(f"Speaking out the response: {text}")
    with tts_lock:
        try:
            # Speak the text
            tts_engine.say(text)

            # Add hotkeys for pausing (F4) and stopping (F5)
            keyboard.add_hotkey('f4', pause_speaking)
            keyboard.add_hotkey('f5', stop_speaking)

            tts_engine.runAndWait()
            logging.info("Text-to-Speech completed.")
        except RuntimeError as e:
            logging.error(f"TTS Error: {e}")
        finally:
            # Remove hotkeys after completion to avoid conflicts
            keyboard.remove_hotkey('f4')
            keyboard.remove_hotkey('f5')

def pause_speaking():
    global tts_engine
    logging.info("Paused speaking...")
    tts_engine.stop()

def stop_speaking():
    global tts_engine
    logging.info("Stopped speaking.")
    tts_engine.stop()

async def async_web_search(query):
    logging.info(f"Searching the web for query: {query}")
    url = f"https://html.duckduckgo.com/html/?q={query}"
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                html = await response.text()
    except Exception as e:
        logging.error(f"Web search failed: {e}")
        return []

    soup = BeautifulSoup(html, 'html.parser')
    results = []
    for result in soup.find_all('div', class_='result__body')[:3]:
        title_element = result.find('a', class_='result__a')
        snippet_element = result.find('a', class_='result__snippet')
        title = title_element.text.strip() if title_element else "No Title"
        snippet = snippet_element.text.strip() if snippet_element else "No Snippet"
        results.append({'title': title, 'snippet': snippet})

    logging.info(f"Web search complete. Found {len(results)} results.")
    return results

def generate_search_query(ocr_text, question_number):
    logging.info("Generating search query with LLM...")
    context = f"""You are an assistant that generates concise and effective search queries based on OCR text and a multiple-choice question number.

OCR Text:
{ocr_text}

Question Number: {question_number}

Generate a search query that accurately represents the question and options to find relevant information for answering the question."""
    
    try:
        response = ollama.generate(model='llama3.2', prompt=context)
        search_query = response.get('response', '').strip()
        logging.info(f"Generated Search Query: {search_query}")
        return search_query
    except Exception as e:
        logging.error(f"LLM Error during search query generation: {e}")
        return ""

def generate_final_response(search_results, ocr_text, question_number):
    logging.info("Generating final response with LLM...")
    # Format the search results for better readability
    formatted_results = "\n".join([f"Title: {res['title']}\nSnippet: {res['snippet']}\n" for res in search_results])

    context = f"""You are an intelligent assistant specialized in answering multiple-choice questions based on OCR text and web search results.

OCR Text:
{ocr_text}

Question Number: {question_number}

Web Search Results:
{formatted_results}

Based on the OCR text and the search results above, select the most appropriate answer option (A, B, C, etc.) and provide only the option letter as the final answer."""
    
    try:
        response = ollama.generate(model='llama3.2', prompt=context)
        final_response = response.get('response', '').strip()
        logging.info(f"Generated Final Response: {final_response}")
        # Validate the response format (should start with A., B., etc.)
        if not re.match(r'^[A-E]\.', final_response):
            logging.error("LLM response format is invalid.")
            speak_text("I received an unexpected answer format. Please try again.")
            return ""
        return final_response
    except Exception as e:
        logging.error(f"LLM Error during final response generation: {e}")
        return ""

def reset_state():
    global speech_text, ocr_text
    speech_text = ""
    ocr_text = ""

# -------------------- Core Functions --------------------

def start_listening():
    global is_listening
    with lock:
        if not is_listening:
            is_listening = True
            logging.info("Started listening...")
            # Perform region selection and OCR in a separate thread
            threading.Thread(target=process_initial_inputs, daemon=True).start()

def process_initial_inputs():
    global speech_text, ocr_text
    # Step 1: Let the user select a region
    region = get_selected_region()
    left, top, right, bottom = region
    width = right - left
    height = bottom - top
    if width <= 0 or height <= 0:
        logging.error("Invalid region selected. Please try again.")
        speak_text("Invalid region selected. Please try again by pressing F2.")
        reset_state()
        return
    
    # Step 2: Capture the selected region
    with mss.mss() as sct:
        monitor = {
            "top": top,
            "left": left,
            "width": width,
            "height": height
        }
        try:
            screenshot = sct.grab(monitor)
            img = Image.frombytes("RGB", screenshot.size, screenshot.bgra, "raw", "BGRX")
        except Exception as e:
            logging.error(f"Screen capture failed: {e}")
            speak_text("Screen capture failed. Please ensure the selected area is correct and try again.")
            reset_state()
            return
    
    # Step 3: Perform OCR on the captured region
    perform_ocr(img)
    if not ocr_text.strip():
        logging.error("OCR returned empty text. Please ensure the selected region contains readable text.")
        speak_text("OCR did not detect any text. Please ensure the selected area contains readable text and try again.")
        reset_state()
        return
    
    # Step 4: Listen for the question number
    transcribe_speech()
    if not speech_text:
        logging.warning("No speech input received.")
        speak_text("I could not receive any speech input. Please try again by pressing F2.")
        reset_state()
        return
    
    # Step 5: Generate and execute the search query
    generate_and_execute_search(speech_text, ocr_text)

def generate_and_execute_search(question_number, ocr_text):
    # Generate search query using LLM
    search_query = generate_search_query(ocr_text, question_number)
    if not search_query:
        logging.error("Failed to generate a search query.")
        speak_text("I failed to generate a search query. Please try again by pressing F2.")
        reset_state()
        return
    
    # Perform web search and process the answer
    asyncio.run(async_process_search(search_query, question_number, ocr_text))

async def async_process_search(search_query, question_number, ocr_text):
    # Perform web search
    search_results = await async_web_search(search_query)
    if not search_results:
        logging.warning("No search results found. Unable to answer the question.")
        speak_text("I'm sorry, I couldn't find an answer to your question.")
        reset_state()
        return
    
    # Generate final response using LLM
    final_response = generate_final_response(search_results, ocr_text, question_number)
    if not final_response:
        logging.error("Failed to generate a final response.")
        speak_text("I couldn't determine the answer to your question. Please try again by pressing F2.")
        reset_state()
        return
    
    # Speak out the answer
    speak_text(final_response)
    reset_state()

def stop_listening():
    global is_listening
    with lock:
        if is_listening:
            is_listening = False
            logging.info("Stopped listening.")
            speak_text("Stopped listening.")

# -------------------- Main Function --------------------

async def main():
    global speech_text, ocr_text

    configure_tts_voice()  # Voice configuration happens before anything else

    # Set up hotkeys
    keyboard.add_hotkey('f2', start_listening)
    keyboard.add_hotkey('f3', stop_listening)

    logging.info("QA Assistant is running.")
    logging.info("Press 'F2' to select a region and answer a question, 'F3' to stop listening, 'Esc' to exit.")

    while True:
        # Wait for a short duration to avoid busy waiting
        await asyncio.sleep(0.1)

        # Check for exit condition
        if keyboard.is_pressed('esc'):
            logging.info("Exiting...")
            break

        # Reset speech_text and ocr_text is handled after processing

    # Clean up keyboard hooks before exiting
    keyboard.unhook_all()

# -------------------- Entry Point --------------------

if __name__ == "__main__":
    # Define global locks
    lock = threading.Lock()
    tts_lock = threading.Lock()
    
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logging.info("Program terminated by user.")
    except Exception as e:
        logging.error(f"An unexpected error occurred: {e}")
    finally:
        # Ensure all keyboard hooks are removed before exiting
        keyboard.unhook_all() : Make changes so that it doesn't expect any specific kind of inpt, if user says something then it should give it to the llm in addition to the ocr text but it does not recieve ANY AUDIO then it should only proceed with ocr text and try to find answer of that text