import weechat
import requests
import json
from requests.auth import HTTPBasicAuth

"""
Ollama Bot for WeeChat (Remote Version with Authentication)

This script automatically responds to mentions in channels and private messages using an Ollama LLM hosted remotely.

Features:
- Responds to mentions in channels.
- Can respond to private messages if enabled.
- Allows manual queries using the /ollama command.
- Configurable via WeeChat /set commands.

Usage:
- To ask a question manually:
  /ollama What is Python?

- To enable or disable automatic responses in channels:
  /set plugins.var.python.ollama.highlight_response on  # Enable responses in channels
  /set plugins.var.python.ollama.highlight_response off # Disable responses in channels

- To enable or disable automatic responses in private messages:
  /set plugins.var.python.ollama.pm_response on  # Enable PM responses
  /set plugins.var.python.ollama.pm_response off # Disable PM responses

Dependencies:
- Requires an Ollama server running at https://ollama.priet.us/api/generate with authentication.
"""

# Script metadata
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 = "https://ollama.priet.us/api/generate"
OLLAMA_USER = "nginx-user"
OLLAMA_PASS = "wasamasa123"  # Replace with the actual password

# Register the script
weechat.register(SCRIPT_NAME, SCRIPT_AUTHOR, SCRIPT_VERSION, SCRIPT_LICENSE, SCRIPT_DESC, "", "")

# Script configuration in Weechat
def setup_config():
    if not weechat.config_is_set_plugin("highlight_response"):
        weechat.config_set_plugin("highlight_response", "on")  # Enable auto-responses by default
    if not weechat.config_is_set_plugin("pm_response"):
        weechat.config_set_plugin("pm_response", "off")  # Disable PM responses by default
setup_config()

def ask_ollama(message):
    """Send a query to Ollama and return the complete response."""
    try:
        data = {"model": "gemma:2b", "prompt": message, "stream": False}
        headers = {"Content-Type": "application/json", "User-Agent": "WeeChat-OllamaBot/1.0"}

        weechat.prnt("", f"[DEBUG] Sending request to Ollama: {OLLAMA_API_URL} with prompt: {message}")

        response = requests.post(
            OLLAMA_API_URL,
            json=data,
            headers=headers,
            auth=HTTPBasicAuth(OLLAMA_USER, OLLAMA_PASS),
            verify=False  # Change to True if you have a valid certificate
        )

        if response.status_code == 401:
            return "Authentication Error: Check username/password."
        elif response.status_code == 403:
            return "Permission Denied: Verify API access."
        elif response.status_code != 200:
            return f"HTTP Error {response.status_code}: {response.text}"

        response_json = response.json()
        return response_json.get("response", "No response received from Ollama.")
    
    except requests.exceptions.RequestException as e:
        return f"Error connecting to Ollama: {str(e)}"

def command_ollama(data, buffer, args):
    """Command /ollama to manually ask Ollama a question."""
    weechat.prnt("", f"[DEBUG] /ollama command received with args: {args}")
    
    if not args:
        weechat.prnt(buffer, "[Ollama] Usage: /ollama <question>")
        return weechat.WEECHAT_RC_OK

    response = ask_ollama(args)
    weechat.prnt(buffer, f"[Ollama] {response}")
    return weechat.WEECHAT_RC_OK

def message_callback(data, buffer, date, tags, displayed, highlight, prefix, message):
    """Detect mentions in channels or private messages and respond automatically with Ollama."""
    
    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", "")  # Get the current IRC username
    is_mentioned = f"@{username.lower()}" in message.lower()  # Ensure @username is explicitly mentioned

    # Ignore private messages if pm_response is off
    if is_private and weechat.config_get_plugin("pm_response") == "off":
        return weechat.WEECHAT_RC_OK

    # Only respond in private messages if it's a direct question
    if is_private and not message.strip().endswith("?"):
        return weechat.WEECHAT_RC_OK

    # Only respond in channels if explicitly mentioned or highlighted
    if not is_private and not is_mentioned and not int(highlight):
        return weechat.WEECHAT_RC_OK

    response = ask_ollama(message)
    
    if is_private:
        weechat.command(buffer, f"/msg {prefix} {response}")  # Reply to private message
    else:
        weechat.command(buffer, f"/say {response}")  # Reply in the channel

    return weechat.WEECHAT_RC_OK

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

# Register configuration with /set
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", "")

# Register commands and hooks
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", "")