"""Tests for decorator-based timing instrumentation."""

from __future__ import annotations

import asyncio
import json
import os
import re
import subprocess
import sys
import textwrap
from pathlib import Path
from typing import TYPE_CHECKING
from unittest.mock import Mock

import pytest

import mindroom.timing as timing_module
from mindroom.timing import (
    DispatchPipelineTiming,
    elapsed_ms_between,
    emit_timing_event,
    timed,
    timed_block,
    timing_enabled,
    timing_scope,
)

if TYPE_CHECKING:
    from collections.abc import AsyncIterator


def _mock_timing_logger(monkeypatch: pytest.MonkeyPatch) -> Mock:
    logger = Mock()
    monkeypatch.setattr(timing_module, "logger", logger)
    return logger


def _assert_timing_logged(logger: Mock, label: str, *, scope: str | None = None) -> None:
    logger.debug.assert_called_once()
    assert logger.debug.call_args.args == ("timing_elapsed",)
    assert logger.debug.call_args.kwargs["label"] == label
    assert isinstance(logger.debug.call_args.kwargs["duration_ms"], float)
    assert logger.debug.call_args.kwargs["duration_ms"] >= 0
    if scope is None:
        assert "timing_scope" not in logger.debug.call_args.kwargs
    else:
        assert logger.debug.call_args.kwargs["timing_scope"] == scope


def test_timed_sync_logs_when_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
    """Sync functions should emit a timing log when timing is enabled."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("sync_label")
    def add(left: int, right: int) -> int:
        return left + right

    assert add(2, 3) == 5
    _assert_timing_logged(logger, "sync_label")


def test_timed_sync_logs_on_exception(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Sync functions should still emit a timing log when they raise."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("sync_error_label")
    def fail() -> None:
        msg = "boom"
        raise RuntimeError(msg)

    with pytest.raises(RuntimeError, match="boom"):
        fail()

    _assert_timing_logged(logger, "sync_error_label")


def test_timed_block_logs_when_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
    """Inline timing blocks should emit the same timing event shape as timed()."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    with timed_block("block_label", scope="scope-123", phase="setup"):
        pass

    _assert_timing_logged(logger, "block_label", scope="scope-123")
    assert logger.debug.call_args.kwargs["phase"] == "setup"


def test_timed_block_logs_on_exception(monkeypatch: pytest.MonkeyPatch) -> None:
    """Inline timing blocks should still emit a timing log when they raise."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    def fail() -> None:
        with timed_block("block_error_label"):
            msg = "boom"
            raise RuntimeError(msg)

    with pytest.raises(RuntimeError, match="boom"):
        fail()

    _assert_timing_logged(logger, "block_error_label")


def test_timed_block_does_not_log_when_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
    """Inline timing blocks should stay quiet when timing is disabled."""
    monkeypatch.delenv("MINDROOM_TIMING", raising=False)
    logger = _mock_timing_logger(monkeypatch)

    with timed_block("block_label"):
        pass

    logger.debug.assert_not_called()


def test_timed_block_uses_context_scope(monkeypatch: pytest.MonkeyPatch) -> None:
    """Inline timing blocks should inherit the active timing scope when omitted."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    token = timing_scope.set("scope-456")
    try:
        with timed_block("block_label"):
            pass
    finally:
        timing_scope.reset(token)

    _assert_timing_logged(logger, "block_label", scope="scope-456")


@pytest.mark.asyncio
async def test_timed_async_logs_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Async functions should emit a timing log when timing is enabled."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("async_label")
    async def compute() -> str:
        return "done"

    assert await compute() == "done"
    _assert_timing_logged(logger, "async_label")


