from __future__ import annotations

import json
from datetime import datetime
from unittest.mock import AsyncMock

import pytest

from aiops.infra.agent.middlewares import time_usage
from aiops.infra.agent.middlewares.time_usage import (
    BatchParseDatetimeToTimestamp,
    BatchParseTimestampToDatetime,
    TimeUsageMiddleware,
)
from tests.infra.agent.middlewares.helpers import DummyModelRequest


class FixedDatetime(datetime):
    @classmethod
    def now(cls, tz=None):  # type: ignore[override]
        return cls(2024, 5, 6, 7, 8, 9)


def test_time_usage_prompt_includes_current_time(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(time_usage, "datetime", FixedDatetime)

    prompt = time_usage.TIME_USAGE_PROMPT()

    assert "# Time Conversion Guide" in prompt
    assert "Current time: 2024-05-06T07:08:09" in prompt
    assert "batch_parse_timestamp_to_datetime" in prompt
    assert "batch_parse_datetime_to_timestamp" in prompt


@pytest.mark.asyncio
async def test_time_usage_middleware_injects_prompt(
    monkeypatch: pytest.MonkeyPatch,
    dummy_model_request,
) -> None:
    middleware = TimeUsageMiddleware()
    request = dummy_model_request(system_prompt="original")
    monkeypatch.setattr(time_usage, "TIME_USAGE_PROMPT", lambda: "TIME_PROMPT")

    handler = AsyncMock(return_value="response")

    result = await middleware.awrap_model_call(request, handler)

    assert result == "response"
    called_request = handler.await_args[0][0]
    assert isinstance(called_request, DummyModelRequest)
    assert called_request.system_prompt == "original\n\nTIME_PROMPT"


def test_time_usage_tools_roundtrip_conversion() -> None:
    ts_to_dt = BatchParseTimestampToDatetime()
    dt_to_ts = BatchParseDatetimeToTimestamp()

    timestamps = ["1700000000", "1700003600"]
    parsed_datetimes = json.loads(ts_to_dt._run(timestamps))

    expected_roundtrip = {v: k for k, v in parsed_datetimes.items()}
    roundtrip = json.loads(dt_to_ts._run(list(parsed_datetimes.values())))

    assert roundtrip == expected_roundtrip