# test_ffmpeg.py
import os, re, shutil, subprocess, hashlib, random, time
from pathlib import Path
from datetime import datetime
from moviepy.editor import VideoFileClip
from TTS.api import TTS
from pydub import AudioSegment
from PIL import Image, ImageDraw, ImageFont
from faster_whisper import WhisperModel
import torch
import asyncio
import mouse
import json

# Імпорти модулів
from telegram_bot import send_video_preview
from preview_generator import PreviewGenerator
from horizontal_video import HorizontalVideoGenerator
from vertical_video import VerticalVideoGenerator
from tiktok_processor import TikTokVideoGenerator, process_tiktok_videos, create_tts_with_censoring, get_tiktok_duration_settings

def safe_filename(name):
    """Створює безпечну назву файлу"""
    name = re.sub(r'[<>:"/\\|?*\x00-\x1F]', "_", name)
    name = re.sub(r'\s+','_', name).strip("_")
    name = name.rstrip(". ")
    return name[:50]

def get_speed_settings():
    """Дозволяє користувачу налаштувати прискорення відео"""
    print("\n" + "="*60)
    print("⚡ НАЛАШТУВАННЯ ПРИСКОРЕННЯ ВІДЕО")
    print("="*60)
    print("Виберіть режим прискорення:")
    print("1. ⚡ Без прискорення (оригінальна швидкість)")
    print("2. ???? Фіксоване прискорення для всіх відео")
    print("3. ???? Прискорення до цільової тривалості")
    print("4. ⚙️ Інтерактивне налаштування для кожного відео")
    print("="*60)
    
    while True:
        choice = input("Ваш вибір (1-4): ").strip()
        if choice in ["1", "2", "3", "4"]:
            break
        print("❌ Невірний вибір. Введіть 1, 2, 3 або 4")
    
    speed_config = {
        'mode': choice,
        'youtube_horizontal': {'enabled': False, 'speed': 1.0, 'target_duration': None},
        'youtube_vertical': {'enabled': False, 'speed': 1.0, 'target_duration': None},
        'tiktok': {'enabled': False, 'speed': 1.0, 'target_duration': None}
    }
    
    if choice == "1":
        # Без прискорення
        print("✅ Вибрано режим без прискорення")
        return speed_config
        
    elif choice == "2":
        # Фіксоване прискорення
        print("\n???? Налаштування фіксованого прискорення:")
        while True:
            try:
                speed = float(input("Введіть коефіцієнт прискорення (1.0-4.0, наприклад 1.5): "))
                if 1.0 <= speed <= 4.0:
                    break
                print("❌ Коефіцієнт має бути від 1.0 до 4.0")
            except ValueError:
                print("❌ Введіть дійсне число!")
        
        speed_config['youtube_horizontal'] = {'enabled': True, 'speed': speed, 'target_duration': None}
        speed_config['youtube_vertical'] = {'enabled': True, 'speed': speed, 'target_duration': None}
        speed_config['tiktok'] = {'enabled': True, 'speed': speed, 'target_duration': None}
        print(f"✅ Всі відео будуть прискорені в {speed}x")
        
    elif choice == "3":
        # Прискорення до цільової тривалості
        print("\n???? Налаштування прискорення до цільової тривалості:")
        
        # YouTube горизонтальне
        print("\n????️ YouTube горизонтальне відео:")
        enable = input("Прискорювати YouTube горизонтальні відео? (y/n): ").lower() == 'y'
        if enable:
            while True:
                try:
                    target = float(input("Цільова тривалість (секунди): "))
                    if target > 10:
                        break
                    print("❌ Тривалість має бути більше 10 секунд")
                except ValueError:
                    print("❌ Введіть дійсне число!")
            speed_config['youtube_horizontal'] = {'enabled': True, 'speed': None, 'target_duration': target}
        
        # YouTube вертикальне
        print("\n???? YouTube вертикальне відео:")
        enable = input("Прискорювати YouTube вертикальні відео? (y/n): ").lower() == 'y'
        if enable:
            while True:
                try:
                    target = float(input("Цільова тривалість (секунди): "))
                    if target > 10:
                        break
                    print("❌ Тривалість має бути більше 10 секунд")
                except ValueError:
                    print("❌ Введіть дійсне число!")
            speed_config['youtube_vertical'] = {'enabled': True, 'speed': None, 'target_duration': target}
        
        # TikTok
        print("\n???? TikTok відео:")
        enable = input("Прискорювати TikTok відео? (y/n): ").lower() == 'y'
        if enable:
            while True:
                try:
                    target = float(input("Цільова тривалість кожної частини (секунди): "))
                    if target > 15:
                        break
                    print("❌ Тривалість має бути більше 15 секунд")
                except ValueError:
                    print("❌ Введіть дійсне число!")
            speed_config['tiktok'] = {'enabled': True, 'speed': None, 'target_duration': target}
            
    elif choice == "4":
        # Інтерактивне налаштування
        print("✅ Вибрано інтерактивний режим")
        print("Для кожного відео буде запропоновано налаштування прискорення")
        speed_config['mode'] = 'interactive'
    
    return speed_config

