import asyncio
import time
from typing import Any

import fal
from fastapi.responses import StreamingResponse

from serverless_comfy.env import get_requirements
from serverless_comfy.models import (
    ComfyInput,
    ComfyOutput,
    PromptType,
)
from serverless_comfy.runtime import ComfyRuntime


class ComfyServer(
    fal.App,
    ComfyRuntime,
    _scheduler="nomad",
    keep_alive=30,
):
    min_concurrency = 1
    max_concurrency = 40
    scaling_delay = 60

    machine_type = ["GPU-H100", "GPU-H200"]
    requirements = [*get_requirements(), "torch<2.6.0"]

    def setup(self):
        from packaging import version

        import torch
        from prometheus_client import Counter

        torch_version = torch.__version__.split("+")[0]

        if version.parse("2.6.0") <= version.parse(torch_version):
            raise ValueError(
                "Breaking torch behavior change for non safetensor files. "
                "Full message: Weights only load failed. This file can still "
                "be loaded, to do so you have two options, do those steps only if "
                "you trust the source of the checkpoint\u001b[0m. \n\t(1) In PyTorch "
                "2.6, we changed the default value of the `weights_only` argument in "
                "`torch.load` from `False` to `True`. Re-running `torch.load` with "
                "`weights_only` set to `False` will likely succeed, but it can result "
                "in arbitrary code execution. Do it only if you got the file from a "
                "trusted source. (2) Alternatively, to load with `weights_only=True` "
                "please check the recommended steps in the following error message."
                "WeightsUnpickler error: Unsupported global: GLOBAL "
                "ultralytics.nn.tasks.DetectionModel was not an allowed global by "
                "default. Please use "
                "`torch.serialization.add_safe_globals([DetectionModel])` or the "
                "`torch.serialization.safe_globals([DetectionModel])` context manager "
                "to allowlist this global if you trust this class/function.\n\nCheck "
                "the documentation of torch.load to learn more about types accepted "
                "by default with weights_only "
                "https://pytorch.org/docs/stable/generated/torch.load.html. "
                "test workflow: https://fal.ai/dashboard/comfy/gokayfem/face-detailer"
            )

        t0 = time.perf_counter()
        commit_hash = "30c259cac8c08ff8d015f9aff3151cb525c9b702"  # ComfyUI 0.3.76
        repo_name = f"serverless-comfy-{commit_hash}-25"
        self.setup_comfy(
            repo_name=repo_name,
            commit_hash=commit_hash,
        )

        self.start_server()
        t1 = time.perf_counter()
        print(f"Setup took {t1 - t0:.2f} seconds")

        self.custom_node_counter = Counter(
            "custom_node_counter", "Custom node counter", labelnames=["node"]
        )

        # will be populated lazily when first needed by track_custom_nodes_async
        self.node_name_to_custom_node_mapping = None

    async def build_node_name_to_custom_node_mapping(self):
        mapping = {}

        nodes_info = await self.get_nodes_info()
        for k, v in nodes_info.items():
            python_module = v["python_module"]
            is_builtin_module = python_module.startswith(
                "nodes"
            ) or python_module.startswith("comfy_extras.")
            if not is_builtin_module:
                custom_node_name = v["python_module"].replace("custom_nodes.", "")
                mapping[k] = custom_node_name
            else:
                mapping[k] = "builtin"

        return mapping

    async def track_custom_nodes_async(self, prompt: PromptType):
        """
        Asynchronously track custom nodes usage by parsing the workflow and incrementing counters.
        This runs in the background and doesn't block the main inference flow.
        """
        try:
            # Build the mapping lazily if not already built
            if self.node_name_to_custom_node_mapping is None:
                self.node_name_to_custom_node_mapping = (
                    await self.build_node_name_to_custom_node_mapping()
                )

            used_custom_nodes = set()
            for node_data in prompt.values():
                node_class_type = node_data.class_type

                if self.node_name_to_custom_node_mapping[node_class_type] != "builtin":
                    used_custom_nodes.add(
                        self.node_name_to_custom_node_mapping[node_class_type]
                    )

            for custom_node in used_custom_nodes:
                self.custom_node_counter.labels(node=custom_node).inc()

        except Exception as e:
            print(f"[debug][Custom Node Tracking] Error tracking custom nodes: {e}")

    @fal.endpoint("/")
    async def inference(self, input: ComfyInput) -> ComfyOutput:
        e2e_start_time = time.perf_counter()
        asyncio.create_task(self.track_custom_nodes_async(input.prompt))

        result = await self.execute_prompt(input)
        e2e_end_time = time.perf_counter()
        result.timings["inference"] = e2e_end_time - e2e_start_time
        return result

    @fal.endpoint("/cancel")
    async def inference_cancel(self):
        print("got cancel request")
        return await self.interrupt()

    @fal.endpoint("/stream")
    async def inference_stream(self, input: ComfyInput) -> StreamingResponse:
        # NOTE: calling sse is on purpose to make it work with fal-js client
        asyncio.create_task(self.track_custom_nodes_async(input.prompt))

        return await self.execute_prompt_sse(input)

    @fal.endpoint("/stream/cancel")
    async def inference_stream_cancel(self):
        print("got stream cancel request")
        return await self.interrupt()

    @fal.endpoint("/object_info")
    async def node_class_mapping(self) -> dict[str, Any]:
        return await self.get_nodes_info()


if __name__ == "__main__":
    app = fal.wrap_app(ComfyServer)
    app()
