from google import genai
from google.genai import types
import time
from pylatexenc.latex2text import LatexNodes2Text
import os
from dotenv import load_dotenv

class gemini():
	
	def __init__(self, api_key):
		self.api_key = api_key
		self.client = genai.Client(api_key=self.api_key)
	
	def save_query(self, query_number, queries, file_name):
		file_name = latex2text(file_name)
		with open(f'/storage/emulated/0/Testing/{file_name}.txt','w') as file:
			file.write(queries[query_number])
		print(f'Saved as {file_name}.txt')
		
	def get_response(self, text):
		return self.client.models.generate_content(
		       model="gemini-2.5-flash",
		       contents=text,
		       config=types.GenerateContentConfig(
        thinking_config=types.ThinkingConfig(thinking_budget=0))).text
    	

def slow_print(text, qn):
	print(f'Gemini[{qn}]: ', end="")
	for word in text:
		print(word , sep="" ,end="" , flush=True)
		time.sleep(0.0055)

def latex2text(latex_str):
	text = LatexNodes2Text().latex_to_text(latex_str)
	target = "*"
	if target in list(text):
		text=text.replace("*", "")
	return text
	


def main(qn, q):
	query_number = qn
	queries = q
	load_dotenv()
	api_key = os.environ.get("GEMINI_API_KEY")
	client = gemini(api_key=api_key)
	
	while True:
		user_query = input(f'User[{query_number}]: ')
		
		if user_query == "/quit":
			print("Closing!")
			break
			
		if user_query == "/save":
			query_input = int(input("Query number: "))
			content = queries[query_input] + "In one word give this a title"
			name = client.get_response(text=content)
			client.save_query(query_input, queries, name)
			
		else:
			
			gemini_response = latex2text(client.get_response(text=user_query))
			queries.append(gemini_response)
			slow_print(gemini_response, query_number)
			print("")
			query_number += 1

main(0, [])