import telebot

### LOAD TOKEN ###
import json, os

if not os.path.exists('db.json'):
	db = {"token": "None"}
	js = json.dumps(db, indent=2)
	with open("db.json", "w") as outfile:
		outfile.write(js)

	print('Input token in "None" (db.json)')
	exit()
with open('db.json', 'r') as openfile:
	db = json.load(openfile)

API_TOKEN = db["token"]
bot = telebot.TeleBot(API_TOKEN)
##################


@bot.message_handler(commands=['help', 'start'])
def send_welcome(message):
	bot.reply_to(message, "Скоро...")



### MAIN ###
from api import *

setted_models = {}
system_prompts = {}

@bot.message_handler(commands=['info'])
def info(message):
	global setted_models, system_prompts
	id = str(message.chat.id)
	if id not in setted_models:
		setted_models[id] = "0.1"
	if id not in system_prompts:
		prompt = "None"
	else:
		prompt = system_prompts[str(message.chat.id)]

	bot.send_message(message.chat.id, f"""____ Информация ____
Версия: {setted_models[id]}
System-prompt: {telebot.formatting.hcode(prompt)}
""", parse_mode="HTML")

@bot.message_handler(commands=['model'])
def set_model(message):
	global setted_models, iddb
	try:
		iddb.pop(str(message.chat.id))
	except:
		pass

	model = message.text.split()[1]
	if model == "0.1" or model == "0.2":
		setted_models[str(message.chat.id)] = model
		bot.reply_to(message, "Установлена новая модель 🤖")
	else:
		bot.reply_to(message, "Неизвестная модель")


@bot.message_handler(commands=['prompt'])
def set_prompt(message):
	global system_prompts
	system_prompts[str(message.chat.id)] = message.text[8:]
	bot.reply_to(message, "Установлен новый system-prompt")
@bot.message_handler(commands=['cprompt'])
def clear_prompt(message):
	global system_prompts
	system_prompts.pop(str(message.chat.id))
	bot.reply_to(message, "System-prompt очищен")

@bot.message_handler(func=lambda message: True)
def echo_message(message):
	global setted_models, system_prompts

	id = str(message.chat.id)
	if id not in setted_models:
		setted_models[id] = "0.1"

	if id in system_prompts:
		prompt = '[INST]' + system_prompts[id] + '[/INST]\n\n' + message.text
	else:
		prompt = message.text

	st = bot.send_message(message.chat.id, "Печатает...")
	bot.reply_to(message, gen(prompt, message.chat.id, setted_models[id]).replace(r'\n', '\n'), parse_mode="HTML")
	bot.delete_message(message.chat.id, st.id)



############

bot.infinity_polling()