import os
import openai
import json

import nextcord
from nextcord.ext import commands


Enabled = True

# Get these from .env file too
SquogAILink = os.environ.get("AI_URL") ## AI_URL = https://
SquogAIToken = os.environ.get("AI_TOKEN") ## AI_TOKEN = ...

SquogAI = openai.OpenAI(api_key=SquogAIToken, base_url=SquogAILink)

SquogDataFile = open("./userinfo/userinfo.json", "r")
SquogData = json.loads(SquogDataFile.read())
SquogDataFile.close()

SquogUserPreset = {"id":  999, "messages":  [{"role": "system", "content": "You must act like a curious creature and never say you are a PC program."}]}

async def getResponse(SquogMessageList):
    SquogResponse = SquogAI.chat.completions.create(messages=SquogMessageList, model="gpt-4o", temperature=0.9)
    return json.loads(SquogResponse.model_dump_json())

async def find(listIterate: list, param: str, value: any):
    found = None
    for g in listIterate:
        if g[param] == value:
            found = g
    return found

async def flush():
    LocalSquogDataFile = open("./userinfo/userinfo.json", "w")
    LocalSquogDataFile.write(json.dumps(SquogData))
    LocalSquogDataFile.flush()
    LocalSquogDataFile.close()

class AI(commands.Cog):
    def __init__(self, bot: commands.Bot):
        self.client = bot

    @commands.Cog.listener()
    async def on_message(self, message: nextcord.Message):

        if not Enabled:
            return

        SquogMention = False

        for Mention in message.mentions:
            if Mention.id == self.client.user.id:
                SquogMention = True

        if not SquogMention:
            return

        message.channel.typing()

        SquogUser = await find(SquogData, "id", message.author.id)

        if not SquogUser:
            SquogUser = SquogUserPreset.copy()
            SquogUser["id"] = message.author.id
            SquogData.append(SquogUser)

        SquogUser["messages"].append({"role": "user", "content": message.content})

        try:
            SquogResponse = await getResponse(SquogUser["messages"])

            SquogUser["messages"].append(SquogResponse["choices"][0]["message"])
            await message.reply(SquogResponse["choices"][0]["message"]["content"])

            await flush()
        except:
            await message.reply("Failed to load a reply.")

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