#!/usr/bin/env python

# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utilities for benchmarking and timing code execution.

This module provides the TimeBenchmark class for measuring execution time
using context managers or decorators in a thread-safe manner.
"""

import threading
import time
from contextlib import ContextDecorator


class TimeBenchmark(ContextDecorator):
    """
    Measures execution time using a context manager or decorator.

    This class supports both context manager and decorator usage, and is thread-safe for multithreaded
    environments.

    Args:
        print: If True, prints the elapsed time upon exiting the context or completing the function.
            Defaults to False.

    Examples:

        Using as a context manager::

            >>> benchmark = TimeBenchmark()
            >>> with benchmark:
            ...     time.sleep(1)
            >>> print(f"Block took {benchmark.result:.4f} seconds")
            Block took approximately 1.0000 seconds

        Using with multithreading::

            import threading

            benchmark = TimeBenchmark()

            def context_manager_example():
                with benchmark:
                    time.sleep(0.01)
                print(f"Block took {benchmark.result_ms:.2f} milliseconds")

            threads = []
            for _ in range(3):
                t1 = threading.Thread(target=context_manager_example)
                threads.append(t1)

            for t in threads:
                t.start()

            for t in threads:
                t.join()

            # Expected output:
            # Block took approximately 10.00 milliseconds
            # Block took approximately 10.00 milliseconds
            # Block took approximately 10.00 milliseconds
    """

    def __init__(self, print=False):
        self.local = threading.local()
        self.print_time = print

    def __enter__(self):
        self.local.start_time = time.perf_counter()
        return self

    def __exit__(self, *exc):
        self.local.end_time = time.perf_counter()
        self.local.elapsed_time = self.local.end_time - self.local.start_time
        if self.print_time:
            print(f"Elapsed time: {self.local.elapsed_time:.4f} seconds")
        return False

    @property
    def result(self):
        return getattr(self.local, "elapsed_time", None)

    @property
    def result_ms(self):
        return self.result * 1e3
