import asyncio
from contextlib import contextmanager
from dataclasses import dataclass
from functools import lru_cache, wraps
import os
import time
import subprocess
from typing import Dict, Literal, Any, Callable, Iterator
import traceback
import string
import mimetypes
import random
import inspect
import warnings
from fastapi import HTTPException

BackoffType = Literal["exponential", "fixed"]


@lru_cache(maxsize=1)
def get_comfyui_version(comfyui_repo_path: str) -> tuple[int, int, int]:
    try:
        comfyui_version = subprocess.check_output(
            ["git", "describe", "--tags", "--abbrev=0"],
            cwd=comfyui_repo_path,
            stderr=subprocess.STDOUT,
        ).decode("utf-8")
        if "latest" in comfyui_version:
            return 0, 0, 0

        major, minor, patch = map(int, comfyui_version.lstrip("v").split("."))
    except subprocess.CalledProcessError as error:
        if "No tags can describe" in error.output:
            print(
                f"No tags can describe the current commit in '{comfyui_repo_path}'. "
                "This is probably due to the currently wanted ComfyUI is a very early "
                "version where the tags weren't introduced yet."
            )
            return 0, 0, 0
        elif "not a git repository" in error.output:
            print(f"Repository '{comfyui_repo_path}' is not a git repository.")

        print(
            f"{error}\nFailed to get the version of the ComfyUI'{comfyui_repo_path}' ."
        )
        raise error

    return major, minor, patch


def get_impact_subpack_location(
    impact_repo_dir: str,
) -> Literal["custom_nodes", "impact"]:
    try:
        # get the commit that the current HEAD is pointed at
        command_output = subprocess.check_output(
            ["git", "rev-list", "--max-count=1", "HEAD"],
            cwd=impact_repo_dir,
            stderr=subprocess.STDOUT,
        ).decode("utf-8")
        current_commit = command_output.strip()

        # this commits is the change where the subpack location was changed
        # https://github.com/ltdrdata/ComfyUI-Impact-Pack/commit/6be257254e90d29bbd2f4732546c7597df6205aa
        change_commit = "6be257254e90d29bbd2f4732546c7597df6205aa"
        command_output = subprocess.check_output(
            ["git", "merge-base", change_commit, current_commit],
            cwd=impact_repo_dir,
            stderr=subprocess.STDOUT,
        ).decode("utf-8")
        ancestor_commit = command_output.strip()

        if ancestor_commit == change_commit:
            print(
                "Detected Impact Pack version is after the subpack location change. Putting it under 'custom_nodes' directory."
            )
            return "custom_nodes"
        elif ancestor_commit == current_commit:
            print(
                "Detected Impact Pack version is before the subpack location change. Putting it under Impact node directory."
            )
            return "impact"
        else:
            raise ValueError(
                f"Unable to determine Impact subpack location. Got unexpected merge-base result. {ancestor_commit}"
            )
    except Exception:
        print(
            "Got exception while trying to determine Impact subpack location. Defaulting to custom_nodes:"
        )
        traceback.print_exc()
        return "custom_nodes"


def async_lru_cache(*lru_cache_args, **lru_cache_kwargs):
    def async_lru_cache_decorator(async_function):
        @lru_cache(*lru_cache_args, **lru_cache_kwargs)
        def cached_async_function(*args, **kwargs):
            coroutine = async_function(*args, **kwargs)
            return asyncio.ensure_future(coroutine)

        return cached_async_function

    return async_lru_cache_decorator


def patch_onnx_runtime():
    try:
        import onnxruntime as ort
    except ImportError:
        print("Onnx runtime is not installed")
        return

    _default_session_options = ort.capi._pybind_state.get_default_session_options()

    OMP_NUM_THREADS = 8

    def get_default_session_options_new():
        _default_session_options.inter_op_num_threads = OMP_NUM_THREADS
        _default_session_options.intra_op_num_threads = OMP_NUM_THREADS
        return _default_session_options

    ort.capi._pybind_state.get_default_session_options = get_default_session_options_new

    os.environ["OMP_NUM_THREADS"] = str(OMP_NUM_THREADS)


@dataclass
class Timings:
    start: float
    end: float | None = None

    @property
    def took(self):
        return (self.end or time.perf_counter()) - self.start


@contextmanager
def timed() -> Iterator[Timings]:
    start = time.perf_counter()
    timer = Timings(start)
    try:
        yield timer
    finally:
        timer.end = time.perf_counter()


def debug_print(*args, **kwargs):
    print("[debug]", *args, **kwargs)


def get_random_filename_for_mime_type(mime_type: str, num_chars: int = 8) -> str:
    """
    Generate a random filename for a given mime type.
    """
    alphanumeric = list(string.ascii_letters + string.digits)
    name = "".join(random.choice(alphanumeric) for _ in range(num_chars))
    extension = mimetypes.guess_extension(mime_type)
    if extension is None:
        extension = {
            "audio/ogg": ".ogg",
            "audio/opus": ".opus",
            "audio/flac": ".flac",
            "audio/aac": ".aac",
            "audio/m4a": ".m4a",
            "audio/wav": ".wav",
        }.get(mime_type, ".bin")
    return f"{name}{extension}"


def async_retry(
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    backoff_type: BackoffType = "exponential",
    jitter: bool = False,
    retry_http_exceptions: bool = False,
    silent: bool = False,
    pass_exceptions_on_retry: bool = False,
) -> Callable:
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            retries = 0
            retried_exceptions: list[Exception] = []
            accepts_retried_exceptions = (
                "retried_exceptions" in inspect.signature(func).parameters
            )

            if pass_exceptions_on_retry:
                if accepts_retried_exceptions:
                    kwargs["retried_exceptions"] = retried_exceptions
                else:
                    warnings.warn(
                        "The function does not accept retried_exceptions as a parameter, but pass_exceptions_on_retry is True. "
                        "The exceptions will not be passed on retry. Please add a retried_exceptions parameter to the function."
                    )

            while retries < max_retries:
                try:
                    return await func(*args, **kwargs)
                except Exception as e:
                    if not retry_http_exceptions and isinstance(e, HTTPException):
                        raise e
                    elif not getattr(e, "retryable", True):
                        raise e

                    retries += 1
                    if retries == max_retries:
                        if not silent:
                            debug_print(
                                f"Max retries reached. Raising {type(e).__name__}: {e}"
                            )
                        raise e
                    if not silent:
                        debug_print(
                            f"Received {type(e).__name__}. Retrying {retries} of {max_retries}..."
                        )

                    retried_exceptions.append(e)

                    if backoff_type == "exponential":
                        delay = min(base_delay * (2 ** (retries - 1)), max_delay)
                    else:  # fixed
                        delay = min(base_delay, max_delay)

                    if jitter:
                        delay *= random.uniform(0.5, 1.5)

                    await asyncio.sleep(delay)

        return wrapper

    return decorator


def _is_cuda_illegal_memory_access_exception(exception: BaseException) -> bool:
    if not isinstance(exception, RuntimeError) or len(exception.args) != 1:
        return False

    msg = exception.args[0]

    if "CUDA error: an illegal memory access was encountered" in msg:
        return True

    return False


class CUDAIllegalMemoryAccessException(HTTPException):
    """Exception raised when a CUDA illegal memory access is detected."""

    def __init__(self, headers: Dict[str, str] | None = None) -> None:
        super().__init__(
            status_code=503,
            detail="CUDA error: an illegal memory access was encountered",
            headers=headers,
        )