def get_interactive_speed_settings(video_type, original_duration):
    """Запитує у користувача налаштування прискорення для конкретного відео"""
    type_names = {
        'youtube_horizontal': 'YouTube горизонтальне',
        'youtube_vertical': 'YouTube вертикальне',
        'tiktok': 'TikTok частина'
    }
    
    print(f"\n⚙️ Налаштування для {type_names.get(video_type, video_type)}")
    print(f"⏱️ Поточна тривалість: {original_duration:.1f} секунд")
    
    choice = input("Прискорювати це відео? (y/n): ").lower()
    if choice != 'y':
        return {'enabled': False, 'speed': 1.0, 'target_duration': None}
    
    print("Виберіть спосіб:")
    print("1. Фіксоване прискорення (наприклад, 1.5x)")
    print("2. До цільової тривалості")
    
    while True:
        method = input("Ваш вибір (1-2): ").strip()
        if method in ["1", "2"]:
            break
        print("❌ Введіть 1 або 2")
    
    if method == "1":
        while True:
            try:
                speed = float(input("Коефіцієнт прискорення (1.0-4.0): "))
                if 1.0 <= speed <= 4.0:
                    break
                print("❌ Коефіцієнт має бути від 1.0 до 4.0")
            except ValueError:
                print("❌ Введіть дійсне число!")
        
        final_duration = original_duration / speed
        print(f"✅ Швидкість {speed}x, фінальна тривалість: {final_duration:.1f}с")
        return {'enabled': True, 'speed': speed, 'target_duration': None}
    else:
        max_duration = 300 if video_type == 'tiktok' else 600
        while True:
            try:
                target = float(input(f"Цільова тривалість (10-{max_duration}с): "))
                if 10 <= target <= max_duration:
                    break
                print(f"❌ Тривалість має бути від 10 до {max_duration} секунд")
            except ValueError:
                print("❌ Введіть дійсне число!")
        
        required_speed = original_duration / target
        if required_speed > 4.0:
            print(f"⚠️ Потрібне прискорення {required_speed:.1f}x перевищує максимум 4.0x")
            required_speed = 4.0
            final_duration = original_duration / required_speed
            print(f"Буде використано максимальне прискорення 4.0x, тривалість: {final_duration:.1f}с")
        else:
            print(f"✅ Прискорення {required_speed:.1f}x до {target}с")
        
        return {'enabled': True, 'speed': required_speed, 'target_duration': target}

def apply_video_speedup(input_video_path, output_video_path, speed_factor, temp_folder):
    """Прискорює відео зі збереженням якості"""
    if speed_factor <= 1.0:
        shutil.copy2(input_video_path, output_video_path)
        return Path(output_video_path)
    
    print(f"⚡ Прискорення відео в {speed_factor:.2f}x...")
    temp_video = temp_folder / f"speedup_temp_{random.randint(1000, 9999)}.mp4"
    
    try:
        cmd = [
            "ffmpeg", "-y", "-i", str(input_video_path),
            "-filter_complex", f"[0:v]setpts={1/speed_factor}*PTS[v];[0:a]atempo={speed_factor}[a]",
            "-map", "[v]", "-map", "[a]", "-c:v", "h264_nvenc", "-preset", "p4", "-b:v", "8M",
            "-c:a", "aac", "-b:a", "128k", "-movflags", "+faststart", str(temp_video)
        ]
        
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        shutil.move(str(temp_video), str(output_video_path))
        print(f"✅ Відео прискорено успішно")
        return Path(output_video_path)
        
    except subprocess.CalledProcessError as e:
        print(f"❌ Помилка при прискоренні відео: {e}")
        print(f"FFmpeg stderr: {e.stderr}")
        shutil.copy2(input_video_path, output_video_path)
        return Path(output_video_path)
    finally:
        if temp_video.exists():
            temp_video.unlink()

