import threading
import time
from functools import wraps
from typing import Optional, Union, Any, Sequence

from .flags import flags
from .metric import transform_tensor_metrics, set_metric_kernels
from .state import metadata_state
from triton._C.libproton import proton as libproton

thread_local_scopes = threading.local()

MetricValueType = Union[float, int, Sequence[float], Sequence[int]]


class scope:
    """
    A context manager and decorator for entering and exiting a scope.

    Usage:
        context manager:
        ```python
        with proton.scope("test0", {metric_name: metric_value}):
            foo[1,](x, y)
        ```

        decorator:
        ```python
        @proton.scope("test0", {metric_name: metric_value})
        def foo(x, y):
            ...
        ```

    Args:
        name (str): The name of the scope.
        metrics (dict[str, float], optional): The metrics of the scope. Default is None.
    """

    def __init__(self, name: str, metrics: Optional[dict[str, Any]] = None) -> None:
        self.name = name
        self.metrics = metrics
        self.id = None

    def _enter_scope(self):
        if not flags.profiling_on:
            return
        self.id = libproton.record_scope()
        libproton.enter_scope(self.id, self.name)
        if self.metrics:
            with metadata_state():
                set_metric_kernels()
                libproton.add_metrics(self.id, *transform_tensor_metrics(self.metrics))

    def _exit_scope(self):
        if not flags.profiling_on or self.id is None:
            return
        libproton.exit_scope(self.id, self.name)

    def __enter__(self):
        self._enter_scope()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._exit_scope()

    def __call__(self, func):

        @wraps(func)
        def wrapper(*args, **kwargs):
            self._enter_scope()
            try:
                return func(*args, **kwargs)
            finally:
                self._exit_scope()

        return wrapper


class cpu_timed_scope(scope):
    """
    A scope that measures elapsed time (cpu_time).

    Args:
        name (str): The name of the scope.
        metrics (dict[str, float], optional): Additional metrics to add. Default is None.
    """

    def __init__(self, name: str, metrics: Optional[dict[str, Any]] = None) -> None:
        super().__init__(name, metrics)
        self.start_time = None
        if metrics and "cpu_time" in metrics:
            raise ValueError("The metric name 'cpu_time' is reserved.")

    def _enter_scope(self):
        if not flags.profiling_on:
            return
        self.start_time = time.perf_counter_ns()
        super()._enter_scope()

    def _exit_scope(self):
        if not flags.profiling_on:
            return
        if self.start_time is not None:
            cpu_time = time.perf_counter_ns() - self.start_time
            libproton.add_metrics(self.id, {"cpu_time (ns)(exc)": cpu_time})
        super()._exit_scope()


def enter_scope(name: str, *, metrics: Optional[dict[str, Any]] = None) -> Optional[int]:
    if not flags.profiling_on:
        return None
    id = libproton.record_scope()
    thread_local_scopes.scopes = getattr(thread_local_scopes, "scopes", [])
    thread_local_scopes.scopes.append((id, name))
    libproton.enter_scope(id, name)
    if metrics:
        with metadata_state():
            set_metric_kernels()
            libproton.add_metrics(id, *transform_tensor_metrics(metrics))
    return id


def exit_scope(name: Optional[str] = None, *, metrics: Optional[dict[str, Any]] = None) -> Optional[int]:
    # `name` is an optional argument here, only to match the counterpart in enter_scope to make the API consistent with `proton.language.exit_scope`
    if not flags.profiling_on:
        return None
    id, popped_name = thread_local_scopes.scopes.pop()
    if name and name != popped_name:
        raise ValueError(f"Scope name mismatch: {name} != {popped_name}")
    elif not name:
        name = popped_name
    if metrics:
        with metadata_state():
            set_metric_kernels()
            libproton.add_metrics(id, *transform_tensor_metrics(metrics))
    libproton.exit_scope(id, name)
    return id
