# 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):
        print(message.content)
        if message.author == self.bot.user:
            return

        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"
        }
        
        messages = Jarvis._build_messages(message)
        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
    def _build_messages(message):
        messages = [{
            "role": "system",
            "content": "your name is Senator Sperma.\n"
                       "your task is to tell the program what to do and how to respond after receiving a message.\n"
                       "message field has to be filled out in polish.\n"
                       "act like jarvis (only behavior wise), be a personal assistant to user, do not use emojis\n"
                       "do not ask the user if they want to do an action. the user will tell you that."
        }]
        
        messages.append({
            "role": "user",
            "content": f"{message.author.display_name}: {message.content}"
            
        })
        
        if Jarvis._is_reply(message):
            messages.extend(Jarvis._build_reply_messages(message))
            
        return messages

    @staticmethod
    def _is_reply(message):
        return message.reference and (message.type == discord.MessageType.reply or message.reference.resolved)

    @staticmethod
    def _build_reply_messages(message):
        original_msg = message.reference.resolved
        return [
            {
                "role": "user",
                "content": f"{original_msg.author.display_name}: {original_msg.content}"
            },
            {
                "role": "user",
                "content": f"{message.author.display_name}: {message.content}"
            }
        ]

    @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"])

    @staticmethod
    async def get_original(message):
        if message.type == discord.MessageType.reply:
            # The original message is already resolved (cached)
            reply = message.reference.resolved
        elif message.reference:
            # If not cached, fetch it from the API
            try:
                reply = await message.channel.fetch_message(message.reference.message_id)
            except discord.NotFound:
                reply = None
        else:
            reply = None

        if reply:
            reply_content = reply.content
        else:
            reply_content = None

        if reply_content:
            return f"{message.content} [which is a reply to] {reply_content}"
        else:
            return message.content

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