def select_background_folder(backgrounds_dir):
    """Дозволяє користувачу вибрати папку з фоновими відео"""
    print("\n" + "="*60)
    print("???? ВИБІР ФОНОВИХ ВІДЕО")
    print("="*60)
    
    bg_folders = [f for f in backgrounds_dir.iterdir() if f.is_dir()]
    if not bg_folders:
        print("❌ Папки з фоновими відео не знайдено!")
        print(f"Створіть папки в: {backgrounds_dir}")
        print("Наприклад: minecraft, cakes, future, тощо")
        return None
    
    for i, folder in enumerate(bg_folders, 1):
        video_files = list(folder.glob("*.mp4")) + list(folder.glob("*.mov")) + list(folder.glob("*.avi"))
        print(f"{i}. ???? {folder.name} ({len(video_files)} відео)")
    print("="*60)
    
    while True:
        try:
            choice = input(f"Виберіть папку (1-{len(bg_folders)}): ").strip()
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(bg_folders):
                selected_folder = bg_folders[choice_idx]
                video_files = list(selected_folder.glob("*.mp4")) + list(selected_folder.glob("*.mov")) + list(selected_folder.glob("*.avi"))
                if not video_files:
                    print(f"❌ У папці {selected_folder.name} немає відео файлів!")
                    continue
                print(f"✅ Вибрано папку: {selected_folder.name} ({len(video_files)} відео)")
                return selected_folder
            else:
                print(f"❌ Невірний вибір. Введіть число від 1 до {len(bg_folders)}")
        except ValueError:
            print(f"❌ Невірний вибір. Введіть число від 1 до {len(bg_folders)}")

def chunk_words(words, n=2):
    """Групує слова для субтитрів"""
    chunks = []
    i = 0
    while i < len(words):
        group = words[i:i+n]
        text = " ".join([w["word"] for w in group])
        start = group[0]["start"]
        end = group[-1]["end"]
        chunks.append({"text": text, "start": start, "end": end})
        i += n
    return chunks

def create_ass_subtitles(chunks, ass_path, font_path=None, font_size=48, vertical=False):
    """Створює ASS файл субтитрів з центрованим текстом та обводкою"""
    font_name = "Komika Axis"
    if vertical:
        font_size = 28
        margin_v = 250
        margin_lr = 100
    else:
        margin_v = 0
        margin_lr = 30
    
    ass_content = f"""[Script Info]
Title: Generated Subtitles
ScriptType: v4.00+

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,{font_name},{font_size},&H00FFFFFF,&H000000FF,&H00000000,&H80000000,1,0,0,0,100,100,0,0,1,3,0,5,{margin_lr},{margin_lr},{margin_v},1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""
    
    for chunk in chunks:
        start_time = format_ass_time(chunk["start"])
        end_time = format_ass_time(chunk["end"])
        text = chunk["text"].replace("\\", "\\\\").replace("{", "\\{").replace("}", "\\}")
        
        if vertical:
            words = text.split()
            lines = []
            current_line = ""
            max_chars = 20
            
            for word in words:
                if len(word) > max_chars:
                    if current_line:
                        lines.append(current_line.strip())
                        current_line = ""
                    while len(word) > max_chars:
                        lines.append(word[:max_chars-1] + "-")
                        word = word[max_chars-1:]
                    if word:
                        current_line = word
                elif len(current_line + " " + word) <= max_chars:
                    current_line += (" " + word if current_line else word)
                else:
                    if current_line:
                        lines.append(current_line.strip())
                    current_line = word
            
            if current_line:
                lines.append(current_line.strip())
            
            text = "\\N".join(lines)
        
        ass_content += f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{text}\n"
    
    with open(ass_path, 'w', encoding='utf-8-sig') as f:
        f.write(ass_content)

def format_ass_time(seconds):
    """Конвертує секунди в формат часу ASS (H:MM:SS.CC)"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = seconds % 60
    return f"{hours}:{minutes:02d}:{secs:05.2f}"

def apply_subtitles_ass(video_path, chunks, story_id, vertical=False, temp_folder=None):
    """Додає субтитри використовуючи ASS файл"""
    if not chunks:
        return video_path
    
    print(f"???? Створення ASS субтитрів для {len(chunks)} фрагментів...")
    suffix = "_vertical" if vertical else ""
    ass_path = temp_folder / f"{story_id}_subtitles{suffix}.ass"
    create_ass_subtitles(chunks, ass_path, font_size=56, vertical=vertical)
    
    output_video = temp_folder / f"{story_id}_with_subtitles{suffix}.mp4"
    ass_path_escaped = str(ass_path).replace("\\", "/").replace(":", "\\:")
    
    cmd = [
        "ffmpeg", "-y", "-i", str(video_path), "-vf", f"ass='{ass_path_escaped}'",
        "-c:v", "h264_nvenc", "-preset", "p4", "-b:v", "5M", "-c:a", "copy", str(output_video)
    ]
    
    try:
        print(f" ???? Додавання субтитрів...")
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        print(f" ✅ Субтитри успішно додано")
        return output_video
    except subprocess.CalledProcessError as e:
        print(f"❌ Помилка при додаванні субтитрів: {e}")
        return video_path

