import os
from ollama import generate

def Prompt_template(query, src_language, trg_language):
    instruction = f'Translate the following sentences from {src_language} to {trg_language}.'
    prompt = (
        'Below is an instruction that describes a task, paired with an input that provides further context. '
        'Write a response that appropriately completes the request.\n'
        f'### Instruction:\n{instruction}\n'
        f'### Input:\n{query}\n### Response:'
    )
    return prompt

# Definirea folderelor
eng_folder = 'eng'
ro_folder = 'ro'

# Creează folderul "ro" dacă nu există
os.makedirs(ro_folder, exist_ok=True)
print("Pregătire: Folderul 'ro' a fost creat sau există deja.")

# Procesează fișierele din folderul "eng"
for filename in os.listdir(eng_folder):
    if filename.endswith('.txt'):
        base = os.path.splitext(filename)[0]
        filepath = os.path.join(eng_folder, filename)
        print(f"\nÎncepe procesarea fișierului: {filename}")
        
        # Citește conținutul fișierului
        with open(filepath, 'r', encoding='utf-8') as f:
            text = f.read()
        print(f"Fișierul {filename} citit ({len(text)} caractere).")
        
        # Împarte în bucăți de 3000 caractere
        chunks = [text[i:i+3000] for i in range(0, len(text), 3000)]
        print(f"Fișierul {filename} împărțit în {len(chunks)} bucăți.")
        
        # Procesează fiecare bucată
        for i, chunk in enumerate(chunks, 1):
            print(f"Traducere bucată {i}/{len(chunks)} din {filename} ({len(chunk)} caractere)...")
            prompt = Prompt_template(chunk, 'English', 'Romanian')
            response = generate(model='mannix/llamax3-8b-alpaca', prompt=prompt)
            
            # Extrage textul tradus
            translated_text = response['response'].split('### Response:\n')[1].strip()
            print(f"Bucată {i} tradusă cu succes.")
            
            # Salvează bucata tradusă
            chunk_filename = f"{base}_{i}.txt"
            ro_filepath = os.path.join(ro_folder, chunk_filename)
            with open(ro_filepath, 'w', encoding='utf-8') as f:
                f.write(translated_text)
            print(f"Bucată {i} salvată ca {chunk_filename} în folderul 'ro'.")
        
        print(f"Fișierul {filename} procesat complet.")

print("\nToate fișierele au fost traduse și salvate în folderul 'ro'.")