import asyncio
from urllib.parse import urlparse
import hashlib
from fal.toolkit.utils.download_utils import _download_file_python
from concurrent.futures import ThreadPoolExecutor
import contextlib
import io
import json
import os
import queue
import tempfile
import threading
import time
import uuid
from collections.abc import AsyncIterator
from contextlib import contextmanager
from functools import partial
from io import BytesIO
from pathlib import Path
from typing import Any
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from fal.toolkit import download_file

try:
    from fal.exceptions.gpu import (
        _is_cuda_oom_exception,
        CUDAOutOfMemoryException,
    )
except ImportError:
    from fal.exceptions._cuda import (
        _is_cuda_oom_exception,
        CUDAOutOfMemoryException,
    )

import httpx
from fastapi import HTTPException
from sse_starlette import EventSourceResponse, ServerSentEvent

from serverless_comfy.cdn import AsyncCDNManager
from serverless_comfy.env import setup_comfy
from serverless_comfy.models import (
    FAL_OUTPUT_TYPE_NODES,
    ComfyInput,
    ComfyOutput,
    ComfyParameterKey,
    FalFile,
    FalInputDevInfoType,
    FalInputsType,
    PrimitiveInputType,
    PromptType,
)
from serverless_comfy.server import start_server
from serverless_comfy.utils import (
    async_lru_cache,
    _is_cuda_illegal_memory_access_exception,
    CUDAIllegalMemoryAccessException,
)


def get_repository():
    repository = os.getenv("FAL_COMFY_REPOSITORY", "fal_v3")
    fallback_repository = os.getenv("FAL_COMFY_FALLBACK_REPOSITORY", "fal")
    return repository, fallback_repository


