import threading
import time
import unittest
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Generator, Optional
from unittest.mock import Mock, call, patch
from urllib.parse import urlparse
from uuid import UUID

import httpx
import pytest
from httpx import ConnectTimeout, HTTPError

from huggingface_hub.constants import ENDPOINT
from huggingface_hub.errors import HfHubHTTPError, OfflineModeIsEnabled
from huggingface_hub.utils._http import (
    _WARNED_TOPICS,
    RateLimitInfo,
    _adjust_range_header,
    _warn_on_warning_headers,
    default_client_factory,
    fix_hf_endpoint_in_url,
    get_async_session,
    get_session,
    hf_raise_for_status,
    http_backoff,
    parse_ratelimit_headers,
    set_client_factory,
)


URL = "https://www.google.com"


class TestHttpBackoff(unittest.TestCase):
    def setUp(self) -> None:
        get_session_mock = Mock()
        self.mock_request = get_session_mock().request

        self.patcher = patch("huggingface_hub.utils._http.get_session", get_session_mock)
        self.patcher.start()

    def tearDown(self) -> None:
        self.patcher.stop()

    def test_backoff_no_errors(self) -> None:
        """Test normal usage of `http_backoff`."""
        data_mock = Mock()
        response = http_backoff("GET", URL, data=data_mock)
        self.mock_request.assert_called_once_with(method="GET", url=URL, data=data_mock)
        self.assertIs(response, self.mock_request())

    def test_backoff_3_calls(self) -> None:
        """Test `http_backoff` with 2 fails."""
        response_mock = Mock()
        self.mock_request.side_effect = (ValueError(), ValueError(), response_mock)
        response = http_backoff(  # retry on ValueError, instant retry
            "GET", URL, retry_on_exceptions=ValueError, base_wait_time=0.0
        )
        self.assertEqual(self.mock_request.call_count, 3)
        self.mock_request.assert_has_calls(
            calls=[
                call(method="GET", url=URL),
                call(method="GET", url=URL),
                call(method="GET", url=URL),
            ]
        )
        self.assertIs(response, response_mock)

    def test_backoff_on_exception_until_max(self) -> None:
        """Test `http_backoff` until max limit is reached with exceptions."""
        self.mock_request.side_effect = ConnectTimeout("Connection timeout")

        with self.assertRaises(ConnectTimeout):
            http_backoff("GET", URL, base_wait_time=0.0, max_retries=3)

        self.assertEqual(self.mock_request.call_count, 4)

    def test_backoff_on_status_code_until_max(self) -> None:
        """Test `http_backoff` until max limit is reached with status codes."""
        mock_503 = Mock()
        mock_503.status_code = 503
        mock_504 = Mock()
        mock_504.status_code = 504
        mock_504.raise_for_status.side_effect = HTTPError("HTTP Error")
        self.mock_request.side_effect = (mock_503, mock_504, mock_503, mock_504)

        with self.assertRaises(HTTPError):
            http_backoff(
                "GET",
                URL,
                base_wait_time=0.0,
                max_retries=3,
                retry_on_status_codes=(503, 504),
            )

        self.assertEqual(self.mock_request.call_count, 4)

    def test_backoff_on_exceptions_and_status_codes(self) -> None:
        """Test `http_backoff` until max limit with status codes and exceptions."""
        mock_503 = Mock()
        mock_503.status_code = 503
        self.mock_request.side_effect = (mock_503, ConnectTimeout("Connection timeout"))

        with self.assertRaises(ConnectTimeout):
            http_backoff("GET", URL, base_wait_time=0.0, max_retries=1)

        self.assertEqual(self.mock_request.call_count, 2)

    def test_backoff_on_valid_status_code(self) -> None:
        """Test `http_backoff` until max limit with a valid status code.

        Quite a corner case: the user wants to retry is status code is 200. Requests are
        retried but in the end, the HTTP 200 response is returned if the server returned
        only 200 responses.
        """
        mock_200 = Mock()
        mock_200.status_code = 200
        self.mock_request.side_effect = (mock_200, mock_200, mock_200, mock_200)

        response = http_backoff("GET", URL, base_wait_time=0.0, max_retries=3, retry_on_status_codes=200)

        self.assertEqual(self.mock_request.call_count, 4)
        self.assertIs(response, mock_200)

    def test_backoff_sleep_time(self) -> None:
        """Test `http_backoff` sleep time goes exponential until max limit.

        Since timing between 2 requests is sleep duration + some other stuff, this test
        can be unstable. However, sleep durations between 10ms and 50ms should be enough
        to make the approximation that measured durations are the "sleep time" waited by
        `http_backoff`. If this is not the case, just increase `base_wait_time`,
        `max_wait_time` and `expected_sleep_times` with bigger values.
        """
        sleep_times = []

        def _side_effect_timer() -> Generator[ConnectTimeout, None, None]:
            t0 = time.time()
            while True:
                yield ConnectTimeout("Connection timeout")
                t1 = time.time()
                sleep_times.append(round(t1 - t0, 1))
                t0 = t1

        self.mock_request.side_effect = _side_effect_timer()

        with self.assertRaises(ConnectTimeout):
            http_backoff("GET", URL, base_wait_time=0.1, max_wait_time=0.5, max_retries=5)

        self.assertEqual(self.mock_request.call_count, 6)

        # Assert sleep times are exponential until plateau
        expected_sleep_times = [0.1, 0.2, 0.4, 0.5, 0.5]
        self.assertListEqual(sleep_times, expected_sleep_times)

    def test_backoff_on_429_uses_ratelimit_header(self) -> None:
        """Test that 429 wait time uses full reset time from ratelimit header."""
        sleep_times = []

        def _side_effect_timer() -> Generator:
            t0 = time.time()
            mock_429 = Mock()
            mock_429.status_code = 429
            mock_429.headers = {"ratelimit": '"api";r=0;t=1'}  # Server says wait 1s
            yield mock_429
            t1 = time.time()
            sleep_times.append(round(t1 - t0, 1))
            t0 = t1
            mock_200 = Mock()
            mock_200.status_code = 200
            yield mock_200

        self.mock_request.side_effect = _side_effect_timer()

        response = http_backoff(
            "GET", URL, base_wait_time=0.1, max_wait_time=0.5, max_retries=3, retry_on_status_codes=429
        )

        assert self.mock_request.call_count == 2
        assert sleep_times == [2.0]
        assert response.status_code == 200


