import types
import unittest
from unittest import mock

from openai import APIStatusError, RateLimitError

import util.translation as T
from util.translation import AdaptiveLimiter, _header_float, _header_int


class FakeClock:
    """Deterministic stand-in for the ``time`` module.

    ``sleep`` advances the clock instead of waiting, so pacing/backoff paths
    resolve instantly and predictably.
    """

    def __init__(self, start=1000.0):
        self.t = start

    def monotonic(self):
        return self.t

    def time(self):
        return self.t

    def sleep(self, seconds):
        self.t += seconds


def make_exc(cls, **attrs):
    """Build an openai exception without running its real constructor."""
    err = cls.__new__(cls)
    for key, value in attrs.items():
        setattr(err, key, value)
    return err


class FakeRaw:
    def __init__(self, headers, parsed):
        self._headers = headers
        self._parsed = parsed

    @property
    def headers(self):
        return self._headers

    def parse(self):
        return self._parsed


class FakeCreate:
    """Callable that replays a scripted sequence of returns/raises."""

    def __init__(self, behaviors):
        self.behaviors = behaviors
        self.calls = []

    def __call__(self, **params):
        self.calls.append(params)
        behavior = self.behaviors[len(self.calls) - 1]
        if isinstance(behavior, Exception):
            raise behavior
        return behavior


def fake_openai(create):
    completions = types.SimpleNamespace(
        with_raw_response=types.SimpleNamespace(create=create)
    )
    return types.SimpleNamespace(chat=types.SimpleNamespace(completions=completions))


class HeaderHelperTests(unittest.TestCase):
    def test_header_int_first_present_and_floored(self):
        headers = {"b": "125.9"}
        self.assertEqual(_header_int(headers, "a", "b"), 125)

    def test_header_int_missing_blank_and_garbage(self):
        self.assertIsNone(_header_int({}, "a"))
        self.assertIsNone(_header_int({"a": "  "}, "a"))
        self.assertIsNone(_header_int({"a": "abc"}, "a"))

    def test_header_float_keeps_fraction(self):
        self.assertAlmostEqual(_header_float({"a": "0.42"}, "a"), 0.42)


class AdaptiveLimiterUpdateTests(unittest.TestCase):
    def _limiter(self):
        # Seeded at 0.5 RPS / 50k TPM, matching the production default.
        return AdaptiveLimiter(0.5, 50000, 4000)

    def test_canonical_headers_correct_pace_and_budget(self):
        lim = self._limiter()
        lim.update({
            "x-ratelimit-limit-requests": "125",
            "x-ratelimit-limit-tokens": "500000",
            "x-ratelimit-remaining-tokens": "498000",
        })
        self.assertAlmostEqual(lim.req_per_sec, 125 / 60.0)
        self.assertAlmostEqual(lim.min_interval, 60.0 / 125)
        self.assertEqual(lim.tok_per_min, 500000)
        self.assertEqual(lim.tokens_remaining, 498000)

    def test_legacy_minute_spelling_still_works(self):
        lim = self._limiter()
        lim.update({
            "x-ratelimit-limit-req-minute": "25",
            "x-ratelimit-limit-tokens-minute": "60000",
            "x-ratelimit-remaining-tokens-minute": "59000",
        })
        self.assertAlmostEqual(lim.req_per_sec, 25 / 60.0)
        self.assertEqual(lim.tok_per_min, 60000)
        self.assertEqual(lim.tokens_remaining, 59000)

    def test_true_per_second_header_used_undivided(self):
        lim = self._limiter()
        lim.update({"x-ratelimit-limit-req-second": "2"})
        self.assertAlmostEqual(lim.req_per_sec, 2.0)

    def test_empty_or_missing_headers_leave_seed_untouched(self):
        for headers in (None, {}):
            lim = self._limiter()
            lim.update(headers)
            self.assertAlmostEqual(lim.req_per_sec, 0.5)
            self.assertEqual(lim.tok_per_min, 50000)

    def test_garbage_values_ignored(self):
        lim = self._limiter()
        lim.update({"x-ratelimit-limit-requests": "abc"})
        self.assertAlmostEqual(lim.req_per_sec, 0.5)

    def test_exhausted_request_budget_holds_off(self):
        clock = FakeClock()
        with mock.patch.object(T, "time", clock):
            lim = self._limiter()
            lim.update({"x-ratelimit-remaining-requests": "0"})
            self.assertGreaterEqual(lim.next_request_at, clock.monotonic() + 60.0)


