import asyncio
import itertools
import os
import threading

from serverless_comfy.comfy_overrides import patch_overrides
from serverless_comfy.utils import get_comfyui_version
import time
import gc


def prompt_worker(q, e, server):
    from main import execution, logging
    import comfy.model_management

    last_gc_collect = 0
    need_gc = False
    gc_collect_interval = 10.0

    current_time = time.perf_counter()

    while True:
        timeout = 1000.0
        if need_gc:
            timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)

        queue_item = q.get(timeout=timeout)
        if queue_item is not None:
            item, item_id = queue_item
            execution_start_time = time.perf_counter()
            prompt_id = item[1]
            server.last_prompt_id = prompt_id

            extra_data = item[3].copy()
            if len(item) > 5 and isinstance(item[5], dict):
                extra_data.update(item[5])

            e.execute(item[2], prompt_id, extra_data, item[4])

            need_gc = True
            status = execution.PromptQueue.ExecutionStatus(
                status_str="success" if e.success else "error",
                completed=e.success,
                messages=e.status_messages,
            )
            if len(item) > 5:

                def remove_sensitive(prompt):
                    return prompt[:5] + prompt[6:]

                try:
                    q.task_done(
                        item_id,
                        e.history_result,
                        status=status,
                        process_item=remove_sensitive,
                    )
                except TypeError:
                    q.task_done(item_id, e.history_result, status=status)
            else:
                q.task_done(item_id, e.history_result, status=status)
            server.send_sync(
                "executing",
                {"node": None, "prompt_id": prompt_id},
                server.client_id,
            )

            current_time = time.perf_counter()
            execution_time = current_time - execution_start_time
            logging.info("Prompt executed in {:.2f} seconds".format(execution_time))

        flags = q.get_flags() if hasattr(q, "get_flags") else {}
        free_memory = flags.get("free_memory", False)

        if flags.get("unload_models", free_memory):
            comfy.model_management.unload_all_models()
            need_gc = True
            last_gc_collect = 0

        if free_memory:
            e.reset()
            need_gc = True
            last_gc_collect = 0

        if need_gc:
            current_time = time.perf_counter()
            if (current_time - last_gc_collect) > gc_collect_interval:
                comfy.model_management.cleanup_models()
                gc.collect()
                comfy.model_management.soft_empty_cache()
                last_gc_collect = current_time
                need_gc = False

        time.sleep(0.001)


# Based on https://github.com/comfyanonymous/ComfyUI/blob/c2cb8e889b6a6d45eba26a534bf57aff86b320eb/main.py
def start_server(outputs: list):
    from main import (
        args,
        cleanup_temp,
        cuda_malloc_warning,
        execution,
        folder_paths,
        hijack_progress,
        nodes,
        run,
        server,
    )

    import utils.extra_config

    patch_overrides()

    if args.temp_directory:
        temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")
        print(f"Setting temp directory to: {temp_dir}")
        folder_paths.set_temp_directory(temp_dir)
    cleanup_temp()

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    server = server.PromptServer(loop)
    # In earlier versions of ComfyUI, PromptServer didn't initialize prompt_queue
    # on class construction, and a seperate PromptQueue should be created that
    # is attached to the server instance.
    comfyui_version = get_comfyui_version(os.getcwd())
    major, minor, patch = comfyui_version
    if (major, minor, patch) < (0, 3, 35):
        q = execution.PromptQueue(server)
    else:
        q = server.prompt_queue

    # Try all known PromptExecutor signatures (handles v0.3.x through v0.9.x+)
    e = None
    for attempt in [
        lambda: execution.PromptExecutor(server, cache_args={"lru": args.cache_lru, "ram": args.cache_ram}),
        lambda: execution.PromptExecutor(server, cache_type=False, cache_args=None),
        lambda: execution.PromptExecutor(server, lru_size=args.cache_lru),
        lambda: execution.PromptExecutor(server, cache_size=args.cache_lru),
        lambda: execution.PromptExecutor(server),
    ]:
        try:
            e = attempt()
            break
        except TypeError:
            continue
    if e is None:
        raise RuntimeError("Could not create PromptExecutor with any known signature")

    outputs.extend([server, q, e])

    extra_model_paths_config_path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml"
    )
    if os.path.isfile(extra_model_paths_config_path):
        utils.extra_config.load_extra_path_config(extra_model_paths_config_path)

    if args.extra_model_paths_config:
        for config_path in itertools.chain(*args.extra_model_paths_config):
            utils.extra_config.load_extra_path_config(config_path)

    if (major, minor, patch) >= (0, 3, 47):
        loop.run_until_complete(nodes.init_extra_nodes(init_custom_nodes=True))
    else:
        nodes.init_extra_nodes(init_custom_nodes=True)

    cuda_malloc_warning()

    server.add_routes()
    hijack_progress(server)

    threading.Thread(
        target=prompt_worker,
        daemon=True,
        args=(
            q,
            e,
            server,
        ),
    ).start()

    if args.output_directory:
        output_dir = os.path.abspath(args.output_directory)
        print(f"Setting output directory to: {output_dir}")
        folder_paths.set_output_directory(output_dir)

    # These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes
    folder_paths.add_model_folder_path(
        "checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints")
    )
    folder_paths.add_model_folder_path(
        "clip", os.path.join(folder_paths.get_output_directory(), "clip")
    )
    folder_paths.add_model_folder_path(
        "vae", os.path.join(folder_paths.get_output_directory(), "vae")
    )

    if args.input_directory:
        input_dir = os.path.abspath(args.input_directory)
        print(f"Setting input directory to: {input_dir}")
        folder_paths.set_input_directory(input_dir)

    if args.quick_test_for_ci:
        exit(0)

    call_on_start = None
    if args.auto_launch:

        def startup_server(address, port):
            import webbrowser

            if os.name == "nt" and address == "0.0.0.0":
                address = "127.0.0.1"
            webbrowser.open(f"http://{address}:{port}")

        call_on_start = startup_server

    try:
        loop.run_until_complete(
            run(
                server,
                address=args.listen,
                port=args.port,
                verbose=not args.dont_print_server,
                call_on_start=call_on_start,
            )
        )
    except KeyboardInterrupt:
        print("\nStopped server")

    cleanup_temp()
