import os
import sys
import argparse
import re
from pathlib import Path
import ollama
from tqdm import tqdm

print("=== Persona2 ContactShit LLM Translator ===\n")

# ====================== НАСТРОЙКИ ======================
MODEL = "qwen3:14b-q4_K_M"
DEFAULT_MAX_DIALOGUES = 0      # 0 = все диалоги
DEFAULT_CHUNK_SIZE = 15
TEMPERATURE = 0.25
CONTEXT = 16384
# ======================================================

parser = argparse.ArgumentParser(description="Перевод скриптов Persona 2 через Ollama")
parser.add_argument('-max', type=int, default=DEFAULT_MAX_DIALOGUES,
                    help='Максимальное количество диалогов для обработки из файла (0 = все)')
parser.add_argument('-chunk', type=int, default=DEFAULT_CHUNK_SIZE,
                    help='Размер чанка в диалогах (по умолчанию 15)')
args = parser.parse_args()

MAX_DIALOGUES = args.max
CHUNK_SIZE = args.chunk

# --------------------- ПРОМПТ ---------------------
SYSTEM_PROMPT = """Ты — профессиональный локализатор японских RPG (Persona) на русский язык.

Правила (строго соблюдать):
- Строки, начинающиеся с \\ — английский оригинал. Никогда их не трогай и не переводи.
- Все теги [xxxx] оставляй точно как есть, в том же порядке и количестве.
- Сохраняй все отступы, пустые строки и структуру.
- Переводи АНГЛИЙСКИЙ ТЕКСТ, который выше, заменяя японский текст на русский, учитывая контекст японского.
- Стиль: естественный, разговорный, соответствующий персонажу.
- Каждая строка диалога — максимум 47 символов.
- Одно диалоговое окно — не более 3 строк + первая строка с именем/тегами.

Выводи ТОЛЬКО готовый скрипт в исходном формате. Ничего лишнего."""

# ======================================================

def get_dialogue_blocks(lines):
    """Разбивает на диалоги по [0611][0211][0311] без \\ в начале"""
    dialogues = []
    current = []
    
    for line in lines:
        stripped = line.strip()
        current.append(line)                     # сохраняем оригинальную строку
        
        if re.search(r'\[0611\]\[0211\]\[0311\]', stripped) and not stripped.startswith('\\'):
            dialogues.append(''.join(current).rstrip('\n'))
            current = []
    
    if current:
        dialogues.append(''.join(current).rstrip('\n'))
    
    return dialogues


# Поиск файлов
txt_files = [f for f in Path('.').glob('*.txt') 
             if f.is_file() and "LLM" not in f.name.upper()]

if not txt_files:
    print("Не найдено .txt файлов для обработки.")
    sys.exit(0)

print(f"Найдено файлов: {len(txt_files)}\n")

for file_path in txt_files:
    print(f"▶ Обработка: {file_path.name}")
    
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    if not lines:
        continue
    
    header = lines[0].strip()
    body_lines = lines[1:]
    
    dialogues = get_dialogue_blocks(body_lines)
    total = len(dialogues)
    
    if MAX_DIALOGUES > 0:
        dialogues = dialogues[:MAX_DIALOGUES]
        print(f"  → Обрабатываем только первые {len(dialogues)} из {total} диалогов")
    
    if not dialogues:
        print("  Нет диалогов.")
        continue
    
    # Разбиваем на чанки
    chunks = [dialogues[i:i + CHUNK_SIZE] for i in range(0, len(dialogues), CHUNK_SIZE)]
    
    translated_chunks = []
    
    for i, chunk in enumerate(tqdm(chunks, desc="  Чанки", leave=False)):
        chunk_text = '\n'.join(chunk)           # без лишних пустых строк между диалогами внутри чанка
        
        try:
            response = ollama.chat(
                model=MODEL,
                messages=[
                    {'role': 'system', 'content': SYSTEM_PROMPT},
                    {'role': 'user', 'content': f"Переведи этот фрагмент скрипта:\n\n{chunk_text}"}
                ],
                options={'temperature': TEMPERATURE, 'num_ctx': CONTEXT}
            )
            translated_chunks.append(response['message']['content'].strip())
        except Exception as e:
            print(f"\n  Ошибка в чанке {i+1}: {e}")
            break
    
    # Финальная сборка БЕЗ лишних пустых строк между чанками
    final_content = header + "\n" + "\n".join(translated_chunks)
    
    output_path = file_path.with_name(f"{file_path.stem}_LLM{file_path.suffix}")
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(final_content)
    
    print(f"  ✓ Готово → {output_path.name} ({len(dialogues)} диалогов)\n")

print("=== Все файлы обработаны! ===")