class AdaptiveLimiterAcquireTests(unittest.TestCase):
    def test_requests_spaced_by_min_interval(self):
        clock = FakeClock()
        with mock.patch.object(T, "time", clock):
            lim = AdaptiveLimiter(2.0, 10_000_000, 4000)  # min_interval = 0.5s
            lim.acquire(10)  # first clears immediately
            self.assertEqual(clock.t, 1000.0)
            lim.acquire(10)  # second waits one interval
            self.assertAlmostEqual(clock.t, 1000.5)
            self.assertEqual(lim.tokens_remaining, 10_000_000 - 20)

    def test_token_gate_waits_for_window_reset(self):
        clock = FakeClock()
        with mock.patch.object(T, "time", clock):
            lim = AdaptiveLimiter(1000.0, 10000, 2000)
            lim.tokens_remaining = 2500  # near depletion: 2500 - 1000 <= 2000
            original_reset = lim.window_reset
            lim.acquire(1000)
            # It blocked until the minute window rolled over.
            self.assertGreaterEqual(clock.t, original_reset)
            # After the reset the budget refills, then this call is charged.
            self.assertEqual(lim.tokens_remaining, 10000 - 1000)


class CallMistralTests(unittest.TestCase):
    def setUp(self):
        self._saved_limiter = T._mistral_limiter
        T._mistral_limiter = None  # force a fresh limiter from the seeds
        self.clock = FakeClock()
        self._patches = [
            mock.patch.object(T, "time", self.clock),
            mock.patch.dict(T.os.environ, {}, clear=False),
        ]
        for patch in self._patches:
            patch.start()

    def tearDown(self):
        for patch in self._patches:
            patch.stop()
        T._mistral_limiter = self._saved_limiter

    def _params(self, **extra):
        base = {"messages": [{"content": "テスト"}], "max_tokens": 100}
        base.update(extra)
        return base

    def test_success_returns_response_and_syncs_limiter(self):
        sentinel = object()
        create = FakeCreate([
            FakeRaw({"x-ratelimit-limit-requests": "125"}, sentinel),
        ])
        with mock.patch.object(T, "openai", fake_openai(create)):
            result = T.callMistral(self._params())
        self.assertIs(result, sentinel)
        self.assertEqual(len(create.calls), 1)
        self.assertAlmostEqual(T._get_mistral_limiter().req_per_sec, 125 / 60.0)

    def test_rate_limit_then_success_retries(self):
        sentinel = object()
        rate_limited = make_exc(
            RateLimitError,
            response=types.SimpleNamespace(headers={"retry-after": "0.01"}),
        )
        create = FakeCreate([
            rate_limited,
            FakeRaw({"x-ratelimit-limit-requests": "60"}, sentinel),
        ])
        with mock.patch.object(T, "openai", fake_openai(create)):
            result = T.callMistral(self._params())
        self.assertIs(result, sentinel)
        self.assertEqual(len(create.calls), 2)

    def test_json_schema_downgraded_to_json_object(self):
        sentinel = object()
        bad_format = make_exc(APIStatusError, status_code=400)
        create = FakeCreate([
            bad_format,
            FakeRaw({}, sentinel),
        ])
        params = self._params(response_format={"type": "json_schema"})
        with mock.patch.object(T, "openai", fake_openai(create)):
            result = T.callMistral(params)
        self.assertIs(result, sentinel)
        self.assertEqual(create.calls[1]["response_format"], {"type": "json_object"})


if __name__ == "__main__":
    unittest.main()