def select_youtube_video_format():
    """Дозволяє користувачу вибрати формат YouTube відео"""
    print("\n" + "="*60)
    print("???? ВИБІР ФОРМАТУ YOUTUBE ВІДЕО")
    print("="*60)
    print("1. ????️ Горизонтальне відео (1920x1080) - класичний YouTube")
    print("2. ???? Вертикальне відео (1080x1920) - повна версія для Shorts")
    print("3. ???? Обидва формати")
    print("="*60)
    
    while True:
        choice = input("Ваш вибір (1-3): ").strip()
        if choice in ["1", "2", "3"]:
            return choice
        print("❌ Невірний вибір. Введіть 1, 2 або 3")

def create_horizontal_video_with_speedup(story_id, preview_txt, audio_preview, final_audio, selected_template, horizontal_generator, chunks, speed_config, output_dir, video_title, temp_folder):
    """Створює горизонтальне відео з прискоренням"""
    combined_video, preview_img_path = horizontal_generator.create_horizontal_video(
        story_id, preview_txt, audio_preview, final_audio, selected_template
    )
    horizontal_generator.preview_img_path = preview_img_path
    
    if chunks:
        final_video = apply_subtitles_ass(combined_video, chunks, story_id, temp_folder=temp_folder)
    else:
        print("⚠️ Субтитри не знайдено")
        final_video = combined_video
    
    horizontal_output_path = output_dir / f"{video_title}.mp4"
    horizontal_speed_settings = speed_config['youtube_horizontal']
    
    if speed_config['mode'] == 'interactive':
        original_duration = final_audio.duration_seconds
        horizontal_speed_settings = get_interactive_speed_settings('youtube_horizontal', original_duration)
    
    if horizontal_speed_settings['enabled'] and horizontal_speed_settings.get('speed', 1.0) > 1.0:
        print(f"⚡ Прискорення горизонтального відео...")
        speed_factor = horizontal_speed_settings.get('speed')
        if horizontal_speed_settings.get('target_duration'):
            speed_factor = final_audio.duration_seconds / horizontal_speed_settings['target_duration']
            speed_factor = min(speed_factor, 4.0)
        
        horizontal_generator.finalize_video(final_video, temp_folder / f"{story_id}_voice.wav", temp_folder / "temp_horizontal.mp4")
        apply_video_speedup(temp_folder / "temp_horizontal.mp4", horizontal_output_path, speed_factor, temp_folder)
    else:
        horizontal_generator.finalize_video(final_video, temp_folder / f"{story_id}_voice.wav", horizontal_output_path)
    
    print(f"????️ Горизонтальне відео створено: {horizontal_output_path}")

def create_full_vertical_video_with_speedup(story_id, story_txt, preview_txt, horizontal_generator, vertical_generator, selected_template, coqui_tts, asr_model, chunks, speed_config, output_dir, video_title, temp_folder):
    """Створює повне вертикальне відео з прискоренням"""
    audio_preview = coqui_tts(preview_txt)
    preview_img_path = temp_folder / f"{story_id}_preview.png"
    horizontal_generator.preview_generator.draw_horizontal_preview(preview_txt, preview_img_path, selected_template)
    
    combined_vertical_video, vertical_audio_path = vertical_generator.create_full_vertical_video(
        story_txt, preview_txt, story_id, preview_img_path, coqui_tts
    )
    
    print(f"???? Розпізнавання мовлення для вертикального відео...")
    segments_vertical, _ = asr_model.transcribe(str(vertical_audio_path), word_timestamps=True)
    
    story_start_time = audio_preview.duration_seconds + 0.5
    story_words = []
    for seg in segments_vertical:
        for w in seg.words:
            if w.start >= story_start_time:
                story_words.append({"word": w.word.strip(), "start": w.start, "end": w.end})
    
    chunks_vertical = chunk_words(story_words)
    
    if chunks_vertical:
        final_vertical_video = apply_subtitles_ass(combined_vertical_video, chunks_vertical, f"{story_id}_full_vertical", vertical=True, temp_folder=temp_folder)
    else:
        print("⚠️ Субтитри не знайдено для вертикального відео")
        final_vertical_video = combined_vertical_video
    
    vertical_path = output_dir / f"vertical_{video_title}.mp4"
    vertical_speed_settings = speed_config['youtube_vertical']
    
    if speed_config['mode'] == 'interactive':
        vertical_duration = AudioSegment.from_wav(vertical_audio_path).duration_seconds
        vertical_speed_settings = get_interactive_speed_settings('youtube_vertical', vertical_duration)
    
    if vertical_speed_settings['enabled'] and vertical_speed_settings.get('speed', 1.0) > 1.0:
        print(f"⚡ Прискорення вертикального відео...")
        speed_factor = vertical_speed_settings.get('speed')
        if vertical_speed_settings.get('target_duration'):
            speed_factor = AudioSegment.from_wav(vertical_audio_path).duration_seconds / vertical_speed_settings['target_duration']
            speed_factor = min(speed_factor, 4.0)
        
        vertical_generator.finalize_shorts_video(final_vertical_video, vertical_audio_path, temp_folder / "temp_vertical.mp4")
        apply_video_speedup(temp_folder / "temp_vertical.mp4", vertical_path, speed_factor, temp_folder)
    else:
        vertical_generator.finalize_shorts_video(final_vertical_video, vertical_audio_path, vertical_path)
    
    print(f"???? Повне вертикальне відео створено: {vertical_path}")

