# curl -fsSL https://ollama.com/install.sh | sh
user@pc:~$ ollama run llama2
user@pc:~$ A circle has the radius 5. What is the area of the circle?
The area of a circle can be calculated using the formula: Area =πr², where r is the radius of the circle. In this case, the radius of the circle is 5, so the area of the circle is: Area = π(5)² = 3.14(5²) = 78.5

So the area of the circle is 78.5 square units.

"""
main.py

This script uses the pynput library to listen for global hotkeys and perform actions on selected text.
It interacts with a local server running Ollama's LLaMA model to correct and improve the text in the clipboard.
The script can fix selected text or the current line in an editor.

Dependencies:
- pynput
- pyperclip
- httpx
- string (Template)
- time
- logging

Hotkeys:
- F9: Fix the selected text in the clipboard
- F10: Fix the current line in the clipboard

Author: Máximo Núñez Alarcón
Date: May 21, 2024
"""

from pynput import keyboard
from pynput.keyboard import Key, Controller
import pyperclip
import time
import logging
from string import Template
import httpx

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Keyboard controller for simulating key presses
controller = Controller()

# Ollama API endpoint and configuration
OLLAMA_ENDPOINT = "http://localhost:11434/api/generate"
OLLAMA_CONFIG = {
    "model": "llama3",
    "keep_alive": "5m",
    "stream": False
}

# Template for the prompt sent to the Ollama API
PROMPT_TEMPLATE = Template(
    """Improve the text in this Markdown contentTo solve an equation of second degree, you must use the following formula: by fixing all typos, spelling mistakes, and grammar errors. Additionally, make sure the text is clear and well-structured, but preserve all new line characters and the original meaning:

    $text

    Return only the corrected text without any preamble or additional commentary.
    """
)

def fix_current_line():
    """
    Selects the current line in the editor and sends it for correction.
    """

    # Select the current line (simulates Ctrl+Shift+Left Arrow): Key.ctrl_l, Key.shift, Key.left
    # However, I am using Arch (Linux):
    controller.press(Key.ctrl_l)
    controller.press(Key.end)
    controller.release(Key.end)
    controller.release(Key.ctrl_l)
    
    # Fix the selected text
    fix_selection()

def fix_text(text):
    """
    Sends the given text to the Ollama API for correction.

    Args:
        text (str): The text to be corrected.

    Returns:
        str: The corrected text if the API call is successful, otherwise None.
    """
    prompt = PROMPT_TEMPLATE.substitute(text=text)
    try:
        response = httpx.post(
            OLLAMA_ENDPOINT,
            json={"prompt": prompt, **OLLAMA_CONFIG},
            headers={"Content-Type": "application/json"},
            timeout=30  # Increased timeout duration
        )
        if response.status_code == 200:
            response_data = response.json()
            if "response" in response_data:
                return response_data["response"].strip()
            else:
                logger.error("Unexpected response format: 'response' key not found.")
        else:
            logger.error(f"HTTP error {response.status_code}: {response.text}")
    except Exception as e:
        logger.error(f"An error occurred: {str(e)}")
    return None

def fix_selection():
    """
    Copies the selected text to the clipboard, sends it for correction, 
    and pastes the corrected text back into the clipboard.
    """
    # 1. Copy the current selection to the clipboard (simulates Ctrl+C)
    with controller.pressed(Key.ctrl_l):
        controller.tap('c')
    
    # Wait for clipboard to update
    time.sleep(0.1)

    # 2. Get the text from clipboard
    text = pyperclip.paste()
    
    if not text:
        logger.warning("No text found in clipboard.")
        return
    
    # 3. Fix the text using the Ollama API
    fixed_text = fix_text(text)
    if fixed_text:
        # 4.  Copy the fixed text back to the clipboard
        pyperclip.copy(fixed_text)
        # Wait for clipboard to update
        time.sleep(0.1)

        # 5. Paste the fixed text (simulates Ctrl+V)
        with controller.pressed(Key.ctrl_l):
            controller.tap('v')
    else:
        logger.warning("Failed to get fixed text.")

def on_f9():
    """
    Callback function triggered by pressing F9 to fix the selected text.
    """
    fix_selection()

def on_f10():
    """
    Callback function triggered by pressing F10 to fix the current line.
    """
    fix_current_line()

# Setting up global hotkeys for F9 and F10 using pynput.
with keyboard.GlobalHotKeys({
        '': on_f9,
        '': on_f10}) as h:
    h.join()