@pytest.mark.asyncio
async def test_timed_async_logs_on_exception(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Async functions should still emit a timing log when they raise."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("async_error_label")
    async def fail() -> None:
        msg = "boom"
        raise RuntimeError(msg)

    with pytest.raises(RuntimeError, match="boom"):
        await fail()

    _assert_timing_logged(logger, "async_error_label")


@pytest.mark.asyncio
async def test_timed_async_generator_logs_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Async generators should emit a timing log after iteration completes."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("async_generator_label")
    async def generate() -> AsyncIterator[str]:
        yield "a"
        yield "b"

    assert [item async for item in generate()] == ["a", "b"]
    _assert_timing_logged(logger, "async_generator_label")


@pytest.mark.asyncio
async def test_timed_async_generator_logs_on_early_close(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Async generators should emit a timing log when iteration stops early."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("async_generator_early_close_label")
    async def generate() -> AsyncIterator[str]:
        yield "a"
        yield "b"

    generator = generate()
    assert await anext(generator) == "a"
    await generator.aclose()

    _assert_timing_logged(logger, "async_generator_early_close_label")


@pytest.mark.asyncio
async def test_timed_async_generator_uses_explicit_scope_on_cross_task_close(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Explicit timing_scope kwargs should survive async-generator close in another task."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("async_generator_cross_task_close_label")
    async def generate(*, timing_scope: str | None = None) -> AsyncIterator[str]:
        del timing_scope
        yield "a"
        yield "b"

    generator = generate(timing_scope="scope-123")
    assert await anext(generator) == "a"
    await asyncio.create_task(generator.aclose())

    _assert_timing_logged(logger, "async_generator_cross_task_close_label", scope="scope-123")


@pytest.mark.asyncio
async def test_timed_async_generator_logs_on_exception(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Async generators should emit a timing log when iteration raises."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("async_generator_error_label")
    async def generate() -> AsyncIterator[str]:
        yield "a"
        msg = "boom"
        raise RuntimeError(msg)

    generator = generate()
    assert await anext(generator) == "a"

    with pytest.raises(RuntimeError, match="boom"):
        await anext(generator)

    _assert_timing_logged(logger, "async_generator_error_label")


def test_timed_returns_original_function_when_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
    """Disabled timing should leave the original callable untouched."""
    monkeypatch.delenv("MINDROOM_TIMING", raising=False)

    def original() -> str:
        return "ok"

    wrapped = timed("disabled_label")(original)

    assert wrapped is original


def test_timed_includes_scope_when_set(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The current timing scope should be rendered in the log line."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("scoped_label")
    def run() -> None:
        return None

    token = timing_scope.set("scope-123")
    try:
        run()
    finally:
        timing_scope.reset(token)

    _assert_timing_logged(logger, "scoped_label", scope="scope-123")


def test_timed_logs_omit_scope_when_unset(monkeypatch: pytest.MonkeyPatch) -> None:
    """Unscoped timing logs should not include the optional timing_scope field."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    @timed("plain_label")
    def run() -> None:
        return None

    run()

    _assert_timing_logged(logger, "plain_label")


def test_emit_timing_event_logs_when_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
    """Structured timing events should inherit the current timing scope."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)

    token = timing_scope.set("scope-123")
    try:
        emit_timing_event("custom_timing_event", value=42, ok=True)
    finally:
        timing_scope.reset(token)

    logger.debug.assert_called_once_with(
        "custom_timing_event",
        value=42,
        ok=True,
        timing_scope="scope-123",
    )


def test_timing_enabled_reflects_env(monkeypatch: pytest.MonkeyPatch) -> None:
    """The public timing-enabled helper should mirror the environment switch."""
    monkeypatch.delenv("MINDROOM_TIMING", raising=False)
    assert timing_enabled() is False

    monkeypatch.setenv("MINDROOM_TIMING", "1")
    assert timing_enabled() is True


def test_elapsed_ms_between_rounds_to_one_decimal_place() -> None:
    """Elapsed-millisecond conversion should use the shared one-decimal policy."""
    assert elapsed_ms_between(1.0, 1.23456) == 234.6
    assert elapsed_ms_between(1.0, 1.23456, ndigits=2) == 234.56


def test_elapsed_millisecond_rounding_policy_is_shared() -> None:
    """Production code should not duplicate the elapsed-millisecond rounding policy."""
    source_root = Path(__file__).parents[1] / "src" / "mindroom"
    direct_rounding_patterns = (
        re.compile(r"round\(\s*\([^)]* - [^)]*\)\s*\* 1000,\s*\d+\s*\)", re.DOTALL),
        re.compile(r"round\(\s*[^,)]+ \* 1000,\s*\d+\s*\)"),
    )

    offenders = [
        path.relative_to(source_root).as_posix()
        for path in source_root.rglob("*.py")
        if path.name != "timing.py" and any(pattern.search(path.read_text()) for pattern in direct_rounding_patterns)
    ]

    assert offenders == []


def test_timed_routes_elapsed_logging_through_shared_helper(monkeypatch: pytest.MonkeyPatch) -> None:
    """timed() should delegate timing_elapsed event shaping to emit_elapsed_timing()."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)
    emit_elapsed_timing = Mock()
    monkeypatch.setattr(timing_module, "emit_elapsed_timing", emit_elapsed_timing)

    @timed("delegated_label")
    def run(*, timing_scope: str | None = None) -> None:
        del timing_scope

    run(timing_scope="scope-123")

    emit_elapsed_timing.assert_called_once()
    assert emit_elapsed_timing.call_args.args[0] == "delegated_label"
    assert isinstance(emit_elapsed_timing.call_args.args[1], float)
    assert emit_elapsed_timing.call_args.kwargs["timing_scope"] == "scope-123"
    logger.debug.assert_not_called()


def _timing_probe_labels(
    *,
    module_name: str,
    call_expression: str,
    patch_target: str,
) -> list[str]:
    """Run one isolated subprocess probe and return emitted timing labels."""
    env = os.environ.copy()
    env["MINDROOM_TIMING"] = "1"
    script = textwrap.dedent(
        f"""
        import json
        from unittest.mock import Mock, patch

        import mindroom.timing as timing_module

        timing_module.logger = Mock()

        import {module_name} as target_module

        with patch("{patch_target}", return_value=object()):
            {call_expression}

        labels = [
            call.kwargs["label"]
            for call in timing_module.logger.debug.call_args_list
            if call.args == ("timing_elapsed",)
        ]
        print(json.dumps(labels))
        """,
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        check=True,
        capture_output=True,
        text=True,
        env=env,
    )
    return json.loads(result.stdout.strip().splitlines()[-1])


def test_load_agent_model_instance_preserves_prompt_assembly_subspan() -> None:
    """Agent model init should keep its dedicated system-prompt timing label."""
    labels = _timing_probe_labels(
        module_name="mindroom.agents",
        call_expression='target_module._load_agent_model_instance(object(), object(), "agent-model")',
        patch_target="mindroom.agents.model_loading.get_model_instance",
    )
    assert labels == ["system_prompt_assembly.agent_create.model_instance"]


def test_load_compaction_model_preserves_history_prepare_subspan() -> None:
    """History compaction model init should keep its dedicated timing label."""
    labels = _timing_probe_labels(
        module_name="mindroom.history.runtime",
        call_expression='target_module._load_compaction_model(object(), object(), "compaction-model")',
        patch_target="mindroom.history.runtime.model_loading.get_model_instance",
    )
    assert labels == ["system_prompt_assembly.history_prepare.compaction_model_init"]


def test_dispatch_pipeline_summary_emits_additive_segments_and_diagnostics() -> None:
    """Pipeline summary should separate additive segments from drill-down diagnostics."""
    logger = Mock()
    timing = DispatchPipelineTiming(source_event_id="$event", room_id="!room")
    timing.metadata["first_visible_kind"] = "stream_update"
    timing.metadata["first_substantive_kind"] = "stream_update"
    timing.marks.update(
        {
            "message_received": 0.0,
            "ingress_cache_append_start": 0.2,
            "ingress_cache_append_ready": 0.7,
            "ingress_normalize_start": 0.9,
            "ingress_normalize_ready": 1.4,
            "gate_enter": 1.0,
            "gate_exit": 3.0,
            "dispatch_prepare_start": 4.0,
            "dispatch_prepare_ready": 6.0,
            "dispatch_plan_start": 6.5,
            "dispatch_plan_ready": 7.5,
            "response_payload_start": 8.0,
            "response_payload_ready": 9.0,
            "lock_wait_start": 10.0,
            "lock_acquired": 12.0,
            "thread_refresh_start": 12.5,
            "thread_refresh_ready": 13.0,
            "response_runtime_start": 14.0,
            "response_runtime_ready": 15.0,
            "ai_prepare_start": 15.5,
            "memory_prepare_start": 15.6,
            "agent_build_start": 15.6,
            "prompt_branches_start": 15.6,
            "memory_prepare_ready": 15.8,
            "agent_build_ready": 16.0,
            "prompt_branches_ready": 16.0,
            "history_classify_start": 16.0,
            "history_classify_ready": 16.2,
            "required_compaction_start": 16.2,
            "required_compaction_ready": 16.6,
            "replay_plan_start": 16.6,
            "replay_plan_ready": 16.8,
            "prompt_assembly_start": 16.8,
            "prompt_assembly_ready": 17.0,
            "history_ready": 17.0,
            "model_request_sent": 18.0,
            "model_first_token": 19.5,
            "first_visible_reply": 20.0,
            "first_substantive_reply": 22.0,
            "streaming_complete": 24.0,
            "response_complete": 25.0,
        },
    )

    timing.emit_summary(logger, outcome="edited")

    logger.debug.assert_called_once()
    assert logger.debug.call_args.args == ("Dispatch pipeline timing",)
    summary = logger.debug.call_args.kwargs
    assert summary["first_visible_kind"] == "stream_update"
    assert summary["first_substantive_kind"] == "stream_update"
    assert summary["seg_ingress_ms"] == 1000.0
    assert summary["seg_coalescing_ms"] == 2000.0
    assert summary["seg_dispatch_ms"] == 7000.0
    assert summary["seg_response_queue_ms"] == 2000.0
    assert summary["seg_first_visible_reply_ms"] == 8000.0
    assert summary["seg_after_first_visible_ms"] == 5000.0
    assert summary["time_to_first_visible_reply_ms"] == 20000.0
    assert summary["time_to_first_substantive_reply_ms"] == 22000.0
    assert summary["total_pipeline_ms"] == 25000.0
    assert summary["diag_ingress_cache_append_ms"] == 500.0
    assert summary["diag_ingress_normalize_ms"] == 500.0
    assert summary["diag_dispatch_prepare_ms"] == 2000.0
    assert summary["diag_dispatch_plan_ms"] == 1000.0
    assert summary["diag_response_payload_setup_ms"] == 1000.0
    assert summary["diag_thread_refresh_ms"] == 500.0
    assert summary["diag_lock_wait_ms"] == 2000.0
    assert summary["diag_runtime_prepare_ms"] == 1000.0
    assert summary["diag_llm_prepare_ms"] == 1500.0
    assert summary["diag_prompt_branch_join_ms"] == 400.0
    assert summary["diag_memory_prepare_ms"] == 200.0
    assert summary["diag_agent_build_ms"] == 400.0
    assert summary["diag_history_classify_ms"] == 200.0
    assert summary["diag_required_compaction_ms"] == 400.0
    assert summary["diag_replay_plan_ms"] == 200.0
    assert summary["diag_prompt_assembly_ms"] == 200.0
    assert summary["diag_history_ready_to_model_request_ms"] == 1000.0
    assert summary["diag_provider_ttft_ms"] == 1500.0
    assert summary["diag_first_visible_to_first_substantive_reply_ms"] == 2000.0
    assert summary["diag_model_first_token_to_first_substantive_reply_ms"] == 2500.0
    assert summary["diag_first_visible_to_stream_complete_ms"] == 4000.0
    assert summary["diag_model_request_to_completion_ms"] == 7000.0
    assert "model_first_token_to_first_visible_stream_update_ms" not in summary
    assert "placeholder_visible_ms" not in summary
    assert "model_request_to_completion_ms" not in summary


def test_dispatch_pipeline_visible_reply_milestones_are_first_write_wins(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Placeholder and substantive reply marks should preserve their earliest milestones."""
    perf_counter = Mock(side_effect=[1.5, 9.0])
    monkeypatch.setattr(timing_module.time, "perf_counter", perf_counter)
    timing = DispatchPipelineTiming(source_event_id="$event", room_id="!room")

    timing.mark_first_visible_reply("placeholder")
    timing.mark_first_visible_reply("stream_update", substantive=True)
    timing.mark_first_visible_reply("final", substantive=True)

    assert timing.marks["first_visible_reply"] == 1.5
    assert timing.metadata["first_visible_kind"] == "placeholder"
    assert timing.marks["first_substantive_reply"] == 9.0
    assert timing.metadata["first_substantive_kind"] == "stream_update"
    assert perf_counter.call_count == 2


def test_dispatch_pipeline_placeholder_only_omits_substantive_reply_metrics(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A failed turn with only a placeholder should not report words as delivered."""
    monkeypatch.setattr(timing_module.time, "perf_counter", Mock(return_value=1.5))
    logger = Mock()
    timing = DispatchPipelineTiming(source_event_id="$event", room_id="!room")
    timing.marks["message_received"] = 0.0

    timing.mark_first_visible_reply("placeholder")
    timing.emit_summary(logger, outcome="failed")

    summary = logger.debug.call_args.kwargs
    assert summary["time_to_first_visible_reply_ms"] == 1500.0
    assert "first_substantive_reply" not in timing.marks
    assert "first_substantive_kind" not in summary
    assert "time_to_first_substantive_reply_ms" not in summary
    assert "diag_first_visible_to_first_substantive_reply_ms" not in summary


def test_emit_timing_event_builds_no_payload_when_debug_is_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """High-frequency timing events must not be assembled just to be dropped downstream."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)
    logger.isEnabledFor.return_value = False

    emit_timing_event("dispatch delivery timing", phase="queued", queue_size=7, progress_hint=True)

    logger.debug.assert_not_called()


def test_emit_timing_event_still_emits_when_debug_is_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The level guard must not change what an enabled timing event reports."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)
    logger.isEnabledFor.return_value = True

    emit_timing_event("dispatch delivery timing", phase="queued", queue_size=7, dropped=None)

    logger.debug.assert_called_once()
    assert logger.debug.call_args.args == ("dispatch delivery timing",)
    assert logger.debug.call_args.kwargs == {"phase": "queued", "queue_size": 7}


def test_dispatch_pipeline_summary_builds_no_payload_when_debug_is_disabled() -> None:
    """The end-to-end summary walks every span pair, so it must not run when dropped."""
    logger = Mock()
    logger.isEnabledFor.return_value = False
    timing = DispatchPipelineTiming(
        source_event_id="$event",
        room_id="!room",
        marks={"message_received": 1.0, "response_complete": 2.0},
    )

    timing.emit_summary(logger, outcome="edited")

    logger.debug.assert_not_called()
    assert timing.summary_emitted


def test_emit_timing_event_emits_when_the_logger_has_no_level_api(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Loggers without the stdlib level API must keep receiving timing events."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = Mock(spec_set=["debug"])
    monkeypatch.setattr(timing_module, "logger", logger)

    emit_timing_event("dispatch delivery timing", phase="queued")

    logger.debug.assert_called_once()


def test_emit_timing_event_emits_when_the_level_check_is_unsupported(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A wrapped logger that cannot answer isEnabledFor must not silence the event."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)
    logger.isEnabledFor.side_effect = AttributeError("'ReturnLogger' object has no attribute 'isEnabledFor'")

    emit_timing_event("dispatch delivery timing", phase="queued")

    logger.debug.assert_called_once()


def test_emit_timing_event_emits_when_the_level_check_is_not_callable(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A non-callable compatibility attribute must not break timing events."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)
    logger.isEnabledFor = False

    emit_timing_event("dispatch delivery timing", phase="queued")

    logger.debug.assert_called_once()


def test_emit_timing_event_emits_when_the_level_check_raises(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A failing compatibility level check must not break timing events."""
    monkeypatch.setenv("MINDROOM_TIMING", "1")
    logger = _mock_timing_logger(monkeypatch)
    logger.isEnabledFor.side_effect = RuntimeError("level check failed")

    emit_timing_event("dispatch delivery timing", phase="queued")

    logger.debug.assert_called_once()


def test_dispatch_pipeline_summary_emits_when_the_level_check_is_unsupported() -> None:
    """The summary must survive loggers that cannot answer isEnabledFor."""
    logger = Mock()
    logger.isEnabledFor.side_effect = AttributeError("'ReturnLogger' object has no attribute 'isEnabledFor'")
    timing = DispatchPipelineTiming(
        source_event_id="$event",
        room_id="!room",
        marks={"message_received": 1.0, "response_complete": 2.0},
    )

    timing.emit_summary(logger, outcome="edited")

    logger.debug.assert_called_once()
