"""Unit tests for the connectr Python SDK.

All tests use mocked HTTP servers and no real gRPC connection.
"""

from __future__ import annotations

import asyncio
import json
import threading
import time
import base64
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any
from urllib.parse import urlparse

import pytest

from service_bridge.http_catalog import extract_trace_from_headers, _new_hex_id, HttpSpan
from service_bridge.servicebridge import (
    EventContext,
    Options,
    ServiceBridge,
    WorkerTLSOpts,
    _FnChannelEntry,
    _TLSMaterials,
    _aiter_grpc_stream,
    _parse_service_key_v2,
    _resolve_ca_pem,
)


_TEST_CA_PEM = """-----BEGIN CERTIFICATE-----
MIIBkDCCATWgAwIBAgIUYRfcgaahUqALFqup7iwPh5YPXpkwCgYIKoZIzj0EAwIw
HTEbMBkGA1UEAwwSc2VydmljZWJyaWRnZS10ZXN0MB4XDTI2MDMxMzEwNTAzMVoX
DTI3MDMxMzEwNTAzMVowHTEbMBkGA1UEAwwSc2VydmljZWJyaWRnZS10ZXN0MFkw
EwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE7apvCGZrlwyXl9CyoEC41MWjoSVxKYJ1
5wqgt2Xy7HLhu42oe3xHi/4EHgjud/FvUvxm9d0qK1VRP5WYp7+D8KNTMFEwHQYD
VR0OBBYEFIFvIQqp9rwGGgB9fWibecpnfYvAMB8GA1UdIwQYMBaAFIFvIQqp9rwG
GgB9fWibecpnfYvAMA8GA1UdEwEB/wQFMAMBAf8wCgYIKoZIzj0EAwIDSQAwRgIh
AMRHdpu4ZsTIbJ0kRZuA4dUbcTsIe2uA9tvGVCT+4g0+AiEA9Jw8vHs3COSKGKF+
GBcG18GGjC7zhBTZTtlQROqngSI=
-----END CERTIFICATE-----
"""
_TEST_CERT_PEM = _TEST_CA_PEM
_TEST_KEY_PEM = """-----BEGIN PRIVATE KEY-----
MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgK3yq3vNMxKGMzF1t
pXm5PB+0ns/4B69cXX4jFv3c0SOhRANCAATtqm8IZmuXDJeX0LKgQLjUxaOhJXEp
gnXnCqC3ZfLscuG7jah7fEeL/gQeCO538W9S/Gb13SorVVE/lZinv4Pw
-----END PRIVATE KEY-----
"""
_TEST_SECRET_SEGMENT = base64.urlsafe_b64encode(b"0123456789abcdef0123456789abcdef").decode().rstrip("=")
_TEST_CA_SEGMENT = base64.urlsafe_b64encode(_TEST_CA_PEM.encode()).decode().rstrip("=")
_TEST_SERVICE_KEY = f"sbv2.k1.{_TEST_SECRET_SEGMENT}.{_TEST_CA_SEGMENT}"


# ─── Helpers ──────────────────────────────────────────────────────────────────

class _MockAdminHandler(BaseHTTPRequestHandler):
    """A minimal HTTP server that records POST /api/http/register calls."""

    calls: list[dict] = []
    responses: dict[str, tuple[int, str]] = {}

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length)
        parsed = {}
        if body:
            try:
                parsed = json.loads(body)
            except json.JSONDecodeError:
                pass
        _MockAdminHandler.calls.append({"path": self.path, "body": parsed, "headers": dict(self.headers)})

        status, response_body = _MockAdminHandler.responses.get(self.path, (200, "{}"))
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(response_body.encode())

    def log_message(self, format, *args):  # noqa: A002
        pass  # suppress output


def make_mock_server() -> tuple[HTTPServer, str]:
    """Start a mock HTTP server on a random port and return (server, base_url)."""
    _MockAdminHandler.calls = []
    _MockAdminHandler.responses = {}
    server = HTTPServer(("localhost", 0), _MockAdminHandler)
    port = server.server_address[1]
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server, f"http://localhost:{port}"


