import time
from contextlib import contextmanager
from prometheus_client import (
    CollectorRegistry,
    Counter,
    Histogram,
    generate_latest,
)

REGISTRY = CollectorRegistry(auto_describe=True)

AGENT_REQUESTS_TOTAL = Counter(
    "raggr_agent_requests_total",
    "Agent generation requests by channel and outcome.",
    labelnames=("channel", "result"),
    registry=REGISTRY,
)
AGENT_RESPONSE_SECONDS = Histogram(
    "raggr_agent_response_seconds",
    "Agent generation time by channel, excluding delivery.",
    labelnames=("channel",),
    buckets=(0.5, 1, 2, 5, 10, 30, 60, 120, 300),
    registry=REGISTRY,
)

CHAT_RESPONSE_SECONDS = Histogram(
    "raggr_chat_response_seconds",
    "End-to-end chat response time in seconds.",
    labelnames=("stream",),
    buckets=(0.5, 1, 2, 5, 10, 30, 60, 120, 300),
    registry=REGISTRY,
)

CHAT_REQUESTS_TOTAL = Counter(
    "raggr_chat_requests_total",
    "Total chat requests by stream mode and result.",
    labelnames=("stream", "result"),
    registry=REGISTRY,
)


def _s(stream: bool) -> str:
    return "true" if stream else "false"


@contextmanager
def time_response(stream: bool):
    t0 = time.monotonic()
    result = "success"
    try:
        yield
    except Exception:
        result = "error"
        raise
    finally:
        CHAT_RESPONSE_SECONDS.labels(stream=_s(stream)).observe(time.monotonic() - t0)
        CHAT_REQUESTS_TOTAL.labels(stream=_s(stream), result=result).inc()


def observe_response(stream: bool, seconds: float) -> None:
    CHAT_RESPONSE_SECONDS.labels(stream=_s(stream)).observe(seconds)


def count_result(stream: bool, result: str) -> None:
    CHAT_REQUESTS_TOTAL.labels(stream=_s(stream), result=result).inc()


def render_metrics() -> bytes:
    return generate_latest(REGISTRY)