from gradio_client import Client

from deep_translator import GoogleTranslator

def predict(prompt, client, model = "0.1"):
	global iddb
	if model == "0.1":
		result = client.predict(
    	prompt,
    	"Default",
    	api_name="/chat"
		)
	elif model == "0.2":
		result = client.predict(
			prompt,
			0.5,  # 'Temperature'
			128,  # 'Max new tokens'
			0.8,  # 'Top-p (nucleus sampling)'
			1.8,  # 'Repetition penalty'
			api_name="/chat"
	  )

	return result

# Detect code
def iscode(text):
	langs = ['sql','php','js','java','c','cpp','python','go']
	is_code = False
	for i in langs:
		if i + r'\n' in text:
			is_code = True
			break
	spec = ['section .']
	if not is_code:
		for i in spec:
			if i in text:
				is_code = True
				break
	return is_code

#  text IN  language IN
def translate(text, source):
	if source == "ru":
		target = "en"
	elif source == "en":
		target = "ru"

	# Fix code translate
	if '```' in text:
		out = ''
		for i in text.split('```'):
			if iscode(i):
				out += '```' + i + '```'
			else:
				out += GoogleTranslator(source = source, target = target).translate(i)
	else:
		out = GoogleTranslator(source = source, target = target).translate(text)
	return out

iddb = {}

def gen(text, id, model):
	global iddb

	if str(id) not in iddb:
		if model == "0.1":
			client = Client("https://afischer1985-ai-interface.hf.space/")
		elif model == "0.2":
			client = Client("https://skier8402-mistral-super-fast.hf.space/")
		iddb[str(id)] = client
	else:
		client = iddb[str(id)]

	prompt = translate(text, "ru")

	success = False
	while not success:
		try:
			predicted = predict(prompt, client, model).replace("</s>", "")
			success = True
		except:
			pass

	fixed = predicted.replace(r'\n', '\n').replace('\\ n', '\n')
	return translate(fixed, "en")