def process_youtube_story(story_path, horizontal_generator, vertical_generator, selected_template, coqui_tts, asr_model, temp_folder, final_dir, speed_config):
    """Обробляє YouTube історію з вибором формату відео"""
    story_txt = story_path.read_text(encoding="utf-8").strip()
    preview_path = story_path.with_name(f"{story_path.stem}_preview.txt")
    description_path = story_path.with_name(f"{story_path.stem}_description.txt")
    
    if not preview_path.exists() or not description_path.exists():
        print(f"❌ Не знайдено preview або description для {story_path.name}")
        return
    
    preview_txt = preview_path.read_text(encoding="utf-8").strip()
    description_txt = description_path.read_text(encoding="utf-8").strip()
    
    if not story_txt or not preview_txt or not description_txt:
        print(f"❌ Порожні файли для {story_path.name}")
        return
    
    story_id = hashlib.md5(story_txt.encode()).hexdigest()
    video_title = safe_filename(preview_txt)
    output_dir = final_dir / video_title
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"???? Використовується template: {selected_template}")
    
    # НОВИЙ ВИБІР ФОРМАТУ
    format_choice = select_youtube_video_format()
    
    print(f"???? Генерація аудіо...")
    audio_preview = coqui_tts(preview_txt)
    audio_story = coqui_tts(story_txt)
    final_audio = audio_preview + AudioSegment.silent(duration=500) + audio_story
    
    original_duration = final_audio.duration_seconds
    audio_path = temp_folder / f"{story_id}_voice.wav"
    final_audio.export(audio_path, format="wav")
    
    print(f"???? Розпізнавання мовлення...")
    story_offset = audio_preview.duration_seconds + 0.5
    segments, _ = asr_model.transcribe(str(audio_path), word_timestamps=True)
    words = [{"word": w.word.strip(), "start": w.start, "end": w.end} 
             for seg in segments for w in seg.words if w.start >= story_offset]
    chunks = chunk_words(words)
    
    # Створюємо відео залежно від вибору
    if format_choice == "1":  # Тільки горизонтальне
        print(f"????️ Створення горизонтального відео...")
        create_horizontal_video_with_speedup(
            story_id, preview_txt, audio_preview, final_audio, selected_template,
            horizontal_generator, chunks, speed_config, output_dir, video_title, temp_folder
        )
        
    elif format_choice == "2":  # Тільки вертикальне (повне)
        print(f"???? Створення повного вертикального відео...")
        create_full_vertical_video_with_speedup(
            story_id, story_txt, preview_txt, horizontal_generator, vertical_generator,
            selected_template, coqui_tts, asr_model, chunks, speed_config, 
            output_dir, video_title, temp_folder
        )
        
    elif format_choice == "3":  # Обидва формати
        print(f"???? Створення обох форматів...")
        create_horizontal_video_with_speedup(
            story_id, preview_txt, audio_preview, final_audio, selected_template,
            horizontal_generator, chunks, speed_config, output_dir, video_title, temp_folder
        )
        create_full_vertical_video_with_speedup(
            story_id, story_txt, preview_txt, horizontal_generator, vertical_generator,
            selected_template, coqui_tts, asr_model, chunks, speed_config, 
            output_dir, video_title, temp_folder
        )
    
    # Зберігаємо додаткові файли
    if hasattr(horizontal_generator, 'preview_img_path'):
        Image.open(horizontal_generator.preview_img_path).convert("RGB").save(output_dir / "preview.jpg", "JPEG")
    (output_dir / "description.txt").write_text(description_txt, encoding="utf-8")
    
    # Очищуємо оригінальні файли
    for f in [story_path, preview_path, description_path]:
        f.write_text("", encoding="utf-8")
    
    print(f"\n✅ {story_path.name} успішно оброблено!")
    
    # Telegram preview
    try:
        print(f"???? Надсилання превью в Telegram...")
        preview_jpg_path = output_dir / "preview.jpg"
        if preview_jpg_path.exists():
            asyncio.run(send_video_preview(preview_txt, str(preview_jpg_path)))
            print(f"✅ Превью успішно надіслано в Telegram")
    except Exception as e:
        print(f"❌ Помилка при надсиланні превью: {e}")

