from time import sleep
import win32clipboard
from optimisewait import optimiseWait, set_autopath
import pyautogui
import time
import pywintypes
import base64
import io
from PIL import Image
import importlib.resources
import tempfile
import shutil
import webbrowser
import os
import pygetwindow as gw

def switch_to_chrome():
    """Switches focus to the Chrome window if available."""
    chrome_windows = [w for w in gw.getWindowsWithTitle("Google Chrome")]
    if chrome_windows:
        chrome_windows[0].activate()
    else:
        print("Chrome window not found.")

def set_image_path(llm: str, debug: bool = False):
    """Dynamically sets the image path for optimisewait based on package installation location."""
    copy_images_to_temp(llm, debug=debug)

def copy_images_to_temp(llm: str, debug: bool = False):
    """Copies the necessary image files to a temporary directory."""
    temp_dir = tempfile.gettempdir()
    image_path = os.path.join(temp_dir, 'talktollm_images', llm)
    os.makedirs(image_path, exist_ok=True)
    if debug:
        print(f"Temporary image directory: {image_path}")

    # Get the path to the original images directory within the package
    original_images_dir = importlib.resources.files('talktollm').joinpath('images')
    original_image_path = original_images_dir / llm
    if debug:
        print(f"Original image directory: {original_image_path}")
    # Copy each file from the original directory to the temporary directory
    for filename in os.listdir(original_image_path):
        source_file = os.path.join(original_image_path, filename)
        destination_file = os.path.join(image_path, filename)
        if not os.path.exists(destination_file):
            if debug:
                print(f"Copying {source_file} to {destination_file}")
            shutil.copy2(source_file, destination_file)
        elif debug:
            print(f"File already exists: {destination_file}")

    set_autopath(image_path)
    if debug:
        print(f"Autopath set to: {image_path}")

def set_clipboard(text: str, retries: int = 3, delay: float = 0.2):
    for i in range(retries):
        try:
            win32clipboard.OpenClipboard()
            win32clipboard.EmptyClipboard()
            try:
                win32clipboard.SetClipboardText(str(text))
            except Exception:
                # Fallback for Unicode characters
                win32clipboard.SetClipboardData(win32clipboard.CF_UNICODETEXT, str(text).encode('utf-16le'))
            win32clipboard.CloseClipboard()
            return  # Success
        except pywintypes.error as e:
            if e.winerror == 5:  # Access is denied
                print(f"Clipboard access denied. Retrying... (Attempt {i+1}/{retries})")
                time.sleep(delay)
            else:
                raise  # Re-raise other pywintypes errors
        except Exception as e:
            raise  # Re-raise other exceptions
    print(f"Failed to set clipboard after {retries} attempts.")

def talkto(llm: str, prompt: str, imagedata: list[str] | None = None, debug: bool = False, tabswitch: bool = True) -> str:
    """
    Interacts with a specified Large Language Model (LLM).

    Args:
        llm: The name of the LLM to interact with ('deepseek' or 'gemini').
        prompt: The text prompt to send to the LLM.
        imagedata: Optional list of base64 encoded image strings to send to the LLM.
        debug: Whether to enable debugging output.
        tabswitch: Whether to switch tabs after interacting with the LLM.

    Returns:
        The LLM's response as a string.
    """
    llm = llm.lower()
    set_image_path(llm, debug=debug)
    urls = {
        'deepseek': 'https://chat.deepseek.com/',
        'gemini': 'https://aistudio.google.com/prompts/new_chat'
    }

    webbrowser.open_new_tab(urls[llm])

    switch_to_chrome()  

    optimiseWait('message',clicks=2)

    # If there are images, paste each one
    if imagedata:
        for img in imagedata:
            set_clipboard_image(img)
            pyautogui.hotkey('ctrl', 'v')
            sleep(7)  # Ensure upload completes before pasting the next image

    set_clipboard(prompt)
    pyautogui.hotkey('ctrl', 'v')

    sleep(1)

    optimiseWait('run')
    
    if llm == 'gemini':
        optimiseWait('done',clicks=0)
        sleep(1.5)

    optimiseWait('copy')
    
    pyautogui.hotkey('ctrl', 'w')
    
    if tabswitch:
        pyautogui.hotkey('alt', 'tab')

    # Get LLM's response
    win32clipboard.OpenClipboard()
    response = win32clipboard.GetClipboardData()
    win32clipboard.CloseClipboard()

    return response

def set_clipboard_image(image_data: str, retries: int = 3, delay: float = 0.2):
    """Set image data to clipboard with retries"""
    for attempt in range(retries):
        try:
            binary_data = base64.b64decode(image_data.split(',')[1])
            image = Image.open(io.BytesIO(binary_data))

            output = io.BytesIO()
            image.convert("RGB").save(output, "BMP")
            data = output.getvalue()[14:]  # Remove bitmap header
            output.close()

            win32clipboard.OpenClipboard()
            win32clipboard.EmptyClipboard()
            win32clipboard.SetClipboardData(win32clipboard.CF_DIB, data)
            win32clipboard.CloseClipboard()
            return True
        except pywintypes.error as e:
            if e.winerror == 5:  # Access is denied
                print(f"Clipboard access denied. Retrying... (Attempt {attempt+1}/{retries})")
                time.sleep(delay)
            else:
                raise  # Re-raise other pywintypes errors
        except Exception as e:
            print(f"Error setting image to clipboard: {e}")
            return False
    return False

if __name__ == "__main__":
    print(talkto('gemini','Hi'))
