from discord import app_commands
import requests
import discord
import asyncio
from discord.ext.commands import is_owner, Context, has_permissions
import os
from typing import Optional

intents = discord.Intents.all()
client = discord.Client(intents=intents)
tree = app_commands.CommandTree(client)

url = "https://huggingface.co/chat/conversation"

# load system prompt
try:
    with open("sysp.txt", "r") as f:
        sysp = f.read()
except:
    sysp = ""

# models
mixtral = "mistralai/Mixtral-8x7B-Instruct-v0.1"
llama = "meta-llama/Llama-2-70b-chat-hf"
openchat = "openchat/openchat-3.5-0106"
mistral = "mistralai/Mistral-7B-Instruct-v0.2"
hermes = "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO"
codellama = "codellama/CodeLlama-70b-Instruct-hf"

async def chat(conversationId, ip, cookie):
    s = requests.Session()
    url = f"https://huggingface.co/chat/conversation/{conversationId}"
    
    payload = {
        "inputs": ip,
        "id": cookie,
        "is_retry": False,
        "is_continue": False,
        "web_search": False,
        "files": []
    }
    
    headers = {
        "authority": "huggingface.co",
        "accept": "*/*",
        "accept-language": "en-US,en;q=0.9",
        "content-type": "application/json",
    }

    response = s.post(url, json=payload, headers=headers)
    return response
    s.close()

@client.tree.command(
    description="Starts the conversation with Lyre.",
    guild=discord.Object(id=1079318749116641320)
)
@app_commands.describe(
    channel='The channel to start the conversation in. Defaults to the channel it was called in.',
)
@has_permissions(manage_channels=True, manage_messages=True)
async def start(ctx: discord.Interaction, channel: Optional[discord.AppCommandOptionType.channel]=None):
    s = requests.Session()
    payload = {
        "model": mixtral,
        "preprompt": sysp
    }
    headers = {
        "authority": "huggingface.co",
        "accept": "*/*",
        "accept-language": "en-US,en;q=0.9",
        "content-type": "application/json"
    }
    if channel == None:
        chping = f"<#{ctx.channel_id}>"
        with open(str(ctx.channel_id), "w") as f:
            f.write(s.cookies.get("hf-chat") + "\n" + response.json.loads(response.text)['conversationId'])
    else:
        chping = f"<#{channel.id}>"
        with open(str(ctx.channel_id), "w") as f:
            f.write(s.cookies.get("hf-chat") + "\n" + response.json.loads(response.text)['conversationId'])
    ctx.response.send_message(f"Started a conversation in {chping}")

@start.error
async def start_error(error, ctx):
    if isinstance(error, commands.MissingPermissions):
        ctx.response.send_message(f"Missing permissions in {chping}")
        
@client.tree.command(
    description="Ends the conversation with Lyre.",
    guild=discord.Object(id=1079318749116641320)
)
@app_commands.describe(
    channel='The channel to end the conversation in. Defaults to the channel it was called in.',
)
@has_permissions(manage_channels=True, manage_messages=True)
async def start(ctx: discord.Interaction, channel: Optional[discord.AppCommandOptionType.channel]=None):
    s = requests.Session()
    payload = {
        "model": mixtral,
        "preprompt": sysp
    }
    headers = {
        "authority": "huggingface.co",
        "accept": "*/*",
        "accept-language": "en-US,en;q=0.9",
        "content-type": "application/json"
    }
    if channel == None:
        chping = f"<#{ctx.channel_id}>"
        os.remove(str(ctx.channel_id))
    else:
        chping = f"<#{channel.id}>"
        os.remove(str(ctx.channel_id))
    ctx.response.send_message(f"Ended the conversation in {chping}")
    s.close()

@end.error
async def end_error(error, ctx):
    if isinstance(error, commands.MissingPermissions):
        ctx.response.send_message(f"Missing permissions in {chping}")

@client.command()
@is_owner()
async def sync(ctx: Context) -> None:
    synced = await client.tree.sync()
    await ctx.reply("{} commands synced".format(len(synced)))

@client.event
async def on_ready():
    print("Ready!")
    synced = await client.tree.sync()

@client.event
async def on_message_create(msg):
    if os.path.exists(str(msg.channel.id)) or isinstance(msg.channel, discord.DMChannel):
        with open(str(msg.channel.id), "r") as f:
            cookie, id = f.read().split("\n")
        r = await chat(id, "{" + str(msg.author) + "}\n" + msg.content, id)
        await msg.reply(json.loads("{" + r.text.split("finalAnswer")[-1][2:].strip())["text"].strip())

client.run(os.environ["TOKEN"])