from pydantic_ai import BinaryImage
from quickmachotkey import mask, quickHotKey
from quickmachotkey.constants import (
    cmdKey,
    controlKey,
    kVK_ANSI_A,
    kVK_ANSI_H,
    optionKey,
    shiftKey,
)
from rumps import Timer

from app.modules.ai import AgentAnswer, agent, model_high, model_low
from app.modules.screenshot import take_screenshot
from app.ui.items import display_text
from app.ui.status import update_statusbar

timer = Timer(None, 0)


def process_task(model):
    timer.stop()
    import app.ui.handlers as handlers

    response = agent.run_sync(
        [BinaryImage(take_screenshot(), media_type="image/png")], model=model
    )
    output: AgentAnswer = response.output

    handlers.long_answer = "\n".join(output.long_answer or [])
    handlers.solution_thinking = output.solution_thinking

    display_text(output.long_answer)
    update_statusbar(output.answer_variant or "?")


@quickHotKey(
    virtualKey=kVK_ANSI_A,
    modifierMask=mask(cmdKey, controlKey, optionKey, shiftKey),
)
def low_handler() -> None:
    update_statusbar("...")
    timer.set_callback(lambda _: process_task(model_low))
    timer.start()


@quickHotKey(
    virtualKey=kVK_ANSI_H,
    modifierMask=mask(cmdKey, controlKey, optionKey, shiftKey),
)
def high_handler() -> None:
    update_statusbar("...")
    timer.set_callback(lambda _: process_task(model_high))
    timer.start()