const express 	= require('express')
const app 		= express()
const port 		= 8080

process.on('uncaughtException', err => log('JayCoach:Exception:', err))

const hfToken = process.env.HF_TOKEN


const MODELS = {
	
	
	'phi3.5': {
		name: 'microsoft/Phi-3.5-mini-instruct'
		,prompt: function(prompt){
			return [
				"<|user|>"
				,prompt
				+"<|end|>"
				,"<|assistant|>"
			].join("\n")
		}
	}
	
	
	,'lama3.1': {
		name: 'meta-llama/Meta-Llama-3.1-8B-Instruct'
		,prompt: function(prompt){
			return [
				"<|start_header_id|>user<|end_header_id|>"
				,prompt+"<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
				+""
			].join("\n")
		}
	}
	
	
	,'lama3.2v': {
		name: 'meta-llama/Llama-3.2-11B-Vision-Instruct'
		,prompt: function(prompt){
			return [
				"<|start_header_id|>user<|end_header_id|>"
				,prompt+"<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
				+""
			].join("\n")
		}
	}
	
	
	
	
	,'mistral-small':{
		name: 'mistralai/Mistral-Small-Instruct-2409'
		,prompt: function(prompt){
			return [
				"[INST]"+prompt+"[/INST]"
			].join("\n")
		}
		
	}
	
	
	,'gemma':{
		name: 'google/gemma-1.1-7b-it'
		,prompt: function(prompt){
			return [
				"<start_of_turn>user"
				,prompt+"<end_of_turn>"
				,"<start_of_turn>model"
				,""
			].join("\n")
		}
		
	}
	
}

function log(m){
	let d = new Date();
	let diso = d.toISOString();
	
	console.log(`${diso} ${m}`)
}


if(!hfToken){
	throw new Error('NO TOKEN!');
}

let LastWorkedModel = '';
const ModelNames = Object.keys(MODELS);
const ModelList = []
let BestModel = {}

for(let modelId in MODELS){
	let model = MODELS[modelId]
	
	model.id 	= modelId
	model.stats = {
		 total:0
		,erros:0
		,parcela:1
		,get errop() { return this.total === 0 ? 0 : this.erros/this.total } 
		,get pok() { return 1 - this.errop } 
	}
	
	ModelList.push(model);
}

// Encontrar o modelo que oferece que tem melhor chances do acerto!
/*
	Se você não entendeu o codigo abaixo, parabens. Eu tb não rsrs.
	brincadeir... 
	Aqui é apenas uma pequena maneira de calcular o melhor modelo rpa ser usado...
	Cada model tem uma prop que chamei de pok (Percentual de OK = % de sucesso quando o model foi usado!)
	Então, vamos somar todos os percentuais de ok que temos, e atribuir uma parcela desse total pra cada model.
	O model que tiver mais % de ok dos demais, tem mais chances de ser escolhido do que um que tem menos... 
	
	Exemplos:
	
	 Phi3  Gemini 			Lama
	|----|--------|-----------------------------|
	  10%    20%			 70%
	   
	    Phi3  	Gemini 				Lama
	|--------|----------------|----------------|
	    20%			 39.9%			40.1%
	   
	Agora, escolhe ai um número entre 0 e 100%, aleatoriamente.
	Se for até 10%, pega  Google... 
	Se for até entre 10 e 20%, vai o Lama...e acma de 20% vai a Microsoft...
	
	Sacou a manha? Com isso, quanto mais um model nao da erro, mais ele tem a chance de ser escolhido!
	Se 2 models tem o mesmo peso, vamos usar um pequeno hack na conta pra nunca dar empate e um deles sempre ter 1 pouco a mais!
	   
	
	O reusumo é: Imagine aquele meme da Nazaré fazendo as contas!
*/
function UpdateProbs(opts){
	BestModel.LastRandom = Math.random();
	
	let AllModels = [];
	
	// total de "oks"
	let Total = ModelList.reduce( (acc,m) => acc + m.stats.pok , 0 )
	
	// calcula parcela de ok desse model!
	ModelList.forEach( m => m.stats.parcela = m.stats.pok/Total  )
	
	// Organiza pela ordem...
	let SortedModels = ModelList.sort( (a,b) => {
		let diff = parseFloat(a.stats.parcela.toPrecision(2)) - parseFloat(b.stats.parcela.toPrecision(2));
		
		if(diff == 0)
			diff = a.stats.parcela*Math.random() - b.stats.parcela*Math.random()
		
		return diff;
	})
	
	BestModel.LastSorted = SortedModels;
	
	let parcAcc = 0;
	for(let [idx,model] of SortedModels.entries()){
		let stats = model.stats;
		parcAcc += stats.parcela;
		
		if(BestModel.LastRandom <= parcAcc){
			BestModel.model = model.id;
			return;
		}
	}
	
	
	return;
}