class TestConfigureSession(unittest.TestCase):
    def setUp(self) -> None:
        # Reconfigure + clear session cache between each test
        set_client_factory(default_client_factory)

    @classmethod
    def tearDownClass(cls) -> None:
        # Clear all sessions after tests
        set_client_factory(default_client_factory)

    @staticmethod
    def _factory() -> httpx.Client:
        client = httpx.Client()
        client.headers.update({"x-test-header": "4"})
        return client

    def test_default_configuration(self) -> None:
        client = get_session()
        # Check httpx.Client default configuration
        self.assertTrue(client.follow_redirects)
        self.assertIsNotNone(client.timeout)

    def test_set_configuration(self) -> None:
        set_client_factory(self._factory)

        # Check headers have been set correctly
        client = get_session()
        self.assertNotEqual(client.headers, {"x-test-header": "4"})
        self.assertEqual(client.headers["x-test-header"], "4")

    def test_get_session_twice(self):
        client_1 = get_session()
        client_2 = get_session()
        self.assertIs(client_1, client_2)  # exact same instance

    def test_get_session_twice_but_reconfigure_in_between(self):
        """Reconfiguring the session clears the cache."""
        client_1 = get_session()
        set_client_factory(self._factory)

        client_2 = get_session()
        self.assertIsNot(client_1, client_2)
        self.assertIsNone(client_1.headers.get("x-test-header"))
        self.assertEqual(client_2.headers["x-test-header"], "4")

    def test_get_session_multiple_threads(self):
        N = 3
        clients = [None] * N

        def _get_session_in_thread(index: int) -> None:
            time.sleep(0.1)
            clients[index] = get_session()

        # Get main thread client
        main_client = get_session()

        # Start 3 threads and get clients in each of them
        threads = [threading.Thread(target=_get_session_in_thread, args=(index,)) for index in range(N)]
        for th in threads:
            th.start()
            print(th)
        for th in threads:
            th.join()

        # Check all clients are the same instance (httpx is thread-safe)
        for i in range(N):
            self.assertIs(main_client, clients[i])
            for j in range(N):
                self.assertIs(clients[i], clients[j])