# ─── ServiceBridge construction ───────────────────────────────────────────────

class TestServiceBridgeInit:
    def test_default_creation(self):
        key = f"sbv2.mysvc.{_TEST_SECRET_SEGMENT}.{_TEST_CA_SEGMENT}"
        svc = ServiceBridge("localhost:14445", key)
        assert svc is not None
        assert svc._service_name == "mysvc"
        assert svc._service_key == key

    def test_with_options(self):
        opts = Options(heartbeat_interval_ms=5000)
        svc = ServiceBridge("localhost:14445", "key", opts)
        assert svc._opts.heartbeat_interval_ms == 5000

    def test_instance_id_is_unique(self):
        a = ServiceBridge("localhost:14445", "key")
        b = ServiceBridge("localhost:14445", "key")
        assert a._instance_id != b._instance_id


# ─── handle_rpc decorator ─────────────────────────────────────────────────────

class TestHandleRpc:
    def test_registers_handler(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_rpc("my.fn")
        async def handler(payload):
            return {"ok": True}

        assert "my.fn" in svc._rpc_handlers
        assert svc._rpc_handlers["my.fn"] is handler

    def test_returns_original_function(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_rpc("fn1")
        def sync_handler(payload):
            return {}

        assert svc._rpc_handlers["fn1"] is sync_handler

    def test_multiple_handlers(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_rpc("fn.a")
        async def h_a(payload): return {}

        @svc.handle_rpc("fn.b")
        async def h_b(payload): return {}

        assert "fn.a" in svc._rpc_handlers
        assert "fn.b" in svc._rpc_handlers
        assert len(svc._rpc_handlers) == 2

    def test_overwrite_handler(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_rpc("fn")
        async def h1(payload): return {"v": 1}

        @svc.handle_rpc("fn")
        async def h2(payload): return {"v": 2}

        assert svc._rpc_handlers["fn"] is h2


# ─── handle_event decorator ───────────────────────────────────────────────────

class TestHandleEvent:
    def test_registers_with_default_group(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_event("order.created")
        async def handler(payload, ec): pass

        assert "consumer.order.created" in svc._event_handlers

    def test_stores_topic(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_event("my.topic")
        async def handler(payload, ec): pass

        key = "consumer.my.topic"
        assert svc._event_handlers[key]["topic"] == "my.topic"

    def test_stores_retry_policy(self):
        svc = ServiceBridge("localhost:14445", "key")
        policy = '{"maxAttempts": 3}'

        @svc.handle_event("t", retry_policy_json=policy)
        async def handler(payload, ec): pass

        key = "consumer.t"
        assert svc._event_handlers[key]["retry_policy_json"] == policy

    def test_returns_original_function(self):
        svc = ServiceBridge("localhost:14445", "key")

        @svc.handle_event("t")
        def sync_handler(payload, ec): pass

        assert svc._event_handlers["consumer.t"]["handler"] is sync_handler


# ─── EventContext ─────────────────────────────────────────────────────────────

class TestEventContext:
    def _make_ec(self):
        retry_calls = []
        reject_calls = []
        ec = EventContext(
            "trace1",
            "span1",
            "t",
            "g",
            message_id="msg1",
            attempt=1,
            headers={},
            _retry_fn=lambda ms: retry_calls.append(ms),
            _reject_fn=lambda r: reject_calls.append(r),
        )
        return ec, retry_calls, reject_calls

    def test_retry_stores_delay(self):
        ec, retry_calls, _ = self._make_ec()
        ec.retry(500)
        assert retry_calls == [500]

    def test_retry_default_delay(self):
        ec, retry_calls, _ = self._make_ec()
        ec.retry()
        assert retry_calls == [1000]

    def test_reject_stores_reason(self):
        ec, _, reject_calls = self._make_ec()
        ec.reject("bad_payload")
        assert reject_calls == ["bad_payload"]

    def test_reject_default_reason(self):
        ec, _, reject_calls = self._make_ec()
        ec.reject()
        assert reject_calls == ["rejected_by_consumer"]


# ─── extract_trace_from_headers ───────────────────────────────────────────────

class TestExtractTrace:
    def test_w3c_traceparent(self):
        headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
        trace_id, parent_span_id = extract_trace_from_headers(headers)
        assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"
        assert parent_span_id == "00f067aa0ba902b7"

    def test_x_trace_id(self):
        headers = {"x-trace-id": "my-trace-123"}
        trace_id, parent_span_id = extract_trace_from_headers(headers)
        assert trace_id == "my-trace-123"
        assert parent_span_id == ""

    def test_no_headers_generates_id(self):
        trace_id, parent_span_id = extract_trace_from_headers({})
        assert trace_id != ""
        assert len(trace_id) == 32  # 16 bytes hex
        assert parent_span_id == ""

    def test_invalid_traceparent_falls_back(self):
        headers = {"traceparent": "bad-value"}
        trace_id, _ = extract_trace_from_headers(headers)
        assert trace_id != ""

    def test_case_insensitive_headers(self):
        headers = {"X-Trace-Id": "upper-case-header"}
        trace_id, _ = extract_trace_from_headers(headers)
        assert trace_id == "upper-case-header"

    def test_traceparent_takes_priority_over_x_trace_id(self):
        headers = {
            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
            "x-trace-id": "should-be-ignored",
        }
        trace_id, _ = extract_trace_from_headers(headers)
        assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736"


# ─── _new_hex_id ──────────────────────────────────────────────────────────────

class TestNewHexId:
    def test_default_length(self):
        result = _new_hex_id()
        assert len(result) == 32  # 16 bytes * 2

    def test_custom_length(self):
        result = _new_hex_id(8)
        assert len(result) == 16

    def test_unique(self):
        ids = {_new_hex_id() for _ in range(100)}
        assert len(ids) == 100


# ─── HttpSpan ─────────────────────────────────────────────────────────────────

class TestHttpSpan:
    def _make_mock_client(self):
        """Return a minimal object that has _report_call."""
        calls = []

        class MockClient:
            def _report_call(self, *args, **kwargs):
                calls.append(args)

        return MockClient(), calls

    def test_start_http_span_returns_span(self):
        svc = ServiceBridge("localhost:14445", "key")
        span = svc.start_http_span("GET", "/users")
        assert span is not None
        assert span.trace_id != ""
        assert span.span_id != ""

    def test_start_http_span_propagates_trace_id(self):
        svc = ServiceBridge("localhost:14445", "key")
        span = svc.start_http_span("GET", "/users", trace_id="fixed-trace-id")
        assert span.trace_id == "fixed-trace-id"

    def test_span_end_calls_report(self):
        mock_client, calls = self._make_mock_client()
        span = HttpSpan(
            trace_id="t1",
            span_id="s1",
            _fn="http:GET:/test",
            _started_at=0,
            _client=mock_client,
            _input_buf=b"{}",
        )
        span.end(status_code=200)
        assert len(calls) == 1

    def test_span_end_success_on_2xx(self):
        mock_client, calls = self._make_mock_client()
        span = HttpSpan(
            trace_id="t", span_id="s", _fn="f", _started_at=0,
            _client=mock_client, _input_buf=b"{}",
        )
        span.end(status_code=201)
        assert calls[0][5] is True  # success=True

    def test_span_end_failure_on_5xx(self):
        mock_client, calls = self._make_mock_client()
        span = HttpSpan(
            trace_id="t", span_id="s", _fn="f", _started_at=0,
            _client=mock_client, _input_buf=b"{}",
        )
        span.end(status_code=500)
        assert calls[0][5] is False  # success=False

    def test_span_end_explicit_success_override(self):
        mock_client, calls = self._make_mock_client()
        span = HttpSpan(
            trace_id="t", span_id="s", _fn="f", _started_at=0,
            _client=mock_client, _input_buf=b"{}",
        )
        span.end(status_code=500, success=True)
        assert calls[0][5] is True


# ─── register_http_endpoint (gRPC) ────────────────────────────────────────────

class _HttpStubOk:
    def __init__(self):
        self.calls: list[Any] = []

    def RegisterHttpEndpoint(self, req, metadata=None, timeout=None):  # noqa: N802
        self.calls.append((req, metadata, timeout))
        return type("Resp", (), {"ok": True})()


class _HttpStubReject:
    def RegisterHttpEndpoint(self, req, metadata=None, timeout=None):  # noqa: N802
        return type("Resp", (), {"ok": False})()


class TestRegisterHttpEndpoint:
    async def test_requires_started_client(self):
        svc = ServiceBridge("localhost:14445", "testkey")
        with pytest.raises(Exception):
            await svc.register_http_endpoint("GET", "/users")

    async def test_sends_grpc_request_payload(self):
        svc = ServiceBridge("localhost:14445", f"sbv2.testsvc.{_TEST_SECRET_SEGMENT}.{_TEST_CA_SEGMENT}")
        stub = _HttpStubOk()
        svc._stub = stub
        svc._worker_address = "127.0.0.1:9999"
        await svc.register_http_endpoint("GET", "/users", allowed_callers=["admin-service"])
        req, _, _ = stub.calls[0]
        assert req.service_name == "testsvc"
        assert req.method == "GET"
        assert req.route_pattern == "/users"
        assert list(req.allowed_callers) == ["admin-service"]

    async def test_raises_when_runtime_rejects(self):
        svc = ServiceBridge("localhost:14445", "key")
        svc._stub = _HttpStubReject()
        with pytest.raises(Exception):
            await svc.register_http_endpoint("GET", "/protected")


class _HangUnaryCall:
    def __init__(self):
        self.cancelled = False

    def result(self, timeout=None):
        _ = timeout
        time.sleep(1.0)
        return type("Resp", (), {"success": True, "output": b"{}"})()

    def cancel(self):
        self.cancelled = True


class _HangHandle:
    def __init__(self):
        self.calls: list[_HangUnaryCall] = []

    def future(self, req, timeout=None, wait_for_ready=None):
        _ = (req, timeout, wait_for_ready)
        call = _HangUnaryCall()
        self.calls.append(call)
        return call


class _HangWorkerStub:
    def __init__(self):
        self.Handle = _HangHandle()


class TestRpcTimeoutGuard:
    @pytest.mark.asyncio
    async def test_hung_rpc_attempt_is_cancelled_and_fails_by_timeout(self):
        svc = ServiceBridge("localhost:14445", "test-key")
        svc._stub = object()  # only used as started/not-started flag in rpc()

        worker_stub = _HangWorkerStub()

        async def fake_get_fn_channel(_fn: str):
            return object(), worker_stub

        svc._get_fn_channel = fake_get_fn_channel  # type: ignore[method-assign]

        with pytest.raises(Exception) as excinfo:
            await svc.rpc("svc", "hang", {"value": 1}, retries=0, timeout_ms=20)

        assert "timeout" in str(excinfo.value).lower()
        assert len(worker_stub.Handle.calls) == 1
        assert worker_stub.Handle.calls[0].cancelled is True


class TestParityValidation:
    def test_parse_service_key_v2_requires_valid_embedded_ca(self):
        _parse_service_key_v2(_TEST_SERVICE_KEY)
        bad_key = f"sbv2.k1.{_TEST_SECRET_SEGMENT}.bm90LWEtdmFsaWQtY2VydA"
        with pytest.raises(RuntimeError, match="malformed embedded ca pem"):
            _parse_service_key_v2(bad_key)

    def test_resolve_ca_pem_rejects_invalid_explicit_ca(self):
        with pytest.raises(RuntimeError, match="control-plane ca cert"):
            _resolve_ca_pem(_TEST_SERVICE_KEY, "not-a-cert")

    def test_handle_hints_are_stored(self):
        svc = ServiceBridge("localhost:14445", _TEST_SERVICE_KEY)

        @svc.handle_rpc(
            "echo",
            timeout_ms=5000,
            retryable=True,
            concurrency=8,
        )
        async def _rpc(_payload):
            return {"ok": True}

        @svc.handle_event("orders.*", concurrency=4, prefetch=32)
        async def _event(_payload, _ctx):
            return None

        assert svc._rpc_opts["echo"]["timeout_ms"] == 5000
        assert svc._rpc_opts["echo"]["retryable"] is True
        assert svc._rpc_opts["echo"]["concurrency"] == 8
        assert svc._event_handlers["k1.orders.*"]["concurrency"] == 4
        assert svc._event_handlers["k1.orders.*"]["prefetch"] == 32

    @pytest.mark.asyncio
    async def test_serve_fails_fast_on_initial_worker_session_error(self):
        svc = ServiceBridge("localhost:14445", _TEST_SERVICE_KEY)

        @svc.handle_rpc("ping")
        async def _ping(_payload):
            return {}

        svc._worker_tls_creds = _TLSMaterials(
            ca_cert=_TEST_CA_PEM.encode(),
            client_cert=_TEST_CERT_PEM.encode(),
            client_key=_TEST_KEY_PEM.encode(),
        )

        async def fake_connect_control_plane():
            svc._stub = object()

        async def fake_start_worker_server(_host: str) -> str:
            return "127.0.0.1:7443"

        def fake_start_worker_session():
            raise RuntimeError("OpenWorkerSession initial connect failed: boom")

        sync_called = False

        async def fake_sync_registrations():
            nonlocal sync_called
            sync_called = True

        svc._connect_control_plane = fake_connect_control_plane  # type: ignore[method-assign]
        svc._start_worker_server = fake_start_worker_server  # type: ignore[method-assign]
        svc._start_worker_session = fake_start_worker_session  # type: ignore[method-assign]
        svc._sync_registrations = fake_sync_registrations  # type: ignore[method-assign]

        with pytest.raises(RuntimeError, match="OpenWorkerSession initial connect failed"):
            await svc.serve(host="localhost", instance_id="i-1", weight=7)

        assert sync_called is False
        assert svc._instance_id == "i-1"
        assert svc._registration_weight == 7

    def test_resolve_explicit_worker_tls_validates_cert_and_key(self):
        svc = ServiceBridge("localhost:14445", _TEST_SERVICE_KEY)
        with pytest.raises(RuntimeError, match="must be provided together"):
            svc._resolve_explicit_worker_tls(WorkerTLSOpts(cert=_TEST_CERT_PEM))
        with pytest.raises(RuntimeError, match="cert/key pair mismatch|invalid worker tls key"):
            svc._resolve_explicit_worker_tls(
                WorkerTLSOpts(cert=_TEST_CERT_PEM, key="not-a-key")
            )
        resolved = svc._resolve_explicit_worker_tls(
            WorkerTLSOpts(cert=_TEST_CERT_PEM, key=_TEST_KEY_PEM)
        )
        assert resolved is not None
        materials, _ = resolved
        assert materials.client_cert.startswith(b"-----BEGIN CERTIFICATE-----")

    def test_fn_channel_round_robin_selector(self):
        svc = ServiceBridge("localhost:14445", _TEST_SERVICE_KEY)
        entry = _FnChannelEntry(
            channels=["ch1", "ch2"],  # type: ignore[list-item]
            stubs=["stub1", "stub2"],
            endpoints=["a:1", "b:2"],
        )
        _, stub1 = svc._pick_fn_channel(entry)
        _, stub2 = svc._pick_fn_channel(entry)
        _, stub3 = svc._pick_fn_channel(entry)
        assert stub1 == "stub1"
        assert stub2 == "stub2"
        assert stub3 == "stub1"

    @pytest.mark.asyncio
    async def test_stream_helper_fails_on_queue_overflow(self):
        class _FastStream:
            def __iter__(self):
                for i in range(128):
                    yield i

            def cancel(self):
                return None

        loop = asyncio.get_event_loop()
        with pytest.raises(RuntimeError, match="overflow"):
            async for _ in _aiter_grpc_stream(
                _FastStream(),
                loop,
                queue_limit=1,
                overflow_error_factory=lambda: RuntimeError("overflow"),
            ):
                await asyncio.sleep(0.05)
