from __future__ import annotations

import contextvars
import logging
import os
import threading
import time
from contextlib import contextmanager
from typing import Any, Iterator


_current_timing: contextvars.ContextVar["TimingCollector | None"] = contextvars.ContextVar(
    "surya_current_timing",
    default=None,
)


def timing_enabled() -> bool:
    return os.getenv("SUYA_TIMING_ENABLED", "true").strip().lower() not in {
        "0",
        "false",
        "no",
        "off",
    }


class TimingCollector:
    def __init__(self, *, request_id: str = "-", batch_size: int = 0) -> None:
        self.request_id = request_id
        self.batch_size = batch_size
        self._lock = threading.Lock()
        self._events: list[dict[str, Any]] = []

    def record(self, name: str, duration_ms: float, **metadata: Any) -> None:
        if not timing_enabled():
            return
        event = {
            "name": name,
            "duration_ms": round(duration_ms, 2),
        }
        if metadata:
            event["metadata"] = {
                key: value
                for key, value in metadata.items()
                if value is not None
            }
        with self._lock:
            self._events.append(event)

    def summary(self) -> list[dict[str, Any]]:
        with self._lock:
            return list(self._events)


def get_current_timing() -> TimingCollector | None:
    return _current_timing.get()


def set_current_timing(collector: TimingCollector | None):
    return _current_timing.set(collector)


def reset_current_timing(token) -> None:
    _current_timing.reset(token)


@contextmanager
def timing_span(name: str, **metadata: Any) -> Iterator[None]:
    collector = get_current_timing()
    if collector is None or not timing_enabled():
        yield
        return
    start = time.perf_counter()
    try:
        yield
    finally:
        collector.record(name, (time.perf_counter() - start) * 1000, **metadata)


def log_timing_summary(
    logger: logging.Logger,
    collector: TimingCollector,
    *,
    message: str = "surya_timing_summary",
) -> None:
    if not timing_enabled():
        return
    logger.info(
        "%s request_id=%s batch_size=%s events=%s",
        message,
        collector.request_id,
        collector.batch_size,
        collector.summary(),
    )