#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#
import os
import tempfile
import time
from datetime import datetime, timedelta
from typing import Any, Iterable, Mapping, Optional

import pytest
import requests
from requests import Request

from airbyte_cdk.models import SyncMode
from airbyte_cdk.sources.streams.call_rate import (
    APIBudget,
    CallRateLimitHit,
    FixedWindowCallRatePolicy,
    HttpRequestMatcher,
    HttpRequestRegexMatcher,
    MovingWindowCallRatePolicy,
    Rate,
    UnlimitedCallRatePolicy,
)
from airbyte_cdk.sources.streams.http import HttpStream
from airbyte_cdk.sources.streams.http.requests_native_auth import TokenAuthenticator
from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH


class StubDummyHttpStream(HttpStream):
    url_base = "https://test_base_url.com"
    primary_key = "some_key"

    def next_page_token(self, response: requests.Response) -> Optional[Mapping[str, Any]]:
        return {"next_page_token": True}  # endless pages

    def path(self, **kwargs) -> str:
        return ""

    def parse_response(self, *args, **kwargs) -> Iterable[Mapping]:
        yield {"data": "some_data"}


class StubDummyCacheHttpStream(StubDummyHttpStream):
    use_cache = True


@pytest.fixture(name="enable_cache")
def enable_cache_fixture():
    prev_cache_path = os.environ.get(ENV_REQUEST_CACHE_PATH)
    with tempfile.TemporaryDirectory(
        # Cleanup can fail on Windows due to file locks. Ignore if so,
        # rather than failing the whole process.
        ignore_cleanup_errors=True,
    ) as temp_dir:
        os.environ[ENV_REQUEST_CACHE_PATH] = temp_dir
        yield

    if prev_cache_path is not None:
        os.environ[ENV_REQUEST_CACHE_PATH] = prev_cache_path


class TestHttpRequestMatcher:
    try_all_types_of_requests = pytest.mark.parametrize(
        "request_factory",
        [Request, lambda *args, **kwargs: Request(*args, **kwargs).prepare()],
    )

    @try_all_types_of_requests
    def test_url(self, request_factory):
        matcher = HttpRequestMatcher(url="http://some_url/")
        assert not matcher(request_factory(url="http://some_wrong_url"))
        assert matcher(request_factory(url="http://some_url"))

    @try_all_types_of_requests
    def test_method(self, request_factory):
        matcher = HttpRequestMatcher(method="GET")
        assert not matcher(request_factory(url="http://some_url"))
        assert not matcher(request_factory(url="http://some_url", method="POST"))
        assert matcher(request_factory(url="http://some_url", method="GET"))

    @try_all_types_of_requests
    def test_params(self, request_factory):
        matcher = HttpRequestMatcher(params={"param1": 10, "param2": 15})
        assert not matcher(request_factory(url="http://some_url/"))
        assert not matcher(
            request_factory(url="http://some_url/", params={"param1": 10, "param3": 100})
        )
        assert not matcher(
            request_factory(url="http://some_url/", params={"param1": 10, "param2": 10})
        )
        assert matcher(
            request_factory(
                url="http://some_url/", params={"param1": 10, "param2": 15, "param3": 100}
            )
        )

    @try_all_types_of_requests
    def test_header(self, request_factory):
        matcher = HttpRequestMatcher(headers={"header1": 10, "header2": 15})
        assert not matcher(request_factory(url="http://some_url"))
        assert not matcher(
            request_factory(url="http://some_url", headers={"header1": "10", "header3": "100"})
        )
        assert not matcher(
            request_factory(url="http://some_url", headers={"header1": "10", "header2": "10"})
        )
        assert matcher(
            request_factory(
                url="http://some_url", headers={"header1": "10", "header2": "15", "header3": "100"}
            )
        )

    @try_all_types_of_requests
    def test_combination(self, request_factory):
        matcher = HttpRequestMatcher(
            method="GET", url="http://some_url/", headers={"header1": 10}, params={"param2": "test"}
        )
        assert matcher(
            request_factory(
                method="GET",
                url="http://some_url",
                headers={"header1": "10"},
                params={"param2": "test"},
            )
        )
        assert not matcher(
            request_factory(method="GET", url="http://some_url", headers={"header1": "10"})
        )
        assert not matcher(request_factory(method="GET", url="http://some_url"))
        assert not matcher(request_factory(url="http://some_url"))


