import unittest
from datetime import datetime, timezone
from unittest.mock import AsyncMock, patch

from tuotuo.metrics import MetricsRoute, PrometheusClient, QueryParam, QueryRange


class FakeResponse:
    ok = True
    status = 200

    def __init__(self, payload):
        self.payload = payload

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        return False

    async def json(self):
        return self.payload

    async def text(self):
        return ""


class FakeRetryClient:
    def __init__(self, responses):
        self.responses = iter(responses)
        self.requests = []

    def get(self, url, **kwargs):
        self.requests.append((url, kwargs))
        return FakeResponse(next(self.responses))

    async def close(self):
        return None


def prometheus_response(result_type):
    return {
        "status": "success",
        "data": {
            "resultType": result_type,
            "result": [],
        },
    }


class PrometheusClientTest(unittest.IsolatedAsyncioTestCase):
    async def test_cluster_scope_prefers_dedicated_data_proxy_endpoint(self):
        session = FakeRetryClient([prometheus_response("vector")])

        with patch.dict("os.environ", {
            "CLINIC_API_TOKEN": "test-token",
            "CLINIC_API_ENDPOINT": "http://clinic-private.example.com:8085",
            "CLINIC_DATA_PROXY_ENDPOINT": "https://clinic.example.com",
        }), patch(
            "tuotuo.metrics.client.aiohttp_retry.RetryClient",
            return_value=session,
        ), patch(
            "tuotuo.metrics.client.cluster_info",
            new=AsyncMock(side_effect=AssertionError("cluster lookup is not needed")),
        ):
            client = PrometheusClient()
            await client.query(QueryParam(
                time_arg=1.0,
                query_template="up",
                cluster_id="cluster-1",
                metrics_route=MetricsRoute.CLUSTER_PROXY,
            ))
            await client.close()

        self.assertEqual(
            "https://clinic.example.com/data-proxy/metrics/cluster-1/api/v1/query",
            session.requests[0][0],
        )

    async def test_cluster_scope_uses_cluster_data_proxy_for_instant_query(self):
        session = FakeRetryClient([prometheus_response("vector")])

        with patch.dict("os.environ", {
            "CLINIC_API_TOKEN": "test-token",
            "CLINIC_API_ENDPOINT": "https://clinic.example.com",
        }, clear=True), patch(
            "tuotuo.metrics.client.aiohttp_retry.RetryClient",
            return_value=session,
        ), patch(
            "tuotuo.metrics.client.cluster_info",
            new=AsyncMock(side_effect=AssertionError("cluster lookup is not needed")),
        ):
            client = PrometheusClient()
            await client.query(QueryParam(
                time_arg=1.0,
                query_template="up",
                cluster_id="cluster-1",
                metrics_route=MetricsRoute.CLUSTER_PROXY,
            ))
            await client.close()

        self.assertEqual(
            "https://clinic.example.com/data-proxy/metrics/cluster-1/api/v1/query",
            session.requests[0][0],
        )
        self.assertEqual(
            {"Authorization": "Bearer test-token"},
            session.requests[0][1]["headers"],
        )

    async def test_cluster_scope_uses_cluster_data_proxy_for_range_query(self):
        session = FakeRetryClient([prometheus_response("matrix")])
        query_range = QueryRange(
            datetime(2026, 8, 27, 9, 58, tzinfo=timezone.utc),
            datetime(2026, 8, 27, 10, 1, tzinfo=timezone.utc),
            "15s",
        )

        with patch.dict("os.environ", {
            "CLINIC_API_TOKEN": "test-token",
            "CLINIC_API_ENDPOINT": "https://clinic.example.com",
        }, clear=True), patch(
            "tuotuo.metrics.client.aiohttp_retry.RetryClient",
            return_value=session,
        ):
            client = PrometheusClient()
            await client.query(QueryParam(
                time_arg=query_range,
                query_template="up",
                cluster_id="cluster-1",
                metrics_route=MetricsRoute.CLUSTER_PROXY,
            ))
            await client.close()

        self.assertEqual(
            "https://clinic.example.com/data-proxy/metrics/cluster-1/api/v1/query_range",
            session.requests[0][0],
        )
        self.assertEqual(1787824680, session.requests[0][1]["params"]["start"])
        self.assertEqual(1787824860, session.requests[0][1]["params"]["end"])

    async def test_default_route_preserves_legacy_control_plane_route(self):
        session = FakeRetryClient([prometheus_response("vector")])

        with patch.dict("os.environ", {
            "CLINIC_API_TOKEN": "test-token",
            "CLINIC_API_ENDPOINT": "https://clinic.example.com",
            "CLINIC_DATA_PROXY_ENDPOINT": "https://data-proxy.example.com",
        }), patch(
            "tuotuo.metrics.client.aiohttp_retry.RetryClient",
            return_value=session,
        ):
            client = PrometheusClient()
            await client.query(QueryParam(
                time_arg=1.0,
                query_template="up",
                cluster_id="cluster-1",
                clinic_org_id="org-1",
                provider="aws",
                region_name="us-west-2",
            ))
            await client.close()

        self.assertEqual(
            "https://clinic.example.com/clinic/api/v1/data/control_plane/metrics",
            session.requests[0][0],
        )
        self.assertEqual("cluster-1", session.requests[0][1]["headers"]["X-ClusterID"])

    async def test_default_route_preserves_legacy_range_routing(self):
        session = FakeRetryClient([
            prometheus_response("matrix"),
            prometheus_response("matrix"),
        ])
        query_range = QueryRange(
            datetime(2026, 8, 27, 9, 58, tzinfo=timezone.utc),
            datetime(2026, 8, 27, 10, 1, tzinfo=timezone.utc),
            "15s",
        )

        with patch.dict("os.environ", {
            "CLINIC_API_TOKEN": "test-token",
            "CLINIC_API_ENDPOINT": "https://clinic.example.com",
        }), patch(
            "tuotuo.metrics.client.aiohttp_retry.RetryClient",
            return_value=session,
        ):
            client = PrometheusClient()
            common = {
                "time_arg": query_range,
                "cluster_id": "cluster-1",
                "clinic_org_id": "org-1",
                "provider": "aws",
                "region_name": "us-west-2",
            }
            await client.query(QueryParam(query_template="up", **common))
            await client.query(QueryParam(
                query_template='up{cluster_id=~".*cluster-1"}',
                **common,
            ))
            await client.close()

        self.assertEqual(
            "https://clinic.example.com/clinic/api/v1/data/control_plane/metrics",
            session.requests[0][0],
        )
        self.assertEqual(
            "https://clinic.example.com/clinic/api/v1/data/metrics",
            session.requests[1][0],
        )


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