from __future__ import annotations

import functools
import inspect
import logging
import time
from collections.abc import Callable


def timed(
    name: str | None = None,
    *,
    log_fn: Callable[..., None] | None = None,
    logger: logging.Logger | None = None,
    unit: str = "s",  # "ms" | "s" | "us"
    warn_after_s: float = 5.0,
):
    """
    Mede o tempo de execução e registra via:
      - log_fn(msg), se fornecido
      - senão logger.info(msg), usando por padrão o logger "agent_internal_stt"

    Se dt > warn_after_s, emite WARNING e adiciona "(FUNÇÃO LENTA)".
    """
    def deco(fn):
        label = name or fn.__qualname__
        base_logger = logger or logging.getLogger("agent_internal_stt")

        def fmt(dt: float) -> str:
            if unit == "us":
                return f"{dt * 1_000_000:.0f} µs"
            if unit == "s":
                return f"{dt:.6f} s"
            return f"{dt * 1000:.2f} ms"

        def emit_timing(dt: float) -> None:
            slow = dt > warn_after_s
            msg = f"[timed] {label} levou {fmt(dt)}" + (" (FUNÇÃO LENTA)" if slow else "")

            # Se não passou log_fn, usamos o logger e aí dá pra logar warning de verdade
            if log_fn is None:
                if slow:
                    base_logger.warning(msg)
                else:
                    base_logger.info(msg)
                return

            # Se passou log_fn, tentamos suportar nível sem quebrar compatibilidade
            try:
                # se o log_fn aceitar algo como log_fn(msg, level="WARNING")
                if slow:
                    log_fn(msg, level="WARNING")
                else:
                    log_fn(msg, level="INFO")
            except TypeError:
                # fallback: mantém assinatura antiga log_fn(msg)
                if slow:
                    log_fn(f"[WARN] {msg}")
                else:
                    log_fn(msg)

        if inspect.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def aw(*args, **kwargs):
                t0 = time.perf_counter()
                try:
                    return await fn(*args, **kwargs)
                finally:
                    emit_timing(time.perf_counter() - t0)
            return aw

        @functools.wraps(fn)
        def w(*args, **kwargs):
            t0 = time.perf_counter()
            try:
                return fn(*args, **kwargs)
            finally:
                emit_timing(time.perf_counter() - t0)
        return w

    return deco