def select_video_type():
    """Дозволяє користувачу вибрати тип відео для генерації"""
    print("\n" + "="*60)
    print("???? ВИБІР ТИПУ ВІДЕО")
    print("="*60)
    print("1. ???? YouTube відео (з вибором формату)")
    print("2. ???? TikTok відео (розділені на частини)")
    print("3. ???? Обидва типи")
    print("="*60)
    
    while True:
        choice = input("Ваш вибір (1-3): ").strip()
        if choice in ["1", "2", "3"]:
            return choice
        print("❌ Невірний вибір. Введіть 1, 2 або 3")

def select_template_for_video(template_mode, default_template, preview_generator):
    """Вибирає template для поточного відео залежно від режиму"""
    if template_mode == "fixed":
        return default_template
    elif template_mode == "interactive":
        template_name, _ = preview_generator.select_template()
        return template_name
    else: # auto
        return default_template

def process_youtube_videos(stories_dir, generators, default_template, template_mode, preview_generator, coqui_tts, asr_model, temp_folder, final_dir, speed_config):
    """Обробляє YouTube відео з прискоренням"""
    story_files = list(stories_dir.glob("*.txt"))
    story_files = [f for f in story_files if not f.name.endswith("_preview.txt") and not f.name.endswith("_description.txt")]
    
    if not story_files:
        print("❌ YouTube історій не знайдено")
        return 0
    
    print(f"???? Знайдено {len(story_files)} YouTube історій для обробки")
    
    processed_count = 0
    for i, story in enumerate(story_files, 1):
        print(f"\n{'='*60}")
        print(f"???? [{i}/{len(story_files)}] Обробка YouTube історії → {story.name}")
        print(f"{'='*60}")
        
        try:
            selected_template = select_template_for_video(template_mode, default_template, preview_generator)
            if selected_template is None:
                print("❌ Template не вибрано, пропускаємо відео")
                continue
            
            process_youtube_story(
                story, generators['horizontal'], generators['vertical'], selected_template,
                coqui_tts, asr_model, temp_folder, final_dir, speed_config
            )
            processed_count += 1
        except Exception as e:
            print(f"❌ Помилка з {story.name}: {e}")
    
    return processed_count