class OfflineModeSessionTest(unittest.TestCase):
    def tearDown(self) -> None:
        return super().tearDown()

    @patch("huggingface_hub.constants.HF_HUB_OFFLINE", True)
    def test_offline_mode(self):
        set_client_factory(default_client_factory)
        client = get_session()
        with self.assertRaises(OfflineModeIsEnabled):
            client.get("https://huggingface.co")


class TestUniqueRequestId(unittest.TestCase):
    api_endpoint = ENDPOINT + "/api/tasks"  # any endpoint is fine

    def test_request_id_is_used_by_server(self):
        response = get_session().get(self.api_endpoint)

        request_id = response.request.headers.get("X-Amzn-Trace-Id")
        response_id = response.headers.get("x-request-id")
        self.assertIn(request_id, response_id)
        self.assertTrue(_is_uuid(request_id))

    def test_request_id_is_unique(self):
        response_1 = get_session().get(self.api_endpoint)
        response_2 = get_session().get(self.api_endpoint)

        request_id_1 = response_1.request.headers["X-Amzn-Trace-Id"]
        request_id_2 = response_2.request.headers["X-Amzn-Trace-Id"]
        self.assertNotEqual(request_id_1, request_id_2)

        self.assertTrue(_is_uuid(request_id_1))
        self.assertTrue(_is_uuid(request_id_2))

    def test_request_id_not_overwritten(self):
        response = get_session().get(self.api_endpoint, headers={"x-request-id": "custom-id"})

        request_id = response.request.headers["x-request-id"]
        self.assertEqual(request_id, "custom-id")

        response_id = response.headers["x-request-id"]
        self.assertEqual(response_id, "custom-id")


def _is_uuid(string: str) -> bool:
    # Taken from https://stackoverflow.com/a/33245493
    try:
        uuid_obj = UUID(string)
    except ValueError:
        return False
    return str(uuid_obj) == string


@pytest.mark.parametrize(
    ("base_url", "endpoint", "expected_url"),
    [
        # Staging url => unchanged
        ("https://hub-ci.huggingface.co/resolve/...", None, "https://hub-ci.huggingface.co/resolve/..."),
        # Prod url => unchanged
        ("https://huggingface.co/resolve/...", None, "https://huggingface.co/resolve/..."),
        # Custom endpoint + staging url => fixed
        ("https://hub-ci.huggingface.co/api/models", "https://mirror.co", "https://mirror.co/api/models"),
        # Custom endpoint + prod url => fixed
        ("https://huggingface.co/api/models", "https://mirror.co", "https://mirror.co/api/models"),
    ],
)
def test_fix_hf_endpoint_in_url(base_url: str, endpoint: Optional[str], expected_url: str) -> None:
    assert fix_hf_endpoint_in_url(base_url, endpoint) == expected_url


def test_adjust_range_header():
    # Basic cases
    assert _adjust_range_header(None, 10) == "bytes=10-"
    assert _adjust_range_header("bytes=0-100", 10) == "bytes=10-100"
    assert _adjust_range_header("bytes=-100", 10) == "bytes=-90"
    assert _adjust_range_header("bytes=100-", 10) == "bytes=110-"

    with pytest.raises(RuntimeError):
        _adjust_range_header("invalid", 10)

    with pytest.raises(RuntimeError):
        _adjust_range_header("bytes=-", 10)

    # Multiple ranges
    with pytest.raises(ValueError):
        _adjust_range_header("bytes=0-100,200-300", 10)

    # Resume size exceeds range
    with pytest.raises(RuntimeError):
        _adjust_range_header("bytes=0-100", 150)
    with pytest.raises(RuntimeError):
        _adjust_range_header("bytes=-50", 100)