async function GetModelAnswer(model, prompt){
	
	let StartIndex;
	
	if(!model){
		UpdateProbs();
		model = BestModel.model;
	}

	
	let i = ModelList.length;
	while(i--){ // pra evitar um loop infinito, vai girar no maximo o numero de models...
		let ModelConfig = MODELS[model];
		let MyStats = ModelConfig.stats;
		log("Stats:");
		console.log(MyStats);
		let InferenceApi = 'https://api-inference.huggingface.co/models/'  + ModelConfig.name;
		
		let data ={
			inputs: ModelConfig.prompt(prompt)
			,parameters:{
				 max_new_tokens: 70
				,return_full_text: false
				,temperature: 0.5
			}
			,options:{
				use_cache: false
				,wait_for_model: false
			}
		}
		
		log("Falando com a IA 🤖")
		console.log(model, ModelConfig.name)
		MyStats.total++;
		let StartTime = new Date();
		const response = await fetch(
			InferenceApi,
			{
				headers: { Authorization: "Bearer "+hfToken, "content-type":"application/json" },
				method: "POST",
				body: JSON.stringify(data),
			}
		);
		let EndTime = new Date();
		let ElapsedTime = EndTime.getTime()  - StartTime.getTime();
		log("Total", ElapsedTime);
		
		if(response.status != 200){
			MyStats.erros++;
			log('FAILED: Escolhendo outro...'+response.status)
			if(StartIndex == null)
				StartIndex = ModelList.map(m => m.id).indexOf(model);
			
			let NextIndex = StartIndex+1;
			
			if(NextIndex >= ModelList.length)
				NextIndex = 0;
		
			if(NextIndex == StartIndex){
				log("Fiz de tudo, mas não deu bom :(");
				throw new Error('SOME_SHIT_HAPPENS');
			}
			
			model = ModelList[NextIndex].id;
			log("Tentando com o ",model);
			continue;
		}
		
		// Validacoes adicionais de erros!
		// Tempo de resposta maior que que 2s?
		// Penaliza
		if(ElapsedTime >= 2500)
			MyStats.erros += 0.100;
		
		if(ElapsedTime < 900){
			MyStats.erros -= 0.100;
			if(MyStats.erros < 0) MyStats.erros = 0;
		}
		
		log("Ok, lendo o json..."+response.status);
		const result = await response.json();
		LastWorkedModel = model;
		return {
			result
			,model
		}
	}
	
	// se chegou aqui é pq todo mundo falhou!
	throw new Error('Nenhum model respondeu! O trem tá feio ou o dev cagou em algo...')
}