def test_http_request_matching(mocker):
    """Test policy lookup based on matchers."""
    users_policy = mocker.Mock(spec=MovingWindowCallRatePolicy)
    groups_policy = mocker.Mock(spec=MovingWindowCallRatePolicy)
    root_policy = mocker.Mock(spec=MovingWindowCallRatePolicy)

    users_policy.matches.side_effect = HttpRequestMatcher(
        url="http://domain/api/users", method="GET"
    )
    users_policy.get_weight.return_value = 1
    groups_policy.matches.side_effect = HttpRequestMatcher(
        url="http://domain/api/groups", method="POST"
    )
    groups_policy.get_weight.return_value = 1
    root_policy.matches.side_effect = HttpRequestMatcher(method="GET")
    root_policy.get_weight.return_value = 1
    api_budget = APIBudget(
        policies=[
            users_policy,
            groups_policy,
            root_policy,
        ]
    )

    (
        api_budget.acquire_call(
            Request("POST", url="http://domain/unmatched_endpoint"), block=False
        ),
        "unrestricted call",
    )
    users_policy.try_acquire.assert_not_called()
    groups_policy.try_acquire.assert_not_called()
    root_policy.try_acquire.assert_not_called()

    users_request = Request("GET", url="http://domain/api/users")
    api_budget.acquire_call(users_request, block=False), "first call, first matcher"
    users_policy.try_acquire.assert_called_once_with(users_request, weight=1)
    groups_policy.try_acquire.assert_not_called()
    root_policy.try_acquire.assert_not_called()

    (
        api_budget.acquire_call(Request("GET", url="http://domain/api/users"), block=False),
        "second call, first matcher",
    )
    assert users_policy.try_acquire.call_count == 2
    groups_policy.try_acquire.assert_not_called()
    root_policy.try_acquire.assert_not_called()

    group_request = Request("POST", url="http://domain/api/groups")
    api_budget.acquire_call(group_request, block=False), "first call, second matcher"
    assert users_policy.try_acquire.call_count == 2
    groups_policy.try_acquire.assert_called_once_with(group_request, weight=1)
    root_policy.try_acquire.assert_not_called()

    (
        api_budget.acquire_call(Request("POST", url="http://domain/api/groups"), block=False),
        "second call, second matcher",
    )
    assert users_policy.try_acquire.call_count == 2
    assert groups_policy.try_acquire.call_count == 2
    root_policy.try_acquire.assert_not_called()

    any_get_request = Request("GET", url="http://domain/api/")
    api_budget.acquire_call(any_get_request, block=False), "first call, third matcher"
    assert users_policy.try_acquire.call_count == 2
    assert groups_policy.try_acquire.call_count == 2
    root_policy.try_acquire.assert_called_once_with(any_get_request, weight=1)


class TestUnlimitedCallRatePolicy:
    def test_try_acquire(self, mocker):
        policy = UnlimitedCallRatePolicy(matchers=[])
        assert policy.matches(mocker.Mock()), "should match anything"
        policy.try_acquire(mocker.Mock(), weight=1)
        policy.try_acquire(mocker.Mock(), weight=10)

    def test_update(self):
        policy = UnlimitedCallRatePolicy(matchers=[])
        policy.update(available_calls=10, call_reset_ts=datetime.now())
        policy.update(available_calls=None, call_reset_ts=datetime.now())
        policy.update(available_calls=10, call_reset_ts=None)


class TestFixedWindowCallRatePolicy:
    def test_limit_rate(self, mocker):
        policy = FixedWindowCallRatePolicy(
            matchers=[], next_reset_ts=datetime.now(), period=timedelta(hours=1), call_limit=100
        )
        policy.try_acquire(mocker.Mock(), weight=1)
        policy.try_acquire(mocker.Mock(), weight=20)
        with pytest.raises(ValueError, match="Weight can not exceed the call limit"):
            policy.try_acquire(mocker.Mock(), weight=101)

        with pytest.raises(CallRateLimitHit) as exc:
            policy.try_acquire(mocker.Mock(), weight=100 - 20 - 1 + 1)

        assert exc.value.time_to_wait
        assert exc.value.weight == 100 - 20 - 1 + 1
        assert exc.value.item

    def test_update_available_calls(self, mocker):
        policy = FixedWindowCallRatePolicy(
            matchers=[], next_reset_ts=datetime.now(), period=timedelta(hours=1), call_limit=100
        )
        # update to decrease number of calls available
        policy.update(available_calls=2, call_reset_ts=None)
        # hit the limit with weight=3
        with pytest.raises(CallRateLimitHit):
            policy.try_acquire(mocker.Mock(), weight=3)
        # ok with less weight=1
        policy.try_acquire(mocker.Mock(), weight=1)

        # update to increase number of calls available, ignored
        policy.update(available_calls=20, call_reset_ts=None)
        # so we still hit the limit with weight=3
        with pytest.raises(CallRateLimitHit):
            policy.try_acquire(mocker.Mock(), weight=3)


