import telebot
from telebot.types import InlineKeyboardMarkup, InlineKeyboardButton

### 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 на минималках.

__ Есть 3 версии:
0.1 - Простейшая, быстрейшая, краткая, не помнит что вы говорили.
0.2 - Умнее, относительно быстрая, помнит что вы говорили.
/m - Выбор модели

__ Список команд:
help - Справка
i - Информация о конфигурации
m - Сменить модель
cc - Очистить контекст
p - Задать инструкцию (системную)
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 *
from img_api import *

import traceback
import logging

# Заданные модели
setted_models = {}
# Системные инструкции
system_prompts = {}

# Размер изображения
img_size = {}


######### INFO AND HELP #################

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


############ MODEL SELECTION ############

@bot.message_handler(commands=['m'])
def set_model(message):
	mm = bot.send_message(message.chat.id, "Выберите новую модель:")
	m_id = mm.id
	bot.edit_message_text("Выберите новую модель:", chat_id = message.chat.id, message_id = m_id
	, reply_markup=gen_markup(str(message.chat.id), m_id, message.chat.id))

def gen_markup(id, m_id, c_id):
	markup = InlineKeyboardMarkup()
	markup.row_width = 3
	markup.add(InlineKeyboardButton("0.1", callback_data=id+"_0.1_"+str(m_id)+"_"+str(c_id)),
						InlineKeyboardButton("0.2", callback_data=id+"_0.2_"+str(m_id)+"_"+str(c_id)),
						InlineKeyboardButton("RWKV", callback_data=id+"_RWKV_"+str(m_id)+"_"+str(c_id)),
						)
						# InlineKeyboardButton("VER", callback_data=id+"_VER_"+str(m_id)+"_"+str(c_id)),
	return markup

@bot.callback_query_handler(func=lambda call: True)
def callback_query(call):
	global setted_models, iddb
	id, m, m_id, c_id = call.data.split("_")
	m_id = int(m_id)
	c_id = int(c_id)

	try:
		iddb.pop(id)
	except:
		pass

	setted_models[id] = m
	bot.edit_message_text(f"Успешно установлена модель {m} 🤖", chat_id = c_id, message_id = m_id)

############### PROMPT ###################

@bot.message_handler(commands=['p'])
def set_prompt(message):
	global system_prompts
	system_prompts[str(message.chat.id)] = message.text[3:]
	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=['cc'])
def clear_context(message):
	global iddb, history
	try:
	  iddb.pop(str(message.chat.id))
	except:
		pass
	try:
		history[id] = []
	except:
		pass
	bot.reply_to(message, "Контекст очищен 🤖")


############### IMG  ####################

@bot.message_handler(commands=['img'])
def draw_image(message):
	mm = bot.send_message(message.chat.id, "Генерируем изображение...")
	try:
		m_id = mm.id
		prompt = " ".join( message.text.split()[1:] )

		img_way = draw(prompt, message.chat.id)
		with open(img_way, 'rb') as f:
			img = f.read()
		f.close()

		bot.send_photo(message.chat.id, img, f'{telebot.formatting.hcode(prompt)}', parse_mode="HTML")
		bot.delete_message(message.chat.id, m_id)
	except:
		bot.send_message(message.chat.id, "Извините, возникла непредвиденная ошибка")


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

@bot.message_handler(func=lambda message: True)
def echo_message(message):
	# Отвечаем в ЛС, либо по команде
	if bot.get_chat(message.chat.id).type == "private" or message.text[:2] == "/a":
		global setted_models, system_prompts

		# Текст генерации при команде
		if message.text[:2] == "/a":
			text = message.text[3:]
		else:
			text = message.text

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


		mm = bot.send_message(message.chat.id, "Печатает...")
		m_id = mm.id


		# Если задана инструкция
		if id in system_prompts:
			if setted_models[id] == "0.1" or setted_models[id] == "0.2":
				prompt = '[INST]' + system_prompts[id] + '[/INST]\n\n' + text
			elif setted_models[id] == "RWKV":
				prompt = f'''
				Instruction: {system_prompts[id]}
				\nInput:{text}
				\nResponse:\n'''
		# Если инструкция не задана
		else:
			if setted_models[id] == "0.1" or setted_models[id] == "0.2":
				prompt = text
			elif setted_models[id] == "RWKV":
				prompt = f'''
				Input: {text}
				\nResponse:\n'''

		try:
		#if 1:
			predicted = gen(prompt, message.chat.id, setted_models[id])
		except:
			bot.send_message(message.chat.id, "Извините, возникла непредвиденная ошибка")

		try:
			bot.edit_message_text(predicted, chat_id=message.chat.id, message_id=m_id, parse_mode="Markdown")
		except Exception as e:
			bot.edit_message_text(predicted, chat_id=message.chat.id, message_id=m_id, parse_mode="HTML")
			logging.error(traceback.format_exc())

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

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

bot.infinity_polling()