class ComfyRuntime:
    def setup_comfy(
        self,
        repo_name: str,
        commit_hash: str = "605e64f6d3da44235498bf9103d7aab1c95ef211",
        custom_nodes: list[list[str]] | None = None,
        models: dict[str, list[list[str]]] | None = None,
        custom_node_models: dict[str, list[list[str]]] | None = None,
        should_cleanup_unused_models: bool = True,
    ):
        self.comfy_client_id = "fal-client-id"

        self.repo_path = setup_comfy(
            commit_hash=commit_hash,
            repo_name=repo_name,
            custom_nodes=custom_nodes,
            models=models,
            custom_node_models=custom_node_models,
            should_cleanup_unused_models=should_cleanup_unused_models,
        )

        self.cdn = AsyncCDNManager()
        self.download_pool = ThreadPoolExecutor(max_workers=36)

        return self.repo_path

    def start_server(self, startup_retries: int = 400):
        self.comfy_input_dir = "/tmp/comfy/input"
        self.comfy_output_dir = "/tmp/comfy/output"

        # Comfy actually creates a new directory named 'temp' inside the provided temp
        # directory. So the actual temp directory is '/tmp/comfy/temp', and the images
        # are stored in there.
        self._comfy_temp_dir = "/tmp/comfy"
        self.comfy_temp_dir = f"{self._comfy_temp_dir}/temp"

        for dir in [self.comfy_temp_dir, self.comfy_input_dir, self.comfy_output_dir]:
            os.makedirs(dir, exist_ok=True)

        os.chdir(self.repo_path)
        outputs = []
        threading.Thread(target=start_server, args=(outputs,)).start()

        self.netloc = "127.0.0.1:8188"
        self.server_url = f"http://{self.netloc}"
        self.setup_websocket(startup_retries=startup_retries)

        (self.prompt_server, self.prompt_queue, self.prompt_executor) = outputs

    def setup_websocket(self, startup_retries: int = 50):
        import websocket

        self.ws = websocket.WebSocket()

        for num_try in range(startup_retries):
            time.sleep(6)
            try:
                self.ws.connect(
                    f"ws://{self.netloc}/ws?clientId={self.comfy_client_id}"
                )
            except ConnectionRefusedError as e:
                if num_try == startup_retries - 1:
                    raise e

                print(e)
                print(f"({num_try + 1}/{startup_retries}) Retrying...")

            else:
                print("Connected to websocket")
                break

    async def preprocess_prompt(
        self,
        prompt: PromptType,
        fal_inputs: FalInputsType | None = None,
        fal_inputs_dev_info: FalInputDevInfoType | None = None,
    ):
        temp_dir = tempfile.mkdtemp()

        if fal_inputs:
            await self.modify_prompt_inputs(
                prompt, fal_inputs, fal_inputs_dev_info, download_dir=temp_dir
            )

        return prompt, temp_dir

    async def set_comfy_server_args(self, server_args: dict[str, Any]):
        from comfy.cli_args import args

        for key, value in server_args.items():
            setattr(args, key, value)

    async def execute_prompt(self, input: ComfyInput) -> ComfyOutput:
        await self.preprocess_prompt(
            input.prompt, input.fal_inputs, input.fal_inputs_dev_info
        )
        return await self.execute_prompt_cached(input)

    async def get_nodes_info(self) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.get(f"{self.server_url}/object_info")
            return response.json()

    @async_lru_cache(maxsize=1)
    async def execute_prompt_cached(self, input: ComfyInput) -> ComfyOutput:
        print("Executing the prompt")
        prompt = input.prompt

        await self.set_comfy_server_args(
            {
                "disable_metadata": input.disable_metadata,
            }
        )

        # TODO: should this be here or outside of lru cache?
        prompt, _ = await self.preprocess_prompt(
            input.prompt, input.fal_inputs, input.fal_inputs_dev_info
        )

        message_queue: queue.Queue[str] = queue.Queue()
        callback = self._make_message_queue_callback(message_queue)

        og_send_sync = self.prompt_server.send_sync
        self.prompt_server.send_sync = self._make_send_sync_tap(og_send_sync, callback)

        try:
            queue_response = await self.queue_prompt(prompt)
            prompt_id = queue_response["prompt_id"]
            events = [
                json.loads(event.data)
                async for event in self._collect_workflow_events(
                    prompt_id, message_queue
                )
            ]
        finally:
            self.prompt_server.send_sync = og_send_sync

        executed_events = []
        fal_execution_errors = []
        comfy_execution_errors = []
        fal_node_timings = {}

        for event in events:
            if event["type"] == "executed":
                executed_events.append(event)
            elif event["type"] == "fal-execution-error":
                fal_execution_errors.append(event)
            elif event["type"] == "execution_error":
                comfy_execution_errors.append(event)
            elif event["type"] == "fal-node-timings":
                event_data = event["data"]
                node_id = event_data["node"]
                if node_id in prompt:
                    node_class_name = prompt[node_id].class_type
                    fal_node_timings[f"{node_class_name}[{node_id}]"] = event_data[
                        "took"
                    ]

        node_outputs = {
            event["data"]["node"]: event["data"]["output"]
            for event in executed_events
            if event["data"]["output"] is not None
        }

        if fal_execution_errors:
            from fastapi import HTTPException

            status_code = 400
            errors = [
                fal_execution_error["data"]["error"]
                for fal_execution_error in fal_execution_errors
            ]

            raise HTTPException(status_code, detail=errors)

        if comfy_execution_errors:
            from fastapi import HTTPException

            errors = [
                execution_error["data"]["exception_message"]
                for execution_error in comfy_execution_errors
            ]
            input_dict = input.dict()
            input_fal_inputs = input_dict.get("fal_inputs", {})

            # Check all errors for CUDA OOM
            for error in errors:
                if isinstance(error, str) and _is_cuda_oom_exception(
                    RuntimeError(error)
                ):
                    raise CUDAOutOfMemoryException()
                elif isinstance(
                    error, str
                ) and _is_cuda_illegal_memory_access_exception(RuntimeError(error)):
                    # CUDA Illegal Memory Access causes the server to hand infedinetely, so
                    # similar to the OOM error, we raise an 503 error to trigger a restart of the server.
                    raise CUDAIllegalMemoryAccessException()

            # Use first error for the response
            error = errors[0]

            if isinstance(error, dict):
                error["input"] = input_fal_inputs
            elif isinstance(error, str):
                error = {"error": error, "input": input_fal_inputs}

            raise HTTPException(400, detail=error)

        fal_outputs = {}
        for event in executed_events:
            node_id = event["data"]["node"]

            node_class_name = prompt[node_id].class_type

            if node_class_name not in FAL_OUTPUT_TYPE_NODES:
                continue

            output_name = prompt[node_id].inputs["output_name"]
            fal_outputs[output_name] = event["data"]["output"]

        return ComfyOutput(
            outputs=node_outputs, fal_outputs=fal_outputs, timings=fal_node_timings
        )

    async def execute_prompt_sse(self, input: ComfyInput) -> EventSourceResponse:
        await self.set_comfy_server_args(
            {
                "disable_metadata": input.disable_metadata,
            }
        )

        prompt, _ = await self.preprocess_prompt(
            input.prompt, input.fal_inputs, input.fal_inputs_dev_info
        )

        message_queue: queue.Queue[str] = queue.Queue()
        callback = self._make_message_queue_callback(message_queue)

        og_send_sync = self.prompt_server.send_sync
        self.prompt_server.send_sync = self._make_send_sync_tap(og_send_sync, callback)

        try:
            queue_response = await self.queue_prompt(prompt)
            prompt_id = queue_response["prompt_id"]
        except Exception:
            self.prompt_server.send_sync = og_send_sync
            raise

        async def _event_stream():
            try:
                async for event in self._collect_workflow_events(
                    prompt_id, message_queue
                ):
                    yield event
            finally:
                self.prompt_server.send_sync = og_send_sync

        return EventSourceResponse(_event_stream(), send_timeout=120.0)

    @staticmethod
    def _truncate_filename(file_name: str, max_length: int = 200) -> str:
        """Truncate a filename to fit within filesystem limits.

        Preserves the file extension and replaces the excess stem with a hash
        to avoid collisions.
        """
        if len(file_name) <= max_length:
            return file_name

        name, ext = os.path.splitext(file_name)
        hash_suffix = hashlib.sha256(file_name.encode()).hexdigest()[:16]
        max_name_length = max_length - len(ext) - len(hash_suffix) - 1

        if max_name_length <= 0:
            return hash_suffix[: max_length - len(ext)] + ext

        return name[:max_name_length] + "_" + hash_suffix + ext

    def _download_file_safe(self, url: str, target_dir: str) -> Path:
        """Download a file, truncating the filename if it exceeds filesystem limits."""
        parsed = urlparse(url)
        original_name = Path(parsed.path).name if parsed.scheme != "data" else ""

        if original_name and len(original_name) > 200:
            truncated = self._truncate_filename(original_name)
            target_path = Path(target_dir) / truncated
            target_path.parent.mkdir(parents=True, exist_ok=True)

            _download_file_python(url=url, target_path=target_path)
            return target_path

        return download_file(url=url, target_dir=target_dir)

    def download_fal_file(self, fal_file: FalFile, target_dir: str):
        from fastapi import HTTPException

        try:
            if fal_file.url.startswith("data:"):
                print(f"Starting download of data URL: {fal_file.url[:10]}[TRUNCATED]")
            else:
                print(f"Starting download of URL: {fal_file.url}")

            # suppress the prints
            with contextlib.redirect_stdout(io.StringIO()):
                file_path = self._download_file_safe(fal_file.url, target_dir)
        except Exception as e:
            print(e)

            raise HTTPException(
                status_code=422,
                detail=f"Failed to download the file from {fal_file.url}",
            )

        return str(file_path)

    async def modify_prompt_inputs(
        self,
        prompt: PromptType,
        inputs: FalInputsType,
        inputs_dev_info: FalInputDevInfoType,
        download_dir: str,
    ):
        futures = []

        for input_name, input_value in inputs.items():
            input_key = inputs_dev_info[input_name].key
            class_type = inputs_dev_info[input_name].class_type
            future = self.download_pool.submit(
                self.modify_prompt_node,
                prompt,
                input_key,
                input_value,
                class_type,
                download_dir,
            )
            futures.append(future)

        for future in futures:
            future.result()

    def modify_prompt_node(
        self,
        prompt: PromptType,
        keys: ComfyParameterKey,
        value: PrimitiveInputType,
        class_type: str,
        download_dir: str,
    ):
        load_image_nodes = ["LoadImage"]
        load_video_nodes = ["VHS_LoadVideoPath", "VHS_LoadVideo"]
        load_audio_nodes = ["LoadAudio", "VHS_LoadAudio", "VHS_LoadAudioUpload"]
        load_nodes = load_image_nodes + load_video_nodes + load_audio_nodes

        if class_type in load_nodes:
            value = self.download_fal_file(FalFile(key=keys, url=value), download_dir)

        input_node = prompt
        for key in keys[:-2]:
            input_node = input_node[key]

        input_name = keys[-1]
        input_node.inputs[input_name] = value

    async def _collect_workflow_events(
        self, prompt_id: str, message_queue: queue.Queue[str]
    ) -> AsyncIterator[ServerSentEvent]:
        t0 = time.perf_counter()

        try:
            while True:
                try:
                    raw_message = message_queue.get_nowait()
                except queue.Empty:
                    await asyncio.sleep(0.1)
                    continue

                message = json.loads(raw_message)

                message_type = message["type"]
                message_data = message["data"]

                async for event in self.dispatch_message(message_type, message_data):
                    yield event

                message_prompt_id = message_data.get("prompt_id")
                is_matching_prompt = message_prompt_id == prompt_id
                is_legacy_termination = (
                    message_type == "executing"
                    and is_matching_prompt
                    and message_data.get("node") is None
                )
                is_terminal_event = (
                    message_type
                    in {"execution_success", "execution_error", "execution_interrupted"}
                    and is_matching_prompt
                )

                if is_legacy_termination or is_terminal_event:
                    break

        finally:
            t1 = time.perf_counter()
            yield ServerSentEvent(
                await self.convert_message_to_comfy_format(
                    "fal-info", {"message": f"Execution took {round(t1 - t0, 2)}s"}
                )
            )

    def _make_message_queue_callback(self, message_queue: queue.Queue[str]):
        def _callback(message: str):
            message_queue.put(message)

        return _callback

    def _make_send_sync_tap(self, og_send_sync, callback):
        def _send_sync(event, data, sid=None):
            try:
                message = {"type": event, "data": data}
                callback(json.dumps(message))
            except Exception:
                print(f"[debug] Failed to mirror prompt event type={event} sid={sid}")
            return og_send_sync(event, data, sid)

        return _send_sync

    async def queue_prompt(self, prompt: PromptType):
        deserialized_prompt = {
            node_id: node_data.model_dump() for node_id, node_data in prompt.items()
        }
        client_id = str(uuid.uuid4())
        p = {"prompt": deserialized_prompt, "client_id": client_id}
        data = json.dumps(p).encode("utf-8")

        try:
            request = Request(f"{self.server_url}/prompt", data=data)
            with urlopen(request) as response:
                response_data = response.read()
                return json.loads(response_data)
        except HTTPError as error:
            response_data = error.read().decode("utf-8")
            print("HTTPError", str(response_data))

            try:
                response_json = json.loads(response_data)
            except json.JSONDecodeError:
                response_json = response_data

            raise HTTPException(400, detail=response_json)

        except asyncio.CancelledError:
            error_message = "Disconnected from client"
            print(error_message)
            raise HTTPException(500, detail=error_message)

        except Exception as error:
            print(error)
            error_message = f"An error occurred while executing the prompt. {error}"
            print(error_message)
            raise HTTPException(500, detail=error_message)

    async def dispatch_message(
        self, type: str, data: dict[str, Any]
    ) -> AsyncIterator[ServerSentEvent]:
        dispatchers = {
            "executed": self.dispatch_executed_message,
        }
        dispatcher = dispatchers.get(type, partial(self.dispatch_generic_message, type))

        async for event in dispatcher(data):
            yield event

    async def dispatch_executed_message(
        self, data: dict[str, Any]
    ) -> AsyncIterator[ServerSentEvent]:
        await self.collect_images_from_executed_message(data)
        await self.collect_gifs_from_executed_message(data)
        await self.collect_3d_objects_from_executed_message(data)
        new_message = await self.convert_message_to_comfy_format("executed", data)

        yield ServerSentEvent(
            data=new_message,
        )

    async def dispatch_generic_message(
        self, type: str, data: dict[str, Any]
    ) -> AsyncIterator[ServerSentEvent]:
        yield ServerSentEvent(
            data=await self.convert_message_to_comfy_format(type, data),
        )

    async def convert_message_to_comfy_format(
        self, type: str, data: dict[str, Any]
    ) -> str:
        return json.dumps(
            {
                "type": type,
                "data": data,
            }
        )

    async def collect_images_from_executed_message(
        self, executed_message_data: dict[str, Any]
    ):
        output_data = executed_message_data.get("output", {}) or {}
        output_images = output_data.get("images", []) or []
        images = []

        with ThreadPoolExecutor(max_workers=36) as executor:
            futures = [
                executor.submit(
                    self.preprocess_image,
                    image_metadata["filename"],
                    image_metadata["type"],
                    image_metadata.get("subfolder", None),
                    image_metadata.get("channel", None),
                )
                for image_metadata in output_images
                if "filename" in image_metadata
            ]

            file_paths = [future.result() for future in futures]

        access_urls = await self.cdn.upload_files(
            content_type="image/png",
            file_paths=file_paths,
            wait_for_upload=True,
        )

        for image_metadata, access_url in zip(output_images, access_urls):
            image_metadata["url"] = access_url
            images.append(image_metadata)

        return images

    async def collect_gifs_from_executed_message(
        self, executed_message_data: dict[str, Any]
    ):
        output_data = executed_message_data.get("output", {}) or {}
        output_gifs = output_data.get("gifs", []) or []
        gifs = []

        with ThreadPoolExecutor(max_workers=36) as executor:
            futures = [
                executor.submit(
                    self.preprocess_image,
                    gif_metadata["filename"],
                    gif_metadata["type"],
                    gif_metadata.get("subfolder", None),
                )
                for gif_metadata in output_gifs
                if "filename" in gif_metadata
            ]

            file_paths = [future.result() for future in futures]

        access_urls = await self.cdn.upload_files(
            content_type="image/gif",
            file_paths=file_paths,
            wait_for_upload=True,
        )

        for gif_metadata, access_url in zip(output_gifs, access_urls):
            gif_metadata["url"] = access_url
            gifs.append(gif_metadata)

        return gifs

    async def collect_3d_objects_from_executed_message(
        self, executed_message_data: dict[str, Any]
    ):
        output_data = executed_message_data.get("output", {}) or {}
        output_previews = output_data.get("previews", []) or []
        objects = []

        file_paths = [
            object_metadata["filepath"]
            for object_metadata in output_previews
            if "filepath" in object_metadata
        ]

        access_urls = await self.cdn.upload_files(
            content_type="application/octet-stream",
            file_paths=file_paths,
            wait_for_upload=True,
        )

        for object_metadata, access_url in zip(output_previews, access_urls):
            object_metadata["url"] = access_url
            objects.append(object_metadata)

        return objects

    def get_comfy_dir(self, type: str):
        if type == "input":
            return self.comfy_input_dir
        elif type == "output":
            return self.comfy_output_dir
        elif type == "temp":
            return self.comfy_temp_dir
        else:
            raise ValueError(
                f"Invalid type: {type} for get_comfy_dir. "
                "Valid types are: input, output, temp."
            )

    def get_file_path(self, filename: str, type: str = "output", subfolder: str = None):
        output_dir = self.get_comfy_dir(type)
        if subfolder is not None:
            full_output_dir = os.path.join(output_dir, subfolder)
            if (
                os.path.commonpath((os.path.abspath(full_output_dir), output_dir))
                != output_dir
            ):
                return None
            output_dir = full_output_dir

        filename = os.path.basename(filename)
        return os.path.join(output_dir, filename), output_dir

    def preprocess_image(
        self,
        filename: str,
        type: str = "output",
        subfolder: str = None,
        channel: str = None,
    ):
        from PIL import Image

        file, _ = self.get_file_path(filename, type, subfolder)

        if not os.path.isfile(file):
            print(f"Called preprocess_image with file that does not exist: {file}")

        if os.path.isfile(file):
            channel = channel or "rgba"

            if channel == "rgb":
                with Image.open(file) as img:
                    if img.mode == "RGBA":
                        r, g, b, a = img.split()
                        new_img = Image.merge("RGB", (r, g, b))
                    else:
                        new_img = img.convert("RGB")

                    buffer = BytesIO()
                    new_img.save(buffer, format="PNG")
                    buffer.seek(0)
                    with open(file, "wb") as f:
                        f.write(buffer.read())

            elif channel == "a":
                with Image.open(file) as img:
                    if img.mode == "RGBA":
                        _, _, _, a = img.split()
                    else:
                        a = Image.new("L", img.size, 255)

                    # alpha img
                    alpha_img = Image.new("RGBA", img.size)
                    alpha_img.putalpha(a)
                    alpha_buffer = BytesIO()
                    alpha_img.save(alpha_buffer, format="PNG")
                    alpha_buffer.seek(0)

                    with open(file, "wb") as f:
                        f.write(alpha_buffer.read())

        return file

    async def interrupt(self):
        async with httpx.AsyncClient() as client:
            await client.post(f"{self.server_url}/interrupt", json={})

        return {"interrupt": "ok"}

    async def free(self, unload_models: bool = False, free_memory: bool = False):
        async with httpx.AsyncClient() as client:
            await client.post(
                f"{self.server_url}/free",
                json={"unload_models": unload_models, "free_memory": free_memory},
            )

        return {"free": "ok"}

    @contextmanager
    def register_comfy_model(
        self, absolute_model_path: str, model_dir: str | Path = "checkpoints"
    ):
        model_path_obj = Path(absolute_model_path)
        comfy_model_dir = self.repo_path / "models" / model_dir

        model_name = model_path_obj.name

        model_target_dir = comfy_model_dir / model_name

        try:
            if model_path_obj.is_absolute():
                model_target_dir.symlink_to(model_path_obj)
            yield model_path_obj.name
        finally:
            if model_path_obj.is_absolute():
                model_target_dir.unlink()
