import os
from llama_cpp import Llama
import tkinter as tk
from tkinter import filedialog

# Carrega o modelo em formato GGUF
model_path = "C:/Users/Chris/.cache/lm-studio/models/hugging-quants/Llama-3.2-3B-Instruct-Q8_0-GGUF/llama-3.2-3b-instruct-q8_0.gguf"
llm = Llama(model_path=model_path)

def summarize_text(text, chunk_size=130):
	# Divide o texto em blocos de chunk_size palavras para evitar o limite de contexto
	words = text.split()
	summaries = []
	
	for i in range(0, len(words), chunk_size):
		chunk = " ".join(words[i:i + chunk_size])
		output = llm(chunk, max_tokens=100)  # Limite de tokens para o resumo
		summaries.append(output["choices"][0]["text"].strip())
	
	# Junta os resumos dos blocos
	return " ".join(summaries)

def process_folder(folder_path):
	# Caminha recursivamente na pasta
	for root, dirs, files in os.walk(folder_path):
		if "FolderSummary.md" in files:
			continue  # Ignora pastas que já têm um resumo

		folder_summary = []
		for file in files:
			if file.endswith(".srt"):
				file_path = os.path.join(root, file)
				try:
					with open(file_path, "r", encoding="utf-8") as f:
						text = f.read()
						summary = summarize_text(text)
						folder_summary.append(f"## {file}\n{summary}\n")
				except Exception as e:
					print(f"Ocorreu um erro ao processar o arquivo {file_path}: {e}")
					continue  # Pula para o próximo arquivo

		# Cria o FolderSummary.md com os resumos
		if folder_summary:
			summary_path = os.path.join(root, "FolderSummary.md")
			with open(summary_path, "w", encoding="utf-8") as f:
				f.write("# Resumo da Pasta\n\n" + "\n".join(folder_summary))
			print(f"Resumo criado em {summary_path}")

while True:
	try:
		# Solicita o diretório ao usuário
		root = tk.Tk()
		root.withdraw()
		folder_selected = filedialog.askdirectory(title="Selecione a pasta para sumarizar")

		if folder_selected:
			process_folder(folder_selected)
			print("Processamento concluído.")
		else:
			print("Nenhuma pasta foi selecionada.")
		
		# Pergunta se o usuário quer processar outra pasta
		retry = input("Deseja selecionar outra pasta? (s/n): ").strip().lower()
		if retry != 's':
			print("Encerrando o programa.")
			break
	except Exception as e:
		print(f"Ocorreu um erro: {e}")
		retry = input("Deseja tentar novamente? (s/n): ").strip().lower()
		if retry != 's':
			print("Encerrando o programa.")
			break