def test_proxy_env_is_used(monkeypatch):
    """Regression test for https://github.com/huggingface/transformers/issues/41301.

    Test is hacky and uses httpx internal attributes, but it works.
    We just need to test that proxies from env vars are used when creating the client.
    """
    monkeypatch.setenv("HTTP_PROXY", "http://proxy.example1.com:8080")
    monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example2.com:8181")

    set_client_factory(default_client_factory)
    client = get_session()
    mounts = client._mounts
    url_patterns = list(mounts.keys())
    assert len(url_patterns) == 2  # http and https

    http_url_pattern = next(url for url in url_patterns if url.pattern == "http://")
    http_proxy_url = mounts[http_url_pattern]._pool._proxy_url
    assert http_proxy_url.scheme == b"http"
    assert http_proxy_url.host == b"proxy.example1.com"
    assert http_proxy_url.port == 8080
    assert http_proxy_url.target == b"/"

    https_url_pattern = next(url for url in url_patterns if url.pattern == "https://")
    https_proxy_url = mounts[https_url_pattern]._pool._proxy_url
    assert https_proxy_url.scheme == b"http"
    assert https_proxy_url.host == b"proxy.example2.com"
    assert https_proxy_url.port == 8181
    assert https_proxy_url.target == b"/"

    # Reset
    set_client_factory(default_client_factory)


def test_client_get_request():
    # Check that sync client works
    client = get_session()
    response = client.get("https://huggingface.co")
    assert response.status_code == 200


@pytest.mark.asyncio
async def test_async_client_get_request():
    # Check that async client works
    client = get_async_session()
    response = await client.get("https://huggingface.co")
    assert response.status_code == 200


class FakeServerHandler(BaseHTTPRequestHandler):
    """Fake server handler to test client behavior."""

    def do_GET(self):
        parsed = urlparse(self.path)

        # Health check endpoint (always succeeds)
        if parsed.path == "/health":
            self._send_response(200, b"OK")
            return

        # Main endpoint (always fails with 500)
        self._send_response(500, b"This is a 500 error")

    def _send_response(self, status_code, body):
        self.send_response(status_code)
        self.send_header("Content-Type", "text/plain")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)


@pytest.fixture(scope="module", autouse=True)
def fake_server():
    # Find a free port
    host, port = "127.0.0.1", 8000
    for port in range(port, 8100):
        try:
            server = HTTPServer((host, port), FakeServerHandler)
            break
        except OSError:
            continue
    else:
        raise RuntimeError("Could not find a free port")

    url = f"http://{host}:{port}"

    # Start server in a separate thread and wait until it's ready
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()

    for _ in range(1000):  # up to 10 seconds
        try:
            if httpx.get(f"{url}/health", timeout=0.01).status_code == 200:
                break
        except httpx.HTTPError:
            pass
        time.sleep(0.01)
    else:
        server.shutdown()
        raise RuntimeError("Fake server failed to start")

    yield url
    server.shutdown()


def _check_raise_status(response: httpx.Response):
    """Common assertions for 500 error tests."""
    with pytest.raises(HfHubHTTPError) as exc_info:
        hf_raise_for_status(response)
    assert exc_info.value.response.status_code == 500
    assert "This is a 500 error" in str(exc_info.value)


def test_raise_on_status_sync_non_stream(fake_server: str):
    response = get_session().get(fake_server)
    _check_raise_status(response)


def test_raise_on_status_sync_stream(fake_server: str):
    with get_session().stream("GET", fake_server) as response:
        _check_raise_status(response)


@pytest.mark.asyncio
async def test_raise_on_status_async_non_stream(fake_server: str):
    response = await get_async_session().get(fake_server)
    _check_raise_status(response)


@pytest.mark.asyncio
async def test_raise_on_status_async_stream(fake_server: str):
    async with get_async_session().stream("GET", fake_server) as response:
        _check_raise_status(response)


