import weechat
import json

SCRIPT_NAME = "ollama"
SCRIPT_AUTHOR = "teraflops"
SCRIPT_VERSION = "2.1"
SCRIPT_LICENSE = "MIT"
SCRIPT_DESC = "Automatically responds to mentions using Ollama and allows manual queries, including PMs"
OLLAMA_API_URL = "http://localhost:11434/api/generate"

weechat.register(SCRIPT_NAME, SCRIPT_AUTHOR, SCRIPT_VERSION, SCRIPT_LICENSE, SCRIPT_DESC, "", "")

def setup_config():
    if not weechat.config_is_set_plugin("highlight_response"):
        weechat.config_set_plugin("highlight_response", "on")
    if not weechat.config_is_set_plugin("pm_response"):
        weechat.config_set_plugin("pm_response", "off")
setup_config()

def ask_ollama_async(prompt, buffer, prefix=""):
    payload = json.dumps({
        "model": "gemma2:2b",
        "prompt": prompt,
        "stream": False
    })

    curl_path = "/usr/bin/curl"

    escaped_payload = payload.replace('"', '\\"')
    cmd = (
        f"{curl_path} -s -X POST "
        "-H 'Content-Type: application/json' "
        f'--data \"{escaped_payload}\" {OLLAMA_API_URL}'
    )

    user_data = f"{buffer}||{prefix}"
    weechat.hook_process_hashtable(
        cmd,
        {"timeout": "10000"},
        10000,
        "ollama_response_callback",
        user_data
    )

def ollama_response_callback(data, command, return_code, out, err):
    buffer, prefix = data.split("||", 1)
    weechat.prnt(buffer, f"[Ollama DEBUG] Command: {command}")
    weechat.prnt(buffer, f"[Ollama DEBUG] Raw output: {out}")
    weechat.prnt(buffer, f"[Ollama DEBUG] Error output: {err}")

    if return_code is None or return_code == weechat.WEECHAT_HOOK_PROCESS_ERROR:
        response = f"[Ollama] Error executing curl. Output was:\n{out}"
    elif out.strip() == "":
        response = "[Ollama] Empty response from Ollama."
    else:
        try:
            parsed = json.loads(out)
            response = parsed.get("response", "[Ollama] No 'response' field in reply.")
        except Exception as e:
            weechat.prnt(buffer, f"[Ollama DEBUG] JSON decode error: {e}")
            response = f"[Ollama] Failed to parse JSON: {e}\nRaw output: {out}"

    if prefix:
        weechat.command(buffer, f"/msg {prefix} {response}")
    else:
        weechat.command(buffer, f"/say {response}")

    return weechat.WEECHAT_RC_OK

def command_ollama(data, buffer, args):
    if not args:
        weechat.prnt(buffer, "[Ollama] Usage: /ollama <question>")
        return weechat.WEECHAT_RC_OK

    ask_ollama_async(args, buffer)
    return weechat.WEECHAT_RC_OK

def message_callback(data, buffer, date, tags, displayed, highlight, prefix, message):
    if weechat.config_get_plugin("highlight_response") == "off":
        return weechat.WEECHAT_RC_OK

    buffer_type = weechat.buffer_get_string(buffer, "localvar_type")
    is_private = buffer_type == "private"
    username = weechat.info_get("irc_nick", "")
    is_mentioned = f"@{username.lower()}" in message.lower()

    if is_private and weechat.config_get_plugin("pm_response") == "off":
        return weechat.WEECHAT_RC_OK

    if is_private and not message.strip().endswith("?"):
        return weechat.WEECHAT_RC_OK

    if not is_private and not is_mentioned and not int(highlight):
        return weechat.WEECHAT_RC_OK

    ask_ollama_async(message, buffer, prefix if is_private else "")
    return weechat.WEECHAT_RC_OK

def config_callback(data, option, value):
    weechat.prnt("", f"[Ollama] Configuration changed: {option} = {value}")
    return weechat.WEECHAT_RC_OK

weechat.config_set_desc_plugin("highlight_response", "Automatically respond to mentions in channels (on/off)")
weechat.config_set_desc_plugin("pm_response", "Automatically respond to private messages (on/off)")
weechat.hook_config("plugins.var.python.ollama.highlight_response", "config_callback", "")
weechat.hook_config("plugins.var.python.ollama.pm_response", "config_callback", "")
weechat.hook_command("ollama", "Ask something to Ollama", "<question>", "Example: /ollama What is Python?", "", "command_ollama", "")
weechat.hook_print("", "notify_highlight", "", 1, "message_callback", "")
weechat.hook_print("", "notify_message", "", 1, "message_callback", "")
weechat.hook_print("", "notify_private", "", 1, "message_callback", "")