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, f"""Это подобие ChatGPT на минималках.

__ Есть 2 версии:
{telebot.formatting.hcode("/m 0.1")} - Простейшая, быстрая, краткая, не помнит что вы говорили.
{telebot.formatting.hcode("/m 0.2")} - Умнее, относительно быстрая, помнит что вы говорили.

__ Список кратких команд:
/info - /i
/model - /m
/ccontext - /cc
/prompt - /p
/cprompt - /cp

Ещё проекты: @just_openbots (тут много интересного)
{telebot.formatting.hlink("Исходный код","https://gitea.gulyaipole.fun/justuser/just_minigpt")}
""", parse_mode = "HTML")
	bot.send_message(message.chat.id, f"""Также настоятельно рекомендую подписаться на канал бота: @justuser31

Обратная связь ( @just_anonchat_bot ) : {telebot.formatting.hcode(":justuser")}""", parse_mode = "HTML")


### MAIN ###
from api import *

import traceback
import logging

setted_models = {}
system_prompts = {}
onoff = {}

@bot.message_handler(commands=['info','i'])
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','m'])
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','p'])
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','cp'])
def clear_prompt(message):
	global system_prompts
	try:
		system_prompts.pop(str(message.chat.id))
	except:
		pass
	bot.reply_to(message, "System-prompt очищен 🤖")


@bot.message_handler(commands=['ccontext','cc'])
def clear_context(message):
	global iddb
	try:
	  iddb.pop(str(message.chat.id))
	except:
		pass
	bot.reply_to(message, "Контекст очищен 🤖")

#@bot.message_handler(commands=['onoff'])
#def set_onoff(message):
#


@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, "Печатает...")
	predicted = gen(prompt, message.chat.id, setted_models[id])
	try:
		bot.reply_to(message, predicted, parse_mode="Markdown")
	except Exception as e:
		bot.reply_to(message, predicted, parse_mode="HTML")
		logging.error(traceback.format_exc())
		print(predicted)

	bot.delete_message(message.chat.id, st.id)


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

bot.infinity_polling()