import os
import json
import base64
import re
from PIL import Image
import requests

def sanitize_filename(prompt):
    """Очистка имени файла с учётом максимальной длины 146 символов"""
    forbidden_chars = r'[\\/:\*\?"<>\|]'
    sanitized = re.sub(forbidden_chars, '_', prompt)
    sanitized = sanitized.strip().rstrip('.')
    
    # Рассчёт максимальной длины для промпта
    max_prompt_length = 125  # 146 - (5+1+10+1+4) = 125
    if len(sanitized) > max_prompt_length:
        sanitized = sanitized[:max_prompt_length].rsplit(' ', 1)[0]
    
    return sanitized or 'generated_image'

def get_next_counter(output_dir):
    max_num = -1
    if os.path.exists(output_dir):
        for fname in os.listdir(output_dir):
            if match := re.match(r'^(\d{5})-\d+', fname):
                num = int(match.group(1))
                max_num = max(max_num, num)
    return max_num + 1 if max_num != -1 else 0

def parse_parameters(parameters_str):
    lines = [line.strip() for line in parameters_str.split('\n') if line.strip()]
    
    prompt = lines[0] if lines else ""
    negative_prompt = ""
    params = {}
    
    current_line = 1
    
    if current_line < len(lines) and lines[current_line].startswith("Negative prompt:"):
        negative_prompt = lines[current_line].split(":", 1)[1].strip()
        current_line += 1
    
    params_block = " ".join(lines[current_line:])
    param_pattern = r"([\w\s\+\-]+?):\s*([^,\n]+)"
    
    for match in re.finditer(param_pattern, params_block):
        key = match.group(1).strip()
        value = match.group(2).strip()
        params[key] = value

    return prompt, negative_prompt, params

def get_input(prompt, validation_func, error_message):
    while True:
        value = input(prompt)
        if validation_func(value):
            return value
        print(error_message)

def generate_images():
    input_dir = get_input(
        "Введите путь к папке с исходными картинками: ",
        lambda x: os.path.isdir(x),
        "Ошибка: директория не существует!"
    )
    
    output_dir = get_input(
        "Введите путь для сохранения новых картинок: ",
        lambda x: True,
        ""
    )
    
    num_generations = int(get_input(
        "Введите количество генераций (включая исходную): ",
        lambda x: x.isdigit() and int(x) > 0,
        "Ошибка: введите целое число больше 0!"
    ))

    os.makedirs(output_dir, exist_ok=True)
    global_counter = get_next_counter(output_dir)
    processed_files = 0
    
    for filename in os.listdir(input_dir):
        input_path = os.path.join(input_dir, filename)
        
        if not os.path.isfile(input_path) or not filename.lower().endswith('.png'):
            continue
            
        print(f"\n{'='*40}\nProcessing file: {filename}")
        
        try:
            with Image.open(input_path) as img:
                parameters_str = img.info.get('parameters', '')
                
            if not parameters_str:
                print("Skipped: No parameters found")
                continue
                
            prompt, negative_prompt, params = parse_parameters(parameters_str)
            
            required_params = ['Seed', 'Steps', 'Sampler', 'CFG scale', 'Size']
            missing = [p for p in required_params if p not in params]
            
            if missing:
                raise KeyError(f"Missing parameters: {', '.join(missing)}")
            
            original_seed = int(params['Seed'])
            steps = int(params['Steps'])
            sampler = f"{params['Sampler']} {params.get('Schedule type', '')}".strip()
            cfg_scale = float(params['CFG scale'])
            width, height = map(int, params['Size'].split('x'))
            
            for i in range(num_generations):
                current_seed = original_seed + i
                
                payload = {
                    "prompt": prompt,
                    "negative_prompt": negative_prompt,
                    "steps": steps,
                    "sampler_name": sampler,
                    "cfg_scale": cfg_scale,
                    "width": width,
                    "height": height,
                    "seed": current_seed,
                }
                
                response = requests.post("http://localhost:7860/sdapi/v1/txt2img", json=payload)
                response.raise_for_status()
                result = response.json()
                
                image_data = base64.b64decode(result['images'][0])
                info = json.loads(result.get('info', '{}'))
                used_seed = info.get('seed', current_seed)
                
                clean_prompt = sanitize_filename(prompt)
                output_filename = f"{global_counter:05}-{used_seed}-{clean_prompt}.png"
                
                # Финальная проверка длины имени файла
                if len(output_filename) > 146:
                    raise ValueError(f"Слишком длинное имя файла: {len(output_filename)} символов (макс. 146)")
                
                output_path = os.path.join(output_dir, output_filename)
                
                with open(output_path, 'wb') as f:
                    f.write(image_data)
                
                print(f"Generated: {output_filename} ({len(output_filename)} chars)")
                global_counter += 1
            
            processed_files += 1
            
        except KeyError as e:
            print(f"ERROR: {e}")
            print("Available parameters:", json.dumps(params, indent=2))
        except Exception as e:
            print(f"ERROR: {str(e)}")
    
    print(f"\n{'='*40}")
    print(f"Обработка завершена!")
    print(f"Всего обработано файлов: {processed_files}")
    print(f"Сгенерировано изображений: {global_counter - get_next_counter(output_dir)}")

if __name__ == "__main__":
    generate_images()