def main():
    """Головна функція з вибором типу відео та налаштуванням прискорення"""
    # Запитуємо про вимкнення ПК
    off = int(input("Введіть 1 щоб вимкнути ПК після завершення скрипта: "))
    
    # --- Шляхи ---
    ROOT = Path(r"K:/test_ffmpeg_drama")
    BACKGROUNDS_DIR = ROOT / "backgrounds"
    STORIES_DIR = ROOT / "stories"
    TIKTOK_STORIES_DIR = ROOT / "tiktok_stories"
    TEMP_FOLDER = ROOT / "temp"
    TEMP_TIKTOK_FOLDER = ROOT / "temp_tiktok"
    FINAL_DIR = ROOT / "final_videos"
    FINAL_TIKTOK_DIR = ROOT / "final_tiktok_videos"
    TEMPLATES_DIR = ROOT / "templates"
    FONT_PREVIEW = ROOT / "fonts/NotoSans-ExtraBold.ttf"
    SUB_FONT_PATH = ROOT / "fonts/KOMIKAX_.ttf"
    BEEP_SOUND_PATH = ROOT / "beep.wav"
    
    # --- Ініціалізація папок ---
    BACKGROUNDS_DIR.mkdir(parents=True, exist_ok=True)
    STORIES_DIR.mkdir(parents=True, exist_ok=True)
    TIKTOK_STORIES_DIR.mkdir(parents=True, exist_ok=True)
    FINAL_DIR.mkdir(parents=True, exist_ok=True)
    FINAL_TIKTOK_DIR.mkdir(parents=True, exist_ok=True)
    TEMPLATES_DIR.mkdir(parents=True, exist_ok=True)
    
    # --- Вибір папки з фоновими відео ---
    SOURCE_FOLDER = select_background_folder(BACKGROUNDS_DIR)
    if SOURCE_FOLDER is None:
        print("❌ Неможливо продовжити без фонових відео")
        return
    
    # Вибір типу відео
    video_type = select_video_type()
    
    # Налаштування прискорення відео
    speed_config = get_speed_settings()
    
    # Налаштування TikTok тривалості (якщо потрібно)
    tiktok_min_duration = 63
    tiktok_target_duration = 105
    if video_type in ["2", "3"]:  # TikTok або обидва
        tiktok_min_duration, tiktok_target_duration = get_tiktok_duration_settings()
    
    # --- Параметри ---
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    if DEVICE == "cuda":
        torch.cuda.empty_cache()
        torch.backends.cudnn.benchmark = True
        torch.set_float32_matmul_precision("high")
        torch.set_default_device(DEVICE)
    
    RESOLUTION = (1920, 1080)
    SHORTS_RESOLUTION = (1080, 1920)
    TIKTOK_RESOLUTION = (1080, 1920)
    FADE_DURATION = 1
    SHORTS_TARGET_DURATION = 59
    
    # Очищуємо temp папки залежно від типу відео
    if video_type in ["1", "3"]:  # YouTube або обидва
        shutil.rmtree(TEMP_FOLDER, ignore_errors=True)
        TEMP_FOLDER.mkdir(parents=True, exist_ok=True)
    
    if video_type in ["2", "3"]:  # TikTok або обидва
        shutil.rmtree(TEMP_TIKTOK_FOLDER, ignore_errors=True)
        TEMP_TIKTOK_FOLDER.mkdir(parents=True, exist_ok=True)
    
    # --- Ініціалізація TTS та ASR ---
    MODEL = "tts_models/en/vctk/vits"
    VOICE = "p312"
    tts = TTS(model_name=MODEL, progress_bar=False)
    ASR_MODEL = WhisperModel("small", device=DEVICE, compute_type="float16" if DEVICE == "cuda" else "int8")
    
    # --- Ініціалізація генераторів ---
    try:
        preview_generator = PreviewGenerator(
            templates_dir=TEMPLATES_DIR,
            font_path=FONT_PREVIEW,
            shorts_resolution=SHORTS_RESOLUTION
        )
        
        # Питаємо користувача який template використовувати
        print("\n???? Налаштування template для відео:")
        print("Ви можете:")
        print("1. Вибрати один template для всіх відео")
        print("2. Вибирати template для кожного відео окремо")
        print("3. Використовувати перший доступний template")
        
        template_choice = input("Ваш вибір (1-3): ").strip()
        
        if template_choice == "1":
            DEFAULT_TEMPLATE = preview_generator.select_template()[0]
            TEMPLATE_MODE = "fixed"
            print(f"✅ Вибрано template: {DEFAULT_TEMPLATE}")
        elif template_choice == "2":
            DEFAULT_TEMPLATE = None
            TEMPLATE_MODE = "interactive"
        else:
            DEFAULT_TEMPLATE = list(preview_generator.templates.keys())[0]
            TEMPLATE_MODE = "auto"
            print(f"✅ Буде використовуватись template: {DEFAULT_TEMPLATE}")
            
    except FileNotFoundError as e:
        print(f"❌ Помилка з template: {e}")
        # Створюємо приклад template
        example_template = TEMPLATES_DIR / "template_default_x40_y224_w1837_h751.png"
        if not example_template.exists():
            example_img = Image.new('RGBA', (1920, 1080), (255, 255, 255, 0))
            example_img.save(example_template)
        
        preview_generator = PreviewGenerator(
            templates_dir=TEMPLATES_DIR,
            font_path=FONT_PREVIEW,
            shorts_resolution=SHORTS_RESOLUTION
        )
        DEFAULT_TEMPLATE = "template_default_x40_y224_w1837_h751"
        TEMPLATE_MODE = "auto"
    
    # --- Створення TTS функції з цензурою ---
    coqui_tts = create_tts_with_censoring(tts, TEMP_FOLDER, BEEP_SOUND_PATH, VOICE)
    
    # --- Ініціалізація генераторів відео ---
    generators = {}
    if video_type in ["1", "3"]:  # YouTube або обидва
        generators['horizontal'] = HorizontalVideoGenerator(
            source_folder=SOURCE_FOLDER,
            temp_folder=TEMP_FOLDER,
            preview_generator=preview_generator,
            resolution=RESOLUTION,
            fade_duration=FADE_DURATION
        )
        generators['vertical'] = VerticalVideoGenerator(
            source_folder=SOURCE_FOLDER,
            temp_folder=TEMP_FOLDER,
            preview_generator=preview_generator,
            shorts_resolution=SHORTS_RESOLUTION,
            target_duration=SHORTS_TARGET_DURATION
        )
    
    if video_type in ["2", "3"]:  # TikTok або обидва
        tiktok_tts = create_tts_with_censoring(tts, TEMP_TIKTOK_FOLDER, BEEP_SOUND_PATH, VOICE)
        generators['tiktok'] = TikTokVideoGenerator(
            source_folder=SOURCE_FOLDER,
            temp_folder=TEMP_TIKTOK_FOLDER,
            preview_generator=preview_generator,
            tts_func=tiktok_tts,
            asr_model=ASR_MODEL,
            tiktok_resolution=TIKTOK_RESOLUTION,
            target_part_duration=tiktok_target_duration,
            min_part_duration=tiktok_min_duration
        )
    
    # Показуємо налаштування прискорення
    print(f"\n⚡ НАЛАШТУВАННЯ ПРИСКОРЕННЯ:")
    if speed_config['mode'] == '1':
        print("⚡ Режим: Без прискорення")
    elif speed_config['mode'] == '2':
        print(f"⚡ Режим: Фіксоване прискорення")
        for video_type_name, settings in speed_config.items():
            if video_type_name != 'mode' and settings['enabled']:
                print(f"  • {video_type_name}: {settings['speed']}x")
    elif speed_config['mode'] == '3':
        print(f"⚡ Режим: До цільової тривалості")
        for video_type_name, settings in speed_config.items():
            if video_type_name != 'mode' and settings['enabled']:
                print(f"  • {video_type_name}: до {settings['target_duration']}с")
    elif speed_config['mode'] == '4':
        print(f"⚡ Режим: Інтерактивне налаштування")
    
    # --- Обробка відео ---
    if video_type == "1":  # Тільки YouTube
        print(f"\n???? РЕЖИМ: Тільки YouTube відео")
        print(f"???? Використовуються фонові відео з: {SOURCE_FOLDER.name}")
        process_youtube_videos(STORIES_DIR, generators, DEFAULT_TEMPLATE, TEMPLATE_MODE, preview_generator, coqui_tts, ASR_MODEL, TEMP_FOLDER, FINAL_DIR, speed_config)
        
    elif video_type == "2":  # Тільки TikTok
        print(f"\n???? РЕЖИМ: Тільки TikTok відео")
        print(f"???? Використовуються фонові відео з: {SOURCE_FOLDER.name}")
        print(f"⚙️ Налаштування: мінімум {tiktok_min_duration}с, ціль {tiktok_target_duration}с")
        process_tiktok_videos(TIKTOK_STORIES_DIR, generators['tiktok'], generators.get('horizontal'), DEFAULT_TEMPLATE, TEMPLATE_MODE, preview_generator, tiktok_tts, ASR_MODEL, speed_config)
        
    elif video_type == "3":  # Обидва типи
        print(f"\n???????? РЕЖИМ: YouTube + TikTok відео")
        print(f"???? Використовуються фонові відео з: {SOURCE_FOLDER.name}")
        print(f"⚙️ TikTok налаштування: мінімум {tiktok_min_duration}с, ціль {tiktok_target_duration}с")
        
        # Спочатку YouTube
        youtube_processed = process_youtube_videos(STORIES_DIR, generators, DEFAULT_TEMPLATE, TEMPLATE_MODE, preview_generator, coqui_tts, ASR_MODEL, TEMP_FOLDER, FINAL_DIR, speed_config)
        
        # Потім TikTok
        tiktok_processed = process_tiktok_videos(TIKTOK_STORIES_DIR, generators['tiktok'], generators.get('horizontal'), DEFAULT_TEMPLATE, TEMPLATE_MODE, preview_generator, tiktok_tts, ASR_MODEL, speed_config)
        
        print(f"\n???? Загальний підсумок:")
        print(f"???? YouTube відео оброблено: {youtube_processed}")
        print(f"???? TikTok відео оброблено: {tiktok_processed}")
    
    print(f"\n???? Всі відео оброблено!")
    
    # Вимкнення ПК
    if off == 1:
        print("???? Вимкнення комп'ютера через 3 хвилини...")
        time.sleep(180)
        os.system("shutdown /s /t 1")
    else:
        print("???? Генератор відео завершив роботу успішно!")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n⚠️ Скрипт перервано користувачем")
    except Exception as e:
        print(f"\n❌ Критична помилка: {e}")
        import traceback
        traceback.print_exc()