import functools
import inspect
import time
from contextlib import contextmanager
from typing import Optional
from app.core.logging import logger


@contextmanager
def benchmark(name: str):
    """
    Context manager to benchmark a block of code.
    Example:
        with benchmark("Database save"):
            await repo.save(data)
    """
    start_time = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start_time
        logger.info(f"[BENCHMARK] {name} completed in {elapsed:.4f} seconds")


def time_it(name: Optional[str] = None):
    """
    Decorator to benchmark synchronous or asynchronous functions.
    Example:
        @time_it("Retrieve exchange rate")
        async def get_rate():
            ...
    """

    def decorator(func):
        label = name or func.__name__

        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            start = time.perf_counter()
            try:
                return await func(*args, **kwargs)
            finally:
                elapsed = time.perf_counter() - start
                logger.info(f"[BENCHMARK] {label} completed in {elapsed:.4f} seconds")

        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            start = time.perf_counter()
            try:
                return func(*args, **kwargs)
            finally:
                elapsed = time.perf_counter() - start
                logger.info(f"[BENCHMARK] {label} completed in {elapsed:.4f} seconds")

        if inspect.iscoroutinefunction(func):
            return async_wrapper
        return sync_wrapper

    return decorator
