# This is the core hotkey application. It has no GUI and runs in the background.
import json
import sqlite3
import threading
import time
from string import Template
from functools import partial
import os
import sys
import logging

# --- PyInstaller Helper ---
def resource_path(relative_path):
    """ Get absolute path to resource, works for dev and for PyInstaller """
    try:
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.abspath(".")
    return os.path.join(base_path, relative_path)

# --- Immediate Logging Setup to the Desktop ---
log_file_path = os.path.join(os.path.expanduser("~"), "Desktop", "tangent_log.txt")
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(log_file_path, mode='w'),
        logging.StreamHandler(sys.stdout)
    ]
)

logging.info("--- SCRIPT STARTED ---")

import google.generativeai as genai
from pynput import keyboard
from pynput.keyboard import Key, Controller
import pyperclip
from PIL import Image, ImageDraw
from pystray import Icon as icon, Menu as menu, MenuItem as item

# --- GLOBAL VARIABLES ---
CONFIG = {}
CACHE_CONN = None
HOTKEY_LISTENER = None
SCRIPT_RUNNING = threading.Event()
CONTROLLER = Controller()
TRAY_ICON = None # Global reference to the icon

# --- Core Logic Functions ---
def load_config():
    """Loads the configuration from config.json."""
    global CONFIG
    try:
        with open(resource_path("config.json"), "r") as f:
            CONFIG = json.load(f)
        logging.info("Configuration loaded successfully.")
    except Exception as e:
        logging.error(f"FATAL: Error loading config.json: {e}")
        time.sleep(10)
        exit()

def setup_cache():
    """Initializes the SQLite database for caching results."""
    global CACHE_CONN
    try:
        db_path = os.path.join(os.path.expanduser("~"), "tangent_cache.db")
        CACHE_CONN = sqlite3.connect(db_path, check_same_thread=False)
        CACHE_CONN.cursor().execute("CREATE TABLE IF NOT EXISTS corrections (original_text TEXT PRIMARY KEY, corrected_text TEXT)")
        CACHE_CONN.commit()
        logging.info("Cache database initialized.")
    except Exception as e:
        logging.error(f"FATAL: Failed to set up cache: {e}")
        time.sleep(10)
        exit()

def get_from_cache(text):
    cursor = CACHE_CONN.cursor()
    res = cursor.execute("SELECT corrected_text FROM corrections WHERE original_text = ?", (text,)).fetchone()
    if res:
        logging.info("Result found in cache.")
        return res[0]
    return None

def add_to_cache(original, corrected):
    cursor = CACHE_CONN.cursor()
    cursor.execute("INSERT OR REPLACE INTO corrections (original_text, corrected_text) VALUES (?, ?)", (original, corrected))
    CACHE_CONN.commit()

def fix_text(text: str, prompt_template: str, mode_name: str) -> str | None:
    if not text.strip(): return None
    cached = get_from_cache(text)
    if cached: return cached
    
    logging.info(f"Calling Gemini API for '{mode_name}'...")
    prompt = Template(prompt_template).substitute(text=text)
    
    try:
        model = genai.GenerativeModel('gemini-2.0-flash-001')
        response = model.generate_content(prompt)
        corrected = response.text.strip()
        
        if 'code' in mode_name.lower():
            if corrected.startswith("```"):
                lines = corrected.split('\n')
                if len(lines) > 2:
                    corrected = '\n'.join(lines[1:-1])
                else:
                    corrected = corrected.strip('`')

        add_to_cache(text, corrected)
        return corrected
    except Exception as e:
        logging.error(f"An error occurred with the Gemini API: {e}")
        return None

def on_hotkey_activated(prompt_template: str, mode_name: str):
    if not SCRIPT_RUNNING.is_set():
        logging.warning("Script is paused. Ignoring hotkey.")
        return
    
    original_clipboard = pyperclip.paste()
    
    selected_text = ""
    for _ in range(3):
        pyperclip.copy('') 
        with CONTROLLER.pressed(Key.ctrl):
            CONTROLLER.tap('c')
        time.sleep(0.15)
        selected_text = pyperclip.paste()
        if selected_text:
            break 
    
    if not selected_text:
        logging.warning("No text selected.")
        pyperclip.copy(original_clipboard) 
        return

    logging.info(f"Processing '{mode_name}'...")
    corrected_text = fix_text(selected_text, prompt_template, mode_name)
    
    if corrected_text:
        pyperclip.copy(corrected_text)
        with CONTROLLER.pressed(Key.ctrl):
            CONTROLLER.tap('v')
        time.sleep(0.1)
        
    pyperclip.copy(original_clipboard)

