from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple
import time

from .logging import log_scope, info, error, exception

SUPPORTED_BACKENDS = {
    "np": "np",
    "numpy": "np",
    "torch": "torch",
    "pytorch": "torch",
    "jax": "jax",
}


@dataclass
class DispatcherContext:
    backend: str
    dtype: Any
    device: Any
    verbose: bool
    debug_level: bool

    def update_from_args(self, **kwargs):
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)


class Dispatcher:
    def __init__(self):
        self._registry: Dict[str, Tuple[Callable, Optional[Callable]]] = {}

    def register(
        self,
        backend_name: str,
        func: Callable,
        transformer: Optional[Callable] = None,
    ):
        if backend_name.lower() in SUPPORTED_BACKENDS:
            self._registry[SUPPORTED_BACKENDS[backend_name.lower()]] = (
                func,
                transformer,
            )

    def detect_backend(self, arg: Any, backend: Optional[str] = None):
        if backend is not None and backend.lower() in SUPPORTED_BACKENDS:
            target_backend = SUPPORTED_BACKENDS[backend.lower()]
        else:
            module = (
                getattr(arg.__class__, "__module__", "").split(".")[0].lower()
            )
            target_backend = SUPPORTED_BACKENDS.get(module, "np")

        if target_backend not in self._registry and self._registry:
            target_backend = next(iter(self._registry.keys()))

        return DispatcherContext(
            backend=target_backend,
            dtype=getattr(arg, "dtype", None),
            device=getattr(arg, "device", None),
            verbose=True,
            debug_level=False,
        )

    def _cast_value(self, value: Any, ctx: DispatcherContext) -> Any:
        if value is None:
            return None

        if isinstance(value, (list, tuple)):
            return [self._cast_value(v, ctx) for v in value]

        if ctx.backend == "torch":
            import torch
            import numpy as np

            if not torch.is_tensor(value):
                try:
                    return torch.as_tensor(
                        value, dtype=ctx.dtype, device=ctx.device
                    )
                except Exception:
                    tensor = torch.as_tensor(np.copy(value))
                    if ctx.device is not None:
                        tensor = tensor.to(ctx.device)
                    return tensor

            return value

        elif ctx.backend == "jax":
            import jax.numpy as jnp

            return jnp.array(value, dtype=ctx.dtype)

        elif ctx.backend == "np":
            import numpy as np

            if not isinstance(value, np.ndarray) or value.dtype != ctx.dtype:
                return np.array(value, dtype=ctx.dtype)
            return value

        return value

    def cast_values(self, ctx: DispatcherContext, *args) -> List[Any]:
        casted_args = []

        for arg in args:
            casted_args.append(self._cast_value(arg, ctx))

        if len(args) == 1:
            return casted_args[0]

        return casted_args

    def _format_time(self, seconds: float) -> str:
        if seconds < 1:
            return f"{seconds * 1000:.2f} ms"
        if seconds < 60:
            return f"{seconds:.2f} s"
        if seconds < 3600:
            m, s = divmod(seconds, 60)
            return f"{int(m)} min {s:.2f} s"

        h, rem = divmod(seconds, 3600)
        m, s = divmod(rem, 60)
        return f"{int(h)} h {int(m)} min {s:.2f} s"

    def _execute_and_time(
        self, ctx: DispatcherContext, func, transformer, *args, **kwargs
    ):
        if transformer:
            args, kwargs = transformer(*args, **kwargs)

        start_time = time.perf_counter()
        try:
            result = func(*args, **kwargs)
            if ctx.verbose:
                elapsed = time.perf_counter() - start_time
                info(f"Finished in {self._format_time(elapsed)}")
            return result
        except Exception as e:
            if ctx.verbose:
                elapsed = time.perf_counter() - start_time
                exception(
                    f"Failed after {self._format_time(elapsed)} with error: {str(e)}"
                )
            raise e

    def __call__(self, ctx: DispatcherContext, *args, **kwargs):
        try:
            func, transformer = self._registry[ctx.backend]
        except KeyError as ke:
            error(f"Backend `{ctx.backend}` not found.")
            raise ke

        if ctx.verbose:
            msg = f"Dispatch: {func.__name__} [{ctx.backend}]"

            with log_scope(msg, ctx.debug_level):
                return self._execute_and_time(
                    ctx, func, transformer, *args, **kwargs
                )
        else:
            return self._execute_and_time(
                ctx, func, transformer, *args, **kwargs
            )
