import whisper
import os
import gc
import psutil

CHUNKS_DIR = "chunks"
OUTPUT_FILE = "metereologia2_transcription.txt"
MODEL_NAME = "tiny"

def get_memory_info():
    mem = psutil.virtual_memory()
    return f"RAM: {mem.used / 1024**2:.0f}MB used / {mem.available / 1024**2:.0f}MB available ({mem.percent}%)"

def main():
    chunk_files = sorted([f for f in os.listdir(CHUNKS_DIR) if f.endswith(".mp3")])
    print(f"Found {len(chunk_files)} chunks to transcribe")
    print(f"Before loading model - {get_memory_info()}")

    model = whisper.load_model(MODEL_NAME)
    print(f"Model loaded - {get_memory_info()}")

    full_text = []

    for i, chunk_file in enumerate(chunk_files):
        chunk_path = os.path.join(CHUNKS_DIR, chunk_file)
        print(f"\n--- Chunk {i+1}/{len(chunk_files)}: {chunk_file} ---")
        print(f"Before transcription - {get_memory_info()}")

        result = model.transcribe(chunk_path, language="pt", fp16=False)
        text = result["text"].strip()
        full_text.append(text)

        print(f"After transcription - {get_memory_info()}")
        print(f"Text preview: {text[:100]}...")

        del result
        gc.collect()

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        f.write("\n\n".join(full_text))

    print(f"\n=== DONE === Transcription saved to {OUTPUT_FILE}")
    print(f"Final - {get_memory_info()}")

if __name__ == "__main__":
    main()