class TestMovingWindowCallRatePolicy:
    def test_no_rates(self):
        """should raise a ValueError when no rates provided"""
        with pytest.raises(ValueError, match="The list of rates can not be empty"):
            MovingWindowCallRatePolicy(rates=[], matchers=[])

    def test_limit_rate(self):
        """try_acquire must respect configured call rate and throw CallRateLimitHit when hit the limit."""
        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])

        for i in range(10):
            policy.try_acquire("call", weight=1), f"{i + 1} call"

        with pytest.raises(CallRateLimitHit) as excinfo1:
            policy.try_acquire("call", weight=1), "call over limit"
        assert excinfo1.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1)

        time.sleep(0.5)

        with pytest.raises(CallRateLimitHit) as excinfo2:
            policy.try_acquire("call", weight=1), "call over limit"
        assert excinfo2.value.time_to_wait < excinfo1.value.time_to_wait, (
            "time to wait must decrease over time"
        )

    def test_limit_rate_support_custom_weight(self):
        """try_acquire must take into account provided weight and throw CallRateLimitHit when hit the limit."""
        policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])

        policy.try_acquire("call", weight=2), "1st call with weight of 2"
        with pytest.raises(CallRateLimitHit) as excinfo:
            policy.try_acquire("call", weight=9), "2nd call, over limit since 2 + 9 = 11 > 10"
        assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1), (
            "should wait 1 minute before next call"
        )

    def test_multiple_limit_rates(self):
        """try_acquire must take into all call rates and apply stricter."""
        policy = MovingWindowCallRatePolicy(
            matchers=[],
            rates=[
                Rate(10, timedelta(minutes=10)),
                Rate(3, timedelta(seconds=10)),
                Rate(2, timedelta(hours=1)),
            ],
        )

        policy.try_acquire("call", weight=2), "1 call"

        with pytest.raises(CallRateLimitHit) as excinfo:
            policy.try_acquire("call", weight=1), "1 call"

        assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(3600, 0.1)
        assert str(excinfo.value) == "Bucket for item=call with Rate limit=2/1.0h is already full"


class TestHttpStreamIntegration:
    def test_without_cache(self, mocker, requests_mock):
        """Test that HttpStream will use call budget when provided"""
        requests_mock.get(f"{StubDummyHttpStream.url_base}/", json={"data": "test"})

        mocker.patch.object(MovingWindowCallRatePolicy, "try_acquire")

        api_budget = APIBudget(
            policies=[
                MovingWindowCallRatePolicy(
                    matchers=[
                        HttpRequestMatcher(url=f"{StubDummyHttpStream.url_base}/", method="GET")
                    ],
                    rates=[
                        Rate(2, timedelta(minutes=1)),
                    ],
                ),
            ]
        )

        stream = StubDummyHttpStream(
            api_budget=api_budget, authenticator=TokenAuthenticator(token="ABCD")
        )
        for i in range(10):
            records = stream.read_records(SyncMode.full_refresh)
            assert next(records) == {"data": "some_data"}

        assert MovingWindowCallRatePolicy.try_acquire.call_count == 10

    @pytest.mark.usefixtures("enable_cache")
    def test_with_cache(self, mocker, requests_mock):
        """Test that HttpStream will use call budget when provided and not cached"""
        requests_mock.get(f"{StubDummyHttpStream.url_base}/", json={"data": "test"})

        mocker.patch.object(MovingWindowCallRatePolicy, "try_acquire")

        api_budget = APIBudget(
            policies=[
                MovingWindowCallRatePolicy(
                    matchers=[
                        HttpRequestMatcher(url=f"{StubDummyHttpStream.url_base}/", method="GET"),
                    ],
                    rates=[
                        Rate(2, timedelta(minutes=1)),
                    ],
                )
            ]
        )

        stream = StubDummyCacheHttpStream(api_budget=api_budget)
        for i in range(10):
            records = stream.read_records(SyncMode.full_refresh)
            assert next(records) == {"data": "some_data"}

        assert MovingWindowCallRatePolicy.try_acquire.call_count == 1


