import json
import os
import subprocess
import shutil

def remove_all_subfolders(directory):
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        if os.path.isdir(item_path):
            shutil.rmtree(item_path)

def combine_gifs_to_mp4(gif_files, output_file):
    # Build the FFmpeg command for creating a 3x3 grid
    input_str = ' '.join(['-i ' + gif for gif in gif_files])
    filter_str = """
    [0:v][1:v][2:v]hstack=inputs=3[top];
    [3:v][4:v][5:v]hstack=inputs=3[middle];
    [6:v][7:v][8:v]hstack=inputs=3[bottom];
    [top][middle][bottom]vstack=inputs=3[v]
    """
    ffmpeg_cmd = f'ffmpeg {input_str} -filter_complex "{filter_str}" -map "[v]" combined.gif'

    # Execute the FFmpeg command to create the 3x3 grid
    subprocess.run(ffmpeg_cmd, shell=True, check=True)

    # Convert the combined GIF to MP4
    ffmpeg_convert_cmd = 'ffmpeg -i combined.gif -c:v libx264 -pix_fmt yuv420p -crf 23 ' + output_file
    subprocess.run(ffmpeg_convert_cmd, shell=True, check=True)

def move_and_rename(src_folder, dest_folder):
    count = 1

    # Check if the destination folder exists. If not, create it.
    if not os.path.exists(dest_folder):
        os.makedirs(dest_folder)

    # Walk through each subfolder in the source folder
    for dirpath, dirnames, filenames in os.walk(src_folder):
        dirnames.sort()  # Sort directories in place
        for file in sorted(filenames):
            if file.endswith('.gif'):
                src_path = os.path.join(dirpath, file)
                dest_path = os.path.join(dest_folder, f'{count}.gif')
                shutil.copy(src_path, dest_path)
                print(f"Moved {file} to {dest_path}")
                count += 1

# Path to the directory with files
dir_path = "data/share/sd_models/"

# Path to the JSON file to be modified
json_file_path = "config/prompts/prompt_travel_multi_controlnet.json"

def update_json_path(filename):
    # Load the current JSON
    with open(json_file_path, 'r') as json_file:
        data = json.load(json_file)

    # Modify the 'path' value
    data['path'] = f"share/sd_models/{filename}"

    # Write back to the JSON file
    with open(json_file_path, 'w') as json_file:
        json.dump(data, json_file, indent=4)

def run_animatediff_command():
    cmd = ["animatediff", "generate", "-c", "config/prompts/prompt_travel_multi_controlnet.json", "-W", "256", "-H", "384", "-L", "48", "-C", "16"]
    
    # Launch the subprocess
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

    # Continuously read from stdout and stderr to show progress
    while True:
        output = process.stdout.readline()
        if output == '' and process.poll() is not None:
            break
        if output:
            print(output.strip())
    rc = process.poll()
    
    # Print any errors
    errors = process.stderr.readlines()
    for error in errors:
        print(error.strip())

def main():
    # Check all files in the directory
    file_names = sorted([f for f in os.listdir(dir_path) if os.path.isfile(os.path.join(dir_path, f))])
 
    for filename in file_names:
        # Update the 'path' in the JSON file
        update_json_path(filename)
        
        # Run the animatediff command
        run_animatediff_command()

prompt_templates = [

    {'prompt_fixed_ratio': 1.0,
 'head_prompt': 'little girl, pink hair, tall, confident, grace, poise, wide open eyes, surprise, curiosity, joyful smile, happiness',
 'prompt_map': {'0': '',
  '8': '((astonishment, wide eyes))',
  '16': '((joyful smile, happiness))',
  '24': '((excitement, sparkling eyes))',
  '32': '((delight, blushing cheeks))',
  '40': '((laughter, joyful atmosphere))'},
 'tail_prompt': ''},

    {'prompt_fixed_ratio': 1.0,
 'head_prompt': 'little girl, pink hair, tall, confident posture, wide open eyes, wonder, surprise, animated face, range of emotions',
 'prompt_map': {'0': '',
  '8': '((astonishment, wide eyes))',
  '16': '((tears welling up, emotional))',
  '24': '((flowing tears, sadness))',
  '32': '((contorted face, sadness))',
  '40': '((streaming tears, sadness))'},
 'tail_prompt': ''},

    {'prompt_fixed_ratio': 1.0,
 'head_prompt': 'little girl, pink hair, confidence, determination, wide eyes, curiosity, rosy cheeks, button nose, cute pink dress, frills, bows',
 'prompt_map': {'0': '',
  '8': '((confusion, squinting eyes))',
  '16': '((fear, widening eyes))',
  '24': '((terror, crying))',
  '32': '((tears streaming down face))',
  '40': '((intense sobs))'},
 'tail_prompt': ''},

    {'prompt_fixed_ratio': 1.0,
 'head_prompt': 'little girl, pink hair, tall, confident posture, closed eyes, deep thought, intense anger, powerful roar, red face',
 'prompt_map': {'0': '',
  '8': '((contemplation))',
  '16': '((furrowed brows))',
  '24': '((anger))',
  '32': '((wide open mouth))',
  '40': '((powerful roar))'},
 'tail_prompt': ''},
 ]

if __name__ == "__main__":

    for idx, prompt_template in enumerate(prompt_templates):

        # Load the current JSON
        with open(json_file_path, 'r') as json_file:
            data = json.load(json_file)

        # Modify the prompts
        data['head_prompt'] = prompt_template['head_prompt']
        data['prompt_map'] = prompt_template['prompt_map']
        data['tail_prompt'] = prompt_template['tail_prompt']

        # Write back to the JSON file
        with open(json_file_path, 'w') as json_file:
            json.dump(data, json_file, indent=4)
            
        main()

        subfolder_path = "output"  # As per your example
        current_folder = f"{os.getcwd()}/output_{idx+1}"

        move_and_rename(subfolder_path, current_folder)

        path = f"{current_folder}"

        # List of GIFs with "output" folder path
        gif_files = [f"{path}/{i}.gif" for i in range(1, 10)]
        output_file = f"{path}/combined.mp4"

        os.remove(f"{os.getcwd()}/combined.gif")

        combine_gifs_to_mp4(gif_files, output_file)

        remove_all_subfolders(f"{os.getcwd()}/{subfolder_path}")