import websocket
import json
from urllib import request
import uuid
from pathlib import Path
# Only used for image uploads, for other requests it mangles the JSON and I don't care to figure out why
import requests
# For seed generation
import random
# For strftime
import time
# For ffmpeg video concatenation
import subprocess
import argparse
from itertools import product
from timeit import default_timer as timer
from datetime import timedelta

# Default to localhost and the default port
endpoint = "127.0.0.1"
port = 8188
# Always generate a random client ID, it doesn't really matter what's used
client_id = str(uuid.uuid4())

# Upload a new image for use in I2V generation
def upload_image(image_path):
    with open(image_path, "rb") as f:
        body = { "image": f }
        data = { "overwrite": "true" }

        resp = requests.post("http://{}:{}/upload/image".format(endpoint, port), files=body, data=data)
        if resp.status_code == 200:
            data = resp.json()
            path = data["name"]
            if "subfolder" in data:
                if data["subfolder"] != "":
                    path = data["subfolder"] + "/" + path
        else:
            print(f"{resp.status_code} - {resp.reason}")
    return path

# Queue a prompt for generation
# Note that prompt here is serialized json defining the entire workflow and inputs to it
def queue_prompt(prompt):
    p = { "prompt": prompt, "client_id": client_id }

    data = json.dumps(p).encode('utf-8')
    req = request.Request("http://{}:{}/prompt".format(endpoint, port), data=data)
    return json.loads(request.urlopen(req).read())['prompt_id']

# Use the websocket to check if a prompt has finished
def wait_for_exec_done(ws, prompt_id):
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    # Done.
                    break
        else:
            # Binary data (preview images)
            continue
    return

# Fetch output for a given job
def get_history(prompt_id):
    with request.urlopen("http://{}:{}/history/{}".format(endpoint, port, prompt_id)) as response:
        return json.loads(response.read())[prompt_id]

# Run a job to generate a video, return results
def run_job(ws, mapping_json, job_json, image_path, comfy_output_root):
    last_frame_sentinel = "LASTFRAME"
    # ComfyUI/output/video/YYYY-MM-DD/HHMMSS.mp4
    prefix = Path("video") / time.strftime("%Y-%m-%d") / time.strftime("%H%M%S")
    prefix = str(prefix)
    # Generate a new random seed every time
    seed = random.randint(1, 999999)
    job_json[mapping_json["seed"]]["inputs"]["seed"] = seed
    # Upload image, replace initial image reference with the one we just uploaded
    job_image = upload_image(image_path)
    job_json[mapping_json["input_image"]]["inputs"]["image"] = job_image
    # Video file destination
    job_json[mapping_json["video_dest"]]["inputs"]["filename_prefix"] = prefix
    # Extract last frame destination
    job_json[mapping_json["last_frame_dest"]]["inputs"]["filename_prefix"] = prefix + last_frame_sentinel
    # Queue up our prompt, wait for it to finish
    job_id = queue_prompt(job_json)
    wait_for_exec_done(ws, job_id)
    # Fetch the video file path and last frame path
    history = get_history(job_id)
    outputs = {}
    for node_id in history['outputs']:
        node_output = history['outputs'][node_id]
        images_output = []
        # Images will contain both the extracted last frame and a preview image
        if 'images' in node_output:
            for image in node_output['images']:
                if last_frame_sentinel in image["filename"]:
                    outputs["lastframe"] = comfy_output_root / image["subfolder"] / image["filename"]
        # Gifs will contain only the generated video file
        if 'gifs' in node_output:
            video = node_output["gifs"][0]
            outputs["video"] = comfy_output_root / video["subfolder"] / video["filename"]
    return outputs