class TestWeightBasedRateLimiting:
    """Tests for weight-based rate limiting where different endpoints consume different amounts from a shared budget."""

    def test_matcher_weight_default_none(self):
        """HttpRequestRegexMatcher weight defaults to None when not specified."""
        matcher = HttpRequestRegexMatcher(url_path_pattern=r"/api/test")
        assert matcher.weight is None

    def test_matcher_weight_is_stored(self):
        """HttpRequestRegexMatcher stores the weight value when provided."""
        matcher = HttpRequestRegexMatcher(url_path_pattern=r"/api/test", weight=60)
        assert matcher.weight == 60

    def test_matcher_rejects_zero_weight(self):
        """HttpRequestRegexMatcher raises ValueError for weight=0."""
        with pytest.raises(ValueError, match="weight must be >= 1"):
            HttpRequestRegexMatcher(url_path_pattern=r"/api/test", weight=0)

    def test_matcher_rejects_negative_weight(self):
        """HttpRequestRegexMatcher raises ValueError for negative weight."""
        with pytest.raises(ValueError, match="weight must be >= 1"):
            HttpRequestRegexMatcher(url_path_pattern=r"/api/test", weight=-5)

    def test_policy_get_weight_returns_matcher_weight(self):
        """BaseCallRatePolicy.get_weight returns weight from the matching matcher."""
        policy = MovingWindowCallRatePolicy(
            matchers=[HttpRequestRegexMatcher(url_path_pattern=r"/api/expensive", weight=120)],
            rates=[Rate(1000, timedelta(hours=1))],
        )
        req = Request("GET", "https://example.com/api/expensive")
        assert policy.get_weight(req) == 120

    def test_policy_get_weight_defaults_to_1(self):
        """BaseCallRatePolicy.get_weight returns 1 when no matcher has a weight set."""
        policy = MovingWindowCallRatePolicy(
            matchers=[HttpRequestRegexMatcher(url_path_pattern=r"/api/default")],
            rates=[Rate(1000, timedelta(hours=1))],
        )
        req = Request("GET", "https://example.com/api/default")
        assert policy.get_weight(req) == 1

    def test_policy_get_weight_no_matching_matcher(self):
        """BaseCallRatePolicy.get_weight returns 1 when no matcher matches the request."""
        policy = MovingWindowCallRatePolicy(
            matchers=[HttpRequestRegexMatcher(url_path_pattern=r"/api/other", weight=50)],
            rates=[Rate(1000, timedelta(hours=1))],
        )
        req = Request("GET", "https://example.com/api/unmatched")
        assert policy.get_weight(req) == 1

    def test_api_budget_uses_weight(self):
        """APIBudget._do_acquire passes the matcher's weight to try_acquire."""
        policy = MovingWindowCallRatePolicy(
            matchers=[HttpRequestRegexMatcher(url_path_pattern=r"/api/heavy", weight=10)],
            rates=[Rate(100, timedelta(hours=1))],
        )
        budget = APIBudget(policies=[policy])

        # Make requests — each weighs 10 from the budget of 100
        for i in range(10):
            budget.acquire_call(Request("GET", "https://example.com/api/heavy"), block=False)

        # The 11th request should exceed the budget (10 * 10 = 100, one more = 110 > 100)
        with pytest.raises(CallRateLimitHit):
            budget.acquire_call(Request("GET", "https://example.com/api/heavy"), block=False)

    def test_weight_1_backward_compatible(self):
        """When weight is not set, behavior is identical to the old hardcoded weight=1."""
        policy = MovingWindowCallRatePolicy(
            matchers=[HttpRequestRegexMatcher(url_path_pattern=r"/api/normal")],
            rates=[Rate(5, timedelta(hours=1))],
        )
        budget = APIBudget(policies=[policy])

        for i in range(5):
            budget.acquire_call(Request("GET", "https://example.com/api/normal"), block=False)

        with pytest.raises(CallRateLimitHit):
            budget.acquire_call(Request("GET", "https://example.com/api/normal"), block=False)

    def test_shared_budget_different_weights(self):
        """Multiple matchers with different weights sharing one policy correctly consume the shared budget."""
        # Shared policy matches both endpoints via regex
        policy = MovingWindowCallRatePolicy(
            matchers=[
                HttpRequestRegexMatcher(url_path_pattern=r"/api/cheap", weight=1),
                HttpRequestRegexMatcher(url_path_pattern=r"/api/expensive", weight=10),
            ],
            rates=[Rate(20, timedelta(hours=1))],
        )
        budget = APIBudget(policies=[policy])

        # Make 1 expensive request (weight 10) and 10 cheap requests (weight 1 each) = total 20
        budget.acquire_call(Request("GET", "https://example.com/api/expensive"), block=False)
        for i in range(10):
            budget.acquire_call(Request("GET", "https://example.com/api/cheap"), block=False)

        # Budget is now at 20/20 — any further request should fail
        with pytest.raises(CallRateLimitHit):
            budget.acquire_call(Request("GET", "https://example.com/api/cheap"), block=False)

    def test_moving_window_rejects_weight_exceeding_limit(self):
        """MovingWindowCallRatePolicy raises ValueError when weight exceeds the lowest configured rate limit."""
        policy = MovingWindowCallRatePolicy(
            matchers=[HttpRequestRegexMatcher(url_path_pattern=r"/api/heavy", weight=50)],
            rates=[Rate(10, timedelta(hours=1)), Rate(100, timedelta(days=1))],
        )
        req = Request("GET", "https://example.com/api/heavy")
        with pytest.raises(
            ValueError, match="Weight can not exceed the lowest configured rate limit"
        ):
            policy.try_acquire(req, weight=50)


