# This file is part of Senator Sperma
# Licensed under the GPL-3.0 License
# Copyright (C) 2026 Salmonik @ salmonik@tuta.io
# https://codeberg.org/salmonik/senator-sperma

import discord
from discord.ext import commands
import logging
import os
from cogs.utils.llm import llm_completion

logger = logging.getLogger(__name__)

class Jarvis(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.keyword = os.getenv("ASSISTANT_KEYWORD", "jarvis").lower()

    @commands.Cog.listener()
    async def on_message(self, message):
        if message.author == self.bot.user:
            return

        logger.info(message.content)

        if message.mentions:
            if message.mentions[0].id == self.bot.user.id:
                await self.assistant(message)

        elif self.keyword in message.content.lower():
            await self.assistant(message)

    @staticmethod
    async def assistant(message):
        chain = await Jarvis._build_message_chain(message)
        messages = Jarvis._build_messages(message, chain)

        extra = None
        api_url = os.getenv("LLM_API_URL", "https://openrouter.ai/api/v1/chat/completions")
        if "openrouter" in api_url.lower():
            extra = {"plugins": [{"id": "response-healing"}]}

        response_text = await llm_completion(messages, extra=extra)
        if response_text is None:
            await message.reply("Przepraszam, nie mogę teraz odpowiedzieć. Spróbuj ponownie później.")
            return

        await message.reply(response_text)

    @staticmethod
    async def _build_message_chain(message):
        chain = []
        current = message
        max_depth = 5

        while current and len(chain) < max_depth:
            chain.append(current)
            if not current.reference:
                break
            try:
                if current.reference.resolved:
                    current = current.reference.resolved
                else:
                    current = await current.channel.fetch_message(current.reference.message_id)
            except discord.NotFound:
                break

        return list(reversed(chain))

    @staticmethod
    def _build_messages(message, chain=None):
        messages = [{
            "role": "system",
            "content": "Jesteś Senator Sperma Armostrong, także znany jako Jarvis."
                       "Twoim zadaniem jest asystowanie użytkownika"
                       "zachowuj sie jak Jarvis, bądź osobistym asystentem użytkownika, nie używaj emotikon"
                       "nie pytaj sie użytkownika czy chce abyś coś zrobił, użytkownik ci sam powie abyś coś wykonał"
                       "Nie pozostawiaj wiadomości pustej oraz nie dawaj swojej wiadomości w pole myślenia"
        }]

        if chain is None:
            chain = [message]

        for msg in chain:
            messages.append({
                "role": "user",
                "content": f"USER: {msg.author.display_name}\nMESSAGE: {msg.content}"
            })

        return messages



def setup(bot):
    bot.add_cog(Jarvis(bot))