#!/usr/bin/env python3
"""
Local Screen-Aware AI Assistant
Optimized for RTX 3060 8GB + 16GB RAM
Hotkey: Shift + /
"""

import base64
import io
import threading
from typing import Set

import mss
from PIL import Image
import ollama
from pynput import keyboard

# ====================== CONFIG ======================
MODEL = "qwen2.5vl:7b"
HOTKEY = {keyboard.Key.shift, keyboard.KeyCode.from_char('/')}
CAPTURE_ALL_MONITORS = False
MAX_IMAGE_SIZE = (1280, 720)
JPEG_QUALITY = 85
# ====================================================

current_keys: Set = set()
is_processing = False


def capture_screen() -> str:
    """Capture screen and return base64-encoded JPEG."""
    with mss.mss() as sct:
        monitor = sct.monitors[0] if CAPTURE_ALL_MONITORS else sct.monitors[1]
        shot = sct.grab(monitor)
        img = Image.frombytes("RGB", shot.size, shot.bgra, "raw", "BGRX")

        if MAX_IMAGE_SIZE:
            img.thumbnail(MAX_IMAGE_SIZE, Image.Resampling.LANCZOS)

        buffer = io.BytesIO()
        img.save(buffer, format="JPEG", quality=JPEG_QUALITY)
        return base64.b64encode(buffer.getvalue()).decode("utf-8")


def ask_vision(question: str, image_b64: str) -> str:
    """Send question + screenshot to the local vision model."""
    try:
        response = ollama.chat(
            model=MODEL,
            messages=[
                {
                    "role": "user",
                    "content": question,
                    "images": [image_b64],
                }
            ],
            options={
                "temperature": 0.2,
                "num_ctx": 4096,
            }
        )
        return response["message"]["content"].strip()
    except Exception as e:
        return f"Error: {e}"


def process_question(question: str):
    """Full pipeline: capture → ask model → show answer."""
    global is_processing
    if is_processing:
        print("Already processing, please wait...")
        return

    is_processing = True
    print("\n📸 Capturing screen...")

    try:
        image_b64 = capture_screen()
        print(f"🤖 Thinking with {MODEL}...")
        answer = ask_vision(question, image_b64)

        print("\n" + "=" * 70)
        print(answer)
        print("=" * 70 + "\n")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        is_processing = False


def on_press(key):
    try:
        current_keys.add(key)
        if HOTKEY.issubset(current_keys) and not is_processing:
            print("\n🎤 Hotkey pressed!")
            print("Press Enter to answer questions on screen, or type a custom question:")
            question = input("> ").strip()
            if not question:
                question = (
                    "Look carefully at the screen. "
                    "Find any questions, problems, exercises, or homework tasks that are visible. "
                    "Answer them clearly and completely. "
                    "If there are multiple questions, answer all of them one by one. "
                    "Show your reasoning when helpful."
                )
            threading.Thread(target=process_question, args=(question,), daemon=True).start()
    except Exception:
        pass


def on_release(key):
    try:
        current_keys.discard(key)
    except Exception:
        pass


def main():
    print("=" * 70)
    print("Local Screen AI Assistant")
    print(f"Model     : {MODEL}")
    print(f"Hotkey    : Shift + /")
    print("Press Shift+/ then just press Enter → it will answer questions on screen")
    print("Type 'q' + Enter to quit.")
    print("=" * 70)

    listener = keyboard.Listener(on_press=on_press, on_release=on_release)
    listener.start()

    try:
        while True:
            user_input = input("\nQuestion (or 'q' to quit): ").strip()
            if user_input.lower() in ("q", "quit", "exit"):
                break
            if user_input:
                process_question(user_input)
    except KeyboardInterrupt:
        pass
    finally:
        listener.stop()
        print("\nGoodbye!")


if __name__ == "__main__":
    main()