def correct_hotkey_format(hotkey_str: str) -> str:
    if '<' in hotkey_str and '>' in hotkey_str and '+' not in hotkey_str:
        parts = hotkey_str.split('>')
        if len(parts) == 2 and len(parts[1]) > 0:
            return f"{parts[0]}>+{parts[1]}"
    return hotkey_str

def setup_hotkey_listener():
    global HOTKEY_LISTENER
    hotkeys = {}
    try:
        for mode in CONFIG['modes']:
            hotkey_str = correct_hotkey_format(mode['hotkey'])
            callback = partial(on_hotkey_activated, prompt_template=mode['prompt'], mode_name=mode['name'])
            hotkeys[hotkey_str] = callback
            logging.info(f"Registered hotkey '{hotkey_str}' for mode '{mode['name']}'.")

        HOTKEY_LISTENER = keyboard.GlobalHotKeys(hotkeys)
        HOTKEY_LISTENER.start()
        logging.info("Hotkey listener started successfully.")
    except Exception as e:
        logging.error(f"FATAL: Failed to set up hotkeys: {e}")
        time.sleep(10)
        exit()

def create_image():
    try:
        return Image.open(resource_path("logo.png")) 
    except Exception as e:
        logging.warning(f"logo.png not found ({e}). Creating a default icon.")
        width, height = 64, 64
        image = Image.new("RGB", (width, height), "black")
        dc = Image.Draw(image)
        dc.text((22, 20), "T", fill="white")
        return image

def toggle_script(icon, item):
    if SCRIPT_RUNNING.is_set(): SCRIPT_RUNNING.clear(); logging.info("Script paused.")
    else: SCRIPT_RUNNING.set(); logging.info("Script resumed.")

def quit_script():
    logging.info("Exit requested. Shutting down.")
    if HOTKEY_LISTENER and HOTKEY_LISTENER.is_alive(): HOTKEY_LISTENER.stop()
    if CACHE_CONN: CACHE_CONN.close()
    if TRAY_ICON: TRAY_ICON.stop()
    SCRIPT_RUNNING.clear() # This will stop the main loop

def setup_and_run_tray_icon():
    """Creates and runs the icon in a background thread."""
    global TRAY_ICON
    image = create_image()
    shortcuts_menu_items = [item(f"{m['name']}: {m.get('hotkey', 'N/A')}", None, enabled=False) for m in CONFIG.get("modes", [])]
    
    tray_menu = menu(
        item('Tangent', None, enabled=False),
        menu.SEPARATOR,
        item('Shortcuts', menu(*shortcuts_menu_items)),
        menu.SEPARATOR,
        item('Status: Running', lambda item: SCRIPT_RUNNING.is_set(), checked=lambda item: SCRIPT_RUNNING.is_set()),
        item('Pause / Resume', toggle_script),
        item('Exit', quit_script)
    )
    
    TRAY_ICON = icon('Tangent', image, 'Tangent', tray_menu)
    logging.info("Starting system tray icon in a separate thread.")
    TRAY_ICON.run_detached()

# --- Main Execution ---
if __name__ == "__main__":
    try:
        SCRIPT_RUNNING.set()
        load_config()
        api_key = CONFIG.get("api_key")
        if not api_key or "PASTE_YOUR" in api_key:
            logging.error("FATAL: Gemini API key not found in config.json.")
            time.sleep(10)
            exit()
        
        genai.configure(api_key=api_key)
        setup_cache()
        setup_hotkey_listener()
        
        # This now starts the icon in the background
        setup_and_run_tray_icon()
        
        # This loop keeps the main script alive until quit_script is called
        while SCRIPT_RUNNING.is_set():
            time.sleep(1)

    except Exception as e:
        logging.critical(f"A fatal, unhandled error occurred in main execution block: {e}", exc_info=True)
        time.sleep(10)

    logging.info("--- Tangent Application Exited ---")