# Use ffmpeg to concatenate the generated videos
# If using Windows you will either need to add ffmpeg to PATH or add the path to ffmpeg.exe here
def concatenate_videos(video_paths, output_path):
    list_file = output_path.with_suffix(".txt")

    with open(list_file, "w") as f:
        for vp in video_paths:
            f.write(f"file '{vp.resolve()}'\n")

    cmd = [
        "ffmpeg",
        "-y",
        "-f", "concat",
        "-safe", "0",
        "-i", str(list_file),
        "-c", "copy",
        str(output_path)
    ]

    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    list_file.unlink()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", dest="initial_image", help="Initial image to generate video from", required=True)
    parser.add_argument("-o", "--output", dest="output_path", default="combined", help="Output path", type=Path)
    parser.add_argument("-c", "--config", dest="config_json", help="JSON file containing prompt and stages configuration")
    parser.add_argument("--delete", action="store_true", help="Delete the noncombined files")
    parser.add_argument("--comfyout", default="ComfyUI/output", help="Path to ComfyUI's output root directory", type=Path)
    parser.add_argument("--apiroot", default="Autogenerate", help="Path to exported JSON API files from the workflow", type=Path)

    args = parser.parse_args()

    with open(args.config_json, "rb") as f:
        d = f.read()
        json_stages = json.loads(d)
    stages = []
    repeats = []
    # See the config JSON file for comments about how the prompts are constructed here
    mapping_json = json_stages["mapping"]
    # Validate mapping JSON
    required_mapping_keys = [
        "positive_prompt",
        "negative_prompt",
        "seed",
        "input_image",
        "video_dest",
        "last_frame_dest"
    ]
    if not all(key in mapping_json for key in required_mapping_keys):
        print("Not all required keys are present under mapping in config JSON!")
        return -2
    for stage in json_stages["stages"]:
        with open(args.apiroot / stage["template"], "rb") as f:
            d = f.read()
            s = json.loads(d)
            s[mapping_json["positive_prompt"]]["inputs"]["text"] = stage["positive_prompt"]
            s[mapping_json["negative_prompt"]]["inputs"]["text"] = stage["negative_prompt"]
            stages.append(s)
            if "repeats" in stage:
                repeats.append(int(stage["repeats"]))
            else:
                repeats.append(1)
    
    outputs = {}
    for i in range(len(stages)):
        outputs[i] = []
    lastframe = None
    # Verify input, create paths for output as needed
    if not args.comfyout.exists() or not args.comfyout.is_dir():
        print("ComfyUI output directory does not exist or isn't a directory!")
        return -1
    args.output_path.mkdir(parents=True, exist_ok=True)
    # ComfyUI uses a websocket to communicate generation progress
    ws = websocket.WebSocket()
    ws.connect("ws://{}:{}/ws?clientId={}".format(endpoint, port, client_id))
    try:
        print("Generating video using image {}".format(args.initial_image))
        start_time = timer()
        for i in range(len(stages)):
            for j in range(repeats[i]):
                print("Generating stage {} ({}/{})...".format(i + 1, j + 1, repeats[i]))
                gen_start_time = timer()
                if lastframe == None:
                    stage_out = run_job(ws, mapping_json, stages[i], args.initial_image, args.comfyout)
                else:
                    stage_out = run_job(ws, mapping_json, stages[i], lastframe, args.comfyout)
                outputs[i].append(stage_out)
                gen_end_time = timer()
                print("Generation time: {}".format(timedelta(seconds=gen_end_time-gen_start_time)))
            lastframe = stage_out["lastframe"]
        values = outputs.values()
        permutations = list(product(*values))
        for permutation in permutations:
            # Output filename is based on Stage1 generation time
            concatenate_videos([p["video"] for p in permutation], Path(args.output_path / (str(permutation[-1]["video"].stem) + "_combined.mp4")))
        # Delete the intermediate files if requested
        if args.delete:
            for output_set in outputs.values():
                for output in output_set:
                    output["lastframe"].unlink()
                    output["video"].unlink()
        end_time = timer()
        print("Total time: {}".format(timedelta(seconds=end_time-start_time)))
    finally:
        # Always close the websocket
        ws.close()

if __name__ == "__main__":
    main()