class TestParseRatelimitHeaders:
    def test_parse_full_headers(self):
        """Test parsing both ratelimit and ratelimit-policy headers."""
        headers = {
            "ratelimit": '"api";r=0;t=55',
            "ratelimit-policy": '"fixed window";"api";q=500;w=300',
        }
        info = parse_ratelimit_headers(headers)
        assert info == RateLimitInfo(
            resource_type="api",
            remaining=0,
            reset_in_seconds=55,
            limit=500,
            window_seconds=300,
        )

    def test_parse_ratelimit_only(self):
        """Test parsing with only ratelimit header (no policy)."""
        headers = {"ratelimit": '"api";r=489;t=189'}
        info = parse_ratelimit_headers(headers)
        assert info is not None
        assert info.resource_type == "api"
        assert info.remaining == 489
        assert info.reset_in_seconds == 189
        assert info.limit is None
        assert info.window_seconds is None

    def test_parse_missing_header(self):
        """Test returns None when ratelimit header is missing."""
        assert parse_ratelimit_headers({}) is None

    def test_parse_malformed_header(self):
        """Test returns None when ratelimit header is malformed."""
        assert parse_ratelimit_headers({"ratelimit": "malformed"}) is None

    def test_parse_case_insensitive(self):
        """Test header lookup is case-insensitive."""
        headers = {"RateLimit": '"api";r=10;t=100', "RateLimit-Policy": '"fixed window";"api";q=500;w=300'}
        info = parse_ratelimit_headers(headers)
        assert info is not None
        assert info.remaining == 10


class TestRateLimitErrorMessage:
    def test_429_with_ratelimit_headers(self):
        """Test 429 error includes rate limit info when headers present."""
        response = Mock(spec=httpx.Response)
        response.status_code = 429
        response.url = "https://huggingface.co/api/models/username/reponame"
        response.headers = httpx.Headers(
            {
                "ratelimit": '"api";r=0;t=55',
                "ratelimit-policy": '"fixed window";"api";q=500;w=300',
            }
        )
        response.raise_for_status.side_effect = httpx.HTTPStatusError("429", request=Mock(), response=response)
        response.json.return_value = {}

        with pytest.raises(HfHubHTTPError) as exc_info:
            hf_raise_for_status(response)

        error_msg = str(exc_info.value)
        assert "429 Too Many Requests" in error_msg
        assert "'api' rate limit" in error_msg
        assert "55 seconds" in error_msg
        assert "0/500" in error_msg
        assert "api/models/username/reponame" in error_msg

    def test_429_without_ratelimit_headers(self):
        """Test 429 error fallback when headers missing."""
        response = Mock(spec=httpx.Response)
        response.status_code = 429
        response.url = "https://huggingface.co/api/models"
        response.headers = httpx.Headers({})
        response.raise_for_status.side_effect = httpx.HTTPStatusError("429", request=Mock(), response=response)
        response.json.return_value = {}

        with pytest.raises(HfHubHTTPError) as exc_info:
            hf_raise_for_status(response)

        assert "429 Too Many Requests" in str(exc_info.value)
        assert "api/models" in str(exc_info.value)


class TestWarnOnWarningHeaders:
    def test_warn_on_warning_headers(self, caplog):
        # Request #1 (multiple warnings)
        response = Mock(spec=httpx.Response)
        response.headers = httpx.Headers(
            [
                ("X-HF-Warning", "Topic1; This is the first warning message."),
                ("X-HF-Warning", "Topic2; This is the second warning message."),
                ("X-HF-Warning", "Topic1; This is a repeated warning message for Topic1."),
                ("X-HF-Warning", "This is a warning without a topic."),
                ("X-HF-Warning", "This is another warning without a topic."),
            ]
        )

        with caplog.at_level("WARNING"):
            _warn_on_warning_headers(response)

        assert _WARNED_TOPICS == {"Topic1", "Topic2", ""}
        warnings = [record.message for record in caplog.records if record.levelname == "WARNING"]
        assert "This is the first warning message." in warnings
        assert "This is the second warning message." in warnings
        assert "This is a repeated warning message for Topic1." not in warnings
        assert "This is a warning without a topic." in warnings
        assert "This is another warning without a topic." not in warnings
        # Request #2 (exact same warnings, should not warn again)
        caplog.clear()
        with caplog.at_level("WARNING"):
            _warn_on_warning_headers(response)
        warnings = [record.message for record in caplog.records if record.levelname == "WARNING"]
        assert len(warnings) == 0  # No new warnings should be added

        # Request #3 (single warning with new topic, should warn)
        response.headers = httpx.Headers({"X-HF-Warning": "Topic4; Another warning."})
        caplog.clear()
        with caplog.at_level("WARNING"):
            _warn_on_warning_headers(response)
        warnings = [record.message for record in caplog.records if record.levelname == "WARNING"]
        assert len(warnings) == 1
        assert warnings == ["Another warning."]
        assert "Topic4" in _WARNED_TOPICS