class TestHttpRequestRegexMatcher:
    """
    Tests for the new regex-based logic:
      - Case-insensitive HTTP method matching
      - Optional url_base (scheme://netloc)
      - Regex-based path matching
      - Query params (must be present)
      - Headers (case-insensitive keys)
    """

    def test_case_insensitive_method(self):
        matcher = HttpRequestRegexMatcher(method="GET")

        req_ok = Request("get", "https://example.com/test/path")
        req_wrong = Request("POST", "https://example.com/test/path")

        assert matcher(req_ok)
        assert not matcher(req_wrong)

    def test_url_base(self):
        matcher = HttpRequestRegexMatcher(url_base="https://example.com")

        req_ok = Request("GET", "https://example.com/test/path?foo=bar")
        req_wrong = Request("GET", "https://another.com/test/path?foo=bar")

        assert matcher(req_ok)
        assert not matcher(req_wrong)

    def test_url_path_pattern(self):
        matcher = HttpRequestRegexMatcher(url_path_pattern=r"/test/")

        req_ok = Request("GET", "https://example.com/test/something")
        req_wrong = Request("GET", "https://example.com/other/something")

        assert matcher(req_ok)
        assert not matcher(req_wrong)

    def test_query_params(self):
        matcher = HttpRequestRegexMatcher(params={"foo": "bar"})

        req_ok = Request("GET", "https://example.com/api?foo=bar&extra=123")
        req_missing = Request("GET", "https://example.com/api?not_foo=bar")

        assert matcher(req_ok)
        assert not matcher(req_missing)

    def test_headers_case_insensitive(self):
        matcher = HttpRequestRegexMatcher(headers={"X-Custom-Header": "abc"})

        req_ok = Request(
            "GET",
            "https://example.com/api?foo=bar",
            headers={"x-custom-header": "abc", "other": "123"},
        )
        req_wrong = Request("GET", "https://example.com/api", headers={"x-custom-header": "wrong"})

        assert matcher(req_ok)
        assert not matcher(req_wrong)

    def test_combined_criteria(self):
        matcher = HttpRequestRegexMatcher(
            method="GET",
            url_base="https://example.com",
            url_path_pattern=r"/test/",
            params={"foo": "bar"},
            headers={"X-Test": "123"},
        )

        req_ok = Request("GET", "https://example.com/test/me?foo=bar", headers={"x-test": "123"})
        req_bad_base = Request(
            "GET", "https://other.com/test/me?foo=bar", headers={"x-test": "123"}
        )
        req_bad_path = Request("GET", "https://example.com/nope?foo=bar", headers={"x-test": "123"})
        req_bad_param = Request(
            "GET", "https://example.com/test/me?extra=xyz", headers={"x-test": "123"}
        )
        req_bad_header = Request(
            "GET", "https://example.com/test/me?foo=bar", headers={"some-other-header": "xyz"}
        )

        assert matcher(req_ok)
        assert not matcher(req_bad_base)
        assert not matcher(req_bad_path)
        assert not matcher(req_bad_param)
        assert not matcher(req_bad_header)
