import discord
import requests
import json
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Replace with your bot token
BOT_TOKEN = os.getenv('DISCORD_TOKEN')

# Ollama configuration
OLLAMA_API_URL = 'http://192.168.1.159:11434/api/generate'  # Adjust if your Ollama setup is different

# Set up the Discord client
intents = discord.Intents.default()
intents.messages = True
intents.message_content = True

client = discord.Client(intents=intents)


# Function to query Ollama
def query_ollama(prompt):
    payload = {
        "prompt": prompt,
        "model": "nollama/mythomax-l2-13b:Q4_K_M"  # Replace with your Ollama model
    }
    try:
        response = requests.post(OLLAMA_API_URL, json=payload, stream=True)
        if response.status_code == 200:
            collected_response = ""
            # Stream and parse each line of JSON from the response
            for line in response.iter_lines(decode_unicode=True):
                if line.strip():  # Skip empty lines
                    try:
                        data = json.loads(line)  # Parse each line as JSON
                        collected_response += data.get("response", "")
                        if data.get("done", False):
                            break
                    except json.JSONDecodeError as e:
                        print(f"Error decoding JSON line: {line}, Error: {e}")
            return collected_response.strip() or "No response from model."
        else:
            return f"Error: {response.status_code} - {response.text}"
    except requests.RequestException as e:
        return f"Error connecting to Ollama: {str(e)}"


# Event for when the bot is ready
@client.event
async def on_ready():
    print(f'We have logged in as {client.user}')


# Event for when a message is sent
@client.event
async def on_message(message):
    # Ignore the bot's own messages
    if message.author == client.user:
        return

    # Respond to all messages except those in DMs
    if not isinstance(message.channel, discord.DMChannel):
        response = query_ollama(message.content.strip())
        await message.channel.send(response)

# Run the bot
client.run(BOT_TOKEN)