# 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 re
import discord
from discord.ext import commands
import requests
import json
import logging
import shutil
from dotenv import load_dotenv
import os
from tempfile import NamedTemporaryFile
from cogs.downloader import Downloader

logger = logging.getLogger(__name__)
load_dotenv()

class Jarvis(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @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.jarvis(message)

        elif "jarvis" in message.content.lower():
            await self.jarvis(message)

    @staticmethod
    async def jarvis(message):
        headers = {
            "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
            "Content-Type": "application/json"
        }
        
        chain = await Jarvis._build_message_chain(message)
        messages = Jarvis._build_messages(message, chain)
        payload = Jarvis._build_payload(messages)
        
        response_data = await Jarvis._make_api_request(payload, headers)
        response_message, action = Jarvis._parse_response(response_data)
        
        channel = message.channel
        await message.reply(response_message)
        
        if action == "download":
            await Jarvis._handle_download_action(message, channel)

    @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, po za asystowaniem użytkownika, jest mówienie programowi nadrzędnemu jak sie zachować"
                       "pole 'message' może być tylko wypełnione w języku Polskim."
                       "zachowuj sie jak Jarvis, bądź osobistym asystentem użytkownika, nie używaj emotikon"
                       "nie pytaj sie użytkownika czy chcesz abyś coś zrobił, użytkownik ci sam powie abyś coś wykonał"
        }]
        
        if chain is None:
            chain = [message]
        
        for msg in chain:
            messages.append({
                "role": "user",
                "content": f"{msg.author.display_name}: {msg.content}"
            })
            
        return messages

    @staticmethod
    def _build_payload(messages):
        return {
            "model": "arcee-ai/trinity-large-preview:free",
            "messages": messages,
            "response_format": Jarvis._get_response_schema(),
            "plugins": [{"id": "response-healing"}]
        }

    @staticmethod
    def _get_response_schema():
        return {
            "type": "json_schema",
            "json_schema": {
                "name": "response",
                "strict": True,
                "schema": {
                    "type": "object",
                    "properties": {
                        "message": {
                            "type": "string",
                            "description": "the message that user will receive",
                        },
                        "action": {
                            "type": "string",
                            "description": "the action to do, either download, compress or none"
                        }
                    },
                    "required": ["message", "action"],
                    "additionalProperties": False,
                }
            }
        }

    @staticmethod
    async def _make_api_request(payload, headers):
        url = "https://openrouter.ai/api/v1/chat/completions"
        response = requests.post(url=url, headers=headers, json=payload)
        response_data = response.json()
        logger.info(response_data)
        return response_data

    @staticmethod
    def _parse_response(response_data):
        content = response_data['choices'][0]['message']['content']
        inner_json_str = content.strip('```json\n').strip('```').strip()
        inner_data = json.loads(inner_json_str)
        
        response_message = inner_data['message']
        action = inner_data['action']
        
        logger.info(response_message)
        logger.info(action)
        
        return response_message, action

    @staticmethod
    async def _handle_download_action(message, channel):
        url_match = Jarvis._extract_url(message.content)
        if not url_match:
            await channel.send("Nie znaleziono poprawnego linku w odpowiedzi.")
            return

        response = await Downloader.download_standalone(url=url_match)
        await Jarvis._process_download_response(response, channel)

    @staticmethod
    def _extract_url(content):
        match = re.search(r'https://(?:[A-Za-z0-9-]+\.)*[A-Za-z0-9-]{1,63}\.[A-Za-z]{2,63}/\S*', content)
        return match.group(0) if match else None

    @staticmethod
    async def _process_download_response(response, channel):
        if not isinstance(response, dict):
            return
            
        if response["status"] == "success":
            await Jarvis._handle_successful_download(response, channel)
        else:
            await Jarvis._handle_failed_download(response)

    @staticmethod
    async def _handle_successful_download(response, channel):
        if response["type"] == "url":
            await channel.send(f"pobioron i wysyłon: {response['content']}")
        elif response["type"] == "file":
            file_data = response["content"]
            await channel.send("wysyłon:", file=discord.File(file_data))
            shutil.rmtree(response["temp_dir"])

    @staticmethod
    async def _handle_failed_download(response):
        logger.error(f"Error: {response['content']}")
        if 'details' in response:
            logger.error(f"Details: {response['details']}")
        shutil.rmtree(response["temp_dir"])

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