async function Prompt(opts){
	
	let error 		= opts.error 
	let tentativas 	= opts.tentativas
	let max 		= opts.max
	let model 		= opts.model
	
	if(!max)
		max = 20
	
	if(tentativas){
		tentativas = 'últimas tentativas:'+tentativas
	} else 
		tentativas = ""
	
	let tom = "";
	let statusErro = ""
	if(error <= 450){
		tom = `Gere mensagens com bastantes elogios e tom de dúvida (e sarcasmo) se foi realmente um humano que fez isso...
			EXEMPLOS:
				- Rapaz, acho que isso foi humanamente impossível...|fim|
				- Você não está usando um script não né?|fim|
				- Não é possível, tá muito baixo pra ter sido um ser humano...|fim|
		`
		statusErro = "Objetivo atingido! Menor que 450"
	} else if(error <= 500){
		tom = `Gere mensagens que parabenizem e elogiem o desempenho.
			EXEMPLOS:
				- Muito, muito, mas muito bom!|fim|
				- ora, ora ora, temos um Vingador da IA aqui|fim|
				- Você é o pica das galáxias da IA hein!|fim|
		`
		statusErro = "Objetivo atingido! Menor que 500"
	}else if(error <= 2000){
		tom = `Gere mensagens inspiradoras, no sentido em que está indo bem!
			EXEMPLOS:
				- Vamos lá, dá pra melhorar, você consegue|fim|
				- Não desista, continue tentando|fim|
			`
		statusErro = "Objetivo não atingido! Maior que 500"
	} else {
		tom = `Gere mensagens sarcástias e engraçadas brincando com a situação. Faça piadas e zoeiras.
			EXEMPLOS:
				- Ei, psiu, volta aqui|fim|
				- Ou, não é pra aí não, volta aqui|fim|
				- Você ainda tá tentando ou tá só de brincadeira mesmo?|fim|
				- Ainda bem que eu não sou você hein...|fim|
				- Nossa, mas esse erro tá sensacionalmente errado!|fim|
				- Muito bom continue assim #sqn|fim|
		`
		statusErro = "Objetivo não atingindo! Muito longe!"
	}
	
	
	let prompt = `
		Um usuário está estudando Redes Neurais e IA e está aprendendo o conceito de Erro (erro quadrático médio).
		Ele está fazendo um exercício onde deve conseguir gerar um erro < 499 (menor que 499).
		
		${tom}
		
		---
		Informações das tentativas:
		
			Erro atual: ${error}
			${tentativas}
			Status: ${statusErro}
			 
		Gere uma mensagem para ser exibida ao usuário com base nas informacoes das tentativas fornecidas.
		Use emojis nas respostas, quando apropriado!
		REGRAS: 
			- máx ${max} palavras
			- Responda como se estivesse falando diretamente com o usuário (use a segunda pessoa "você").
			- encerrar com |fim|
			
			
		---
	`
	
	log("Prompt Info: ")
	console.log(prompt.length, prompt);
	let answer = await GetModelAnswer(model, prompt);
	return answer;
}

app.get('/error', async (req, res) => {
	
	let tentativas = req.query.tentativas;
	let max = 20;
	
	if(tentativas && tentativas.length >= 100){
		res.json({error:"Tentando atacar né?"})
		return;
	}
	
	let error = req.query.error;
	
	if(/^[^0-9]+$/g.test(error)){
		res.json({error:"Tentando atacar né?"})
		return;
	}
	
	if(tentativas)
		tentativas = tentativas.split(",").map(Number).join(",");
	
	let StartTime = new Date();
	let result = await Prompt({
		 error:req.query.error
		,tentativas
		,model:req.query.model
		,max
	});
	let EndTime = new Date();
	let TotalMs = EndTime.getTime() - StartTime.getTime();
	
	let ModelInfo = MODELS[result.model]
	
	log("Respondido:")
	console.log(TotalMs, result);
	
	let resp = result.result;
	if(!resp || !Array.isArray(resp)){
		res.json({text:":("});
		ModelInfo.stats.erros += 0.2;
		return;
	}
	
	let gentext = resp[0].generated_text
	
	let textParts = gentext.split('|fim|');
	
	if(textParts.length < 2){
		ModelInfo.stats.erros += 0.1;
	}
	
	let txtFinal = textParts[0].trim();
	
	let estimatedChars = max*8;
	if(txtFinal.length >= estimatedChars){
		txtFinal = txtFinal.slice(0,estimatedChars);
		ModelInfo.stats.erros += 0.05;
	}
	

	log("FullResp:"+gentext);
	log(`Final:${txtFinal}`);
	res.json({text:txtFinal, model:result.model, hfName:ModelInfo.name, TotalMs})
})

app.get('/test', async (req, res) => {
	res.send("Working!")
})

app.get('/models', async (req, res) => {
	//UpdateProbs()
	
	res.json({
		 BestModel
	})
})

app.get('/', async (req, res) => {
	res.send('JayCoach ON! Veja mais no blog IA Talking: <a href="https://iatalk.ing/tag/jay-trainer">https://iatalk.ing</a>')
})

app.use(function(err, req, res, next) {
  console.error(err.stack);
  res.json({error:'Server error, admin must check logs',status:res.status})
});


app.listen(port, () => {
  log(`JayCoach running`)
})