# ComfyUI模板--Base Auth
import json
import os
import shutil
import subprocess
import uuid
from typing import Dict

import loguru
import modal
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

image = (  # build up a Modal Image to run ComfyUI, step by step
    modal.Image.debian_slim(  # start from basic Linux with Python
        python_version="3.10"
    )
    .apt_install("git")  # install git to clone ComfyUI
    .apt_install("gcc")
    .apt_install("libportaudio2")
    .pip_install("fastapi[standard]==0.115.4")  # install web dependencies
    .pip_install("comfy-cli==1.3.5")  # install comfy-cli
    .pip_install("cos-python-sdk-v5")
    .pip_install("sqlalchemy")
    .pip_install("ultralytics")
    .pip_install("tencentcloud-sdk-python")
    .pip_install("pymysql")
    .pip_install("Pillow")
    .pip_install("ffmpy")
    .pip_install("opencv-python")
    .pip_install("av")
    .pip_install("imageio")
    .pip_install("loguru")
    .pip_install("conformer==0.3.2",extra_options="--no-dependencies")
    .pip_install("einops>0.6.1",extra_options="--no-dependencies")
    .pip_install("openai-whisper")
    .run_commands(  # use comfy-cli to install ComfyUI and its dependencies
        "comfy --skip-prompt install --nvidia --version 0.3.10"
    )
)

image = (
    image.run_commands("comfy node install https://github.com/M1kep/ComfyLiterals")
    .run_commands("comfy node install https://github.com/evanspearman/ComfyMath")
    .run_commands("comfy node install https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved")
    .run_commands("comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-Bowong.git")
    .run_commands("comfy node install https://github.com/crystian/ComfyUI-Crystools")
    .run_commands("comfy node install https://github.com/pythongosssss/ComfyUI-Custom-Scripts")
    .run_commands("comfy node install https://github.com/BennyKok/comfyui-deploy")
    .run_commands("comfy node install https://github.com/yolain/ComfyUI-Easy-Use")
    .run_commands("comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-VideoHelperSuite.git")
    .run_commands("comfy node install https://github.com/jags111/efficiency-nodes-comfyui")
    .run_commands("comfy node install https://github.com/WASasquatch/was-node-suite-comfyui")
    .run_commands("comfy node install https://github.com/rgthree/rgthree-comfy")
    .run_commands("comfy node install https://github.com/cubiq/ComfyUI_essentials")
    .run_commands("comfy node install https://github.com/melMass/comfy_mtb")
    .run_commands("echo 3 && comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI_SparkTTS.git")
    .run_commands("comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-CustomNode.git")
    .run_commands("echo 04111 && comfy node install https://e.coding.net/g-ldyi2063/dev/cosyvoice_comfyui.git")
    .run_commands("echo 4 && comfy node install https://e.coding.net/g-ldyi2063/dev/ComfyUI-LatentSync-Node.git")
    .run_commands(
        "mkdir -p /root/comfy/ComfyUI/models/ComfyUI-CustomNode/model && rm -rf /root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/model && ln -s /root/comfy/ComfyUI/models/ComfyUI-CustomNode/model /root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/model"
    ).run_commands(
        "mkdir -p /root/comfy/ComfyUI/models/ComfyUI-LatentSync-Node/checkpoints && rm -rf /root/comfy/ComfyUI/custom_nodes/ComfyUI-LatentSync-Node/checkpoints && ln -s /root/comfy/ComfyUI/models/ComfyUI-LatentSync-Node/checkpoints /root/comfy/ComfyUI/custom_nodes/ComfyUI-LatentSync-Node/checkpoints"
    ).run_commands(
        "mkdir -p /root/comfy/ComfyUI/models/.cache/torch/hub/checkpoints && mkdir -p /root/.cache/torch/hub/checkpoints && rm -rf /root/.cache/torch/hub/checkpoints && ln -s /root/comfy/ComfyUI/models/.cache/torch/hub/checkpoints /root/.cache/torch/hub/checkpoints"
    ).run_commands(
        "mkdir -p /root/comfy/ComfyUI/models/stabilityai && ln -s /root/comfy/ComfyUI/models/stabilityai /root/comfy/ComfyUI/stabilityai"
    ).run_commands(
        "mkdir -p /root/comfy/ComfyUI/models/CosyVoice-ComfyUI/pretrained_models/CosyVoice-300M && mkdir -p /root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/CosyVoice-300M && rm -rf /root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/CosyVoice-300M && ln -s /root/comfy/ComfyUI/models/CosyVoice-ComfyUI/pretrained_models/CosyVoice-300M /root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/CosyVoice-300M"
    ).run_commands(
        "rm -rf /root/comfy/ComfyUI/models"
    ).run_commands(
        "apt update && apt install -y ffmpeg && ffmpeg -version"
    ).add_local_file("config.yaml", "/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", copy=True
    ).add_local_file("config.py", "/root/comfy/ComfyUI/custom_nodes/cosyvoice_comfyui/pretrained_models/tools/config.py", copy=True
    ).workdir("/root/comfy")

    # Add .run_commands(...) calls for any other custom nodes you want to download
)

app = modal.App(name="highlight-comfyui-s3", image=image)
vol = modal.Volume.from_name("comfyui-model", create_if_missing=True)
bucket_dict = modal.Dict.from_name("aws_s3_bucket", create_if_missing=False)
bucket_input = str(bucket_dict.get("INPUT"))
bucket_output = str(bucket_dict.get("OUTPUT"))
secret = modal.Secret.from_name("aws-s3-secret")
# completed workflows write output images to this directory
output_dir = "/root/comfy/ComfyUI/output"
auth_scheme = HTTPBearer()


@app.cls(
    allow_concurrent_inputs=1,  # allow 10 concurrent API calls
    max_containers=200,
    min_containers=0,
    buffer_containers=0,
    scaledown_window=120,
    # 5 minute container keep alive after it processes an input; increasing this value is a great way to reduce ComfyUI cold start times
    timeout=900,
    gpu=["L4", "T4"],
    cpu=(2,16),
    # memory=(32768, 32768), # (内存预留量, 内存使用上限)
    memory=(20480,40960),
    enable_memory_snapshot=False,
    secrets=[secret, modal.Secret.from_name("web_auth_token")],
    volumes={
        "/root/comfy/ComfyUI/models": vol,
        "/root/comfy/ComfyUI/input_s3": modal.CloudBucketMount(
            bucket_name=bucket_input,
            # bucket_endpoint_url="https://s3.%s.amazonaws.com" % os.environ["AWS_REGION"],
            secret=secret,
            key_prefix="/"
        ),
        "/root/comfy/ComfyUI/output_s3": modal.CloudBucketMount(
            bucket_name=bucket_output,
            # bucket_endpoint_url="https://s3.%s.amazonaws.com" % os.environ["AWS_REGION"],
            secret=secret,
            key_prefix="/"
        ),
    },
)
class ComfyUI:
    @modal.enter()
    def launch_comfy_background(self):
        # starts the ComfyUI server in the background exactly once when the first input is received
        self.session_id = str(uuid.uuid4())
        cmd = "echo client_uuid: {}&&rm -rf /root/comfy/ComfyUI/user&&mkdir -p /root/comfy/ComfyUI/output_s3/logs/{}&&ln -s /root/comfy/ComfyUI/output_s3/logs/{} /root/comfy/ComfyUI/user && comfy launch --background".format(self.session_id,self.session_id, self.session_id)
        subprocess.run(cmd, shell=True, check=True)

    @modal.method()
    def infer(self, workflow_json: str = ""):
        self.poll_server_health()
        self.prompt_uuid = str(uuid.uuid4())
        # runs the comfy run --workflow command as a subprocess
        workflow = json.loads(workflow_json)
        print("Workflow JSON:")
        print(json.dumps(workflow, indent=4, ensure_ascii=False))
        for node in workflow.values():
            if node.get("class_type") == "ComfyUIDeployExternalText":
                if node.get("inputs").get("input_id") == "file_path":
                    file_to_move = node.get("inputs").get("default_value")
                    if file_to_move:
                        os.makedirs(os.path.dirname(file_to_move.replace("input_s3", "input")), exist_ok=True)
                        try:
                            shutil.copy(file_to_move, file_to_move.replace("input_s3", "input"))
                        except:
                            try:
                                print("Try download file to S3 manually")
                                # S3 Fallback
                                import boto3
                                import yaml
                                with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml",encoding="utf-8",mode="r+") as config:
                                    yaml_config = yaml.load(config, Loader=yaml.FullLoader)
                                awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"],
                                                    aws_secret_access_key=yaml_config["aws_access_key"])
                                awss3.meta.client.download_file(bucket_input, file_to_move.split("input_s3/")[1], file_to_move.replace("input_s3", "input"))
                            except:
                                raise Exception("Failed to download file from S3 manually")
                        node["inputs"]["default_value"] = node["inputs"]["default_value"].replace("input_s3", "input")
        with open(f"/root/{self.prompt_uuid}.json", "w", encoding="utf-8") as fi:
            fi.write(json.dumps(workflow, ensure_ascii=False))
        cmd = f"comfy run --workflow /root/{self.prompt_uuid}.json --wait --timeout 890 --verbose"
        subprocess.run(cmd, shell=True, check=True, timeout=895)

        # looks up the name of the output image file based on the workflow
        file_prefix = [
            node.get("inputs")
            for node in workflow.values()
            if node.get("class_type") == "VHS_VideoCombine"
        ][0]["filename_prefix"]

        # returns the image as bytes
        file_list = os.listdir(output_dir)
        # 获取按照文件时间创建排序的列表，默认是按时间升序
        new_file_list = sorted(file_list, key=lambda file: os.path.getctime(os.path.join(output_dir, file)),
                               reverse=True)
        # print("file_list", new_file_list)
        for f in new_file_list:
            if f.startswith(file_prefix):
                os.makedirs(os.path.dirname(os.path.join(output_dir.replace("output", "output_s3"), f)), exist_ok=True)
                try:
                    shutil.copy(os.path.join(output_dir, f), os.path.join(output_dir.replace("output", "output_s3"), f))
                except:
                    try:
                        print("Try move file to S3 manually")
                        # S3 Fallback
                        import boto3
                        import yaml
                        with open("/root/comfy/ComfyUI/custom_nodes/ComfyUI-CustomNode/config.yaml", encoding="utf-8",mode="r+") as config:
                            yaml_config = yaml.load(config, Loader=yaml.FullLoader)
                        awss3 = boto3.resource('s3', aws_access_key_id=yaml_config["aws_key_id"],
                                            aws_secret_access_key=yaml_config["aws_access_key"])
                        awss3.meta.client.upload_file(os.path.join(output_dir, f), bucket_output, f)
                    except:
                        raise Exception("Failed to move file to S3 manually")
                return f

    @modal.fastapi_endpoint(method="POST")
    def api(self, item: Dict, token: HTTPAuthorizationCredentials = Depends(auth_scheme)):
        if token.credentials != os.environ["AUTH_TOKEN"]:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Incorrect bearer token",
                headers={"WWW-Authenticate": "Bearer"},
            )
        try:
            # save this updated workflow to a new file
            new_workflow_file = item["prompt"]
            # run inference on the currently running container
            fname = self.infer.local(new_workflow_file)
            if fname is None:
                raise RuntimeError("Output File not found")
            j = {"status":"success", "file_name": fname}
            loguru.logger.success(j)
            return j
        except Exception as e:
            j = {"status":"fail", "msg": str(e)}
            loguru.logger.error(j)
            return j
        finally:
            print("Purge Garbage By Restart Comfy")
            cmd = "comfy stop"
            try:
                subprocess.run(cmd, shell=True, check=True)
            except:
                pass
            cmd = "comfy launch --background"
            try:
                subprocess.run(cmd, shell=True, check=True)
            except:
                pass

    def poll_server_health(self):
        import socket
        import urllib

        try:
            # dummy request to check if the server is healthy
            req = urllib.request.Request("http://127.0.0.1:8188/system_stats")
            urllib.request.urlopen(req, timeout=5)
            print("ComfyUI server is healthy")
        except (socket.timeout, urllib.error.URLError) as e:
            # if no response in 5 seconds, stop the container; Modal will schedule queued inputs on a new container
            print(f"Server health check failed: {str(e)} restarting ComfyUI...")
            try:
                cmd = "comfy stop"
                try:
                    subprocess.run(cmd, shell=True, check=True)
                except:
                    pass
                cmd = "comfy launch --background"
                try:
                    subprocess.run(cmd, shell=True, check=True)
                except:
                    raise Exception("Failed to launch ComfyUI")
            except:
                modal.experimental.stop_fetching_inputs()
                raise Exception("ComfyUI server is not healthy, restart failed, stopping container")