import json
import socket
import threading
import pytest
from starlette.testclient import TestClient

from tools.havoc_mcp_server import (
    DebugAgentClient,
    DebugAgentError,
    build_asgi_app,
    build_mcp_server,
)


class _StubAgentServer:
    """A minimal in-process stand-in for tools/havoc_debug_agent_server.py, so this
    test never touches ctypes or a real Windows process."""

    def __init__(self, response: dict):
        self._response = response
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._sock.bind(("127.0.0.1", 0))
        self._sock.listen(1)
        self.port = self._sock.getsockname()[1]
        self._thread = threading.Thread(target=self._serve_one, daemon=True)
        self._thread.start()

    def _serve_one(self):
        conn, _ = self._sock.accept()
        conn.recv(4096)
        conn.sendall((json.dumps(self._response) + "\n").encode("utf-8"))
        conn.close()


def test_client_returns_result_on_success():
    stub = _StubAgentServer({"ok": True, "result": {"status": "running"}})
    client = DebugAgentClient(host="127.0.0.1", port=stub.port)
    result = client.call("get_status", {})
    assert result == {"status": "running"}


def test_client_raises_on_error_response():
    stub = _StubAgentServer({"ok": False, "error": "all 4 hardware debug registers (DR0-DR3) are in use"})
    client = DebugAgentClient(host="127.0.0.1", port=stub.port)
    with pytest.raises(DebugAgentError, match="DR0-DR3"):
        client.call("set_watchpoint", {"addr": "0x5000", "size": 4, "mode": "write"})


from tools.havoc_mcp_server import check_bearer_token

def test_check_bearer_token_accepts_matching_token():
    assert check_bearer_token("Bearer secret123", expected="secret123") is True

def test_check_bearer_token_rejects_wrong_token():
    assert check_bearer_token("Bearer wrong", expected="secret123") is False

def test_check_bearer_token_rejects_missing_header():
    assert check_bearer_token(None, expected="secret123") is False

def test_check_bearer_token_rejects_malformed_header():
    assert check_bearer_token("secret123", expected="secret123") is False  # missing "Bearer " prefix

def test_check_bearer_token_rejects_empty_expected_token():
    # A set-but-empty HAVOC_MCP_TOKEN must not turn auth into a no-op.
    assert check_bearer_token("Bearer ", expected="") is False
    assert check_bearer_token("Bearer anything", expected="") is False


def _build_test_app(token: str = "secret123"):
    """Build the real ASGI app (auth middleware wrapping the FastMCP Streamable HTTP
    app) the same way __main__ does, but with a DebugAgentClient that never needs to
    make a real connection -- these tests only exercise the auth middleware layer,
    they never let a request reach the point of actually calling the debug agent."""
    client = DebugAgentClient(host="127.0.0.1", port=1)
    server = build_mcp_server(client)
    return build_asgi_app(server, token)


def test_asgi_app_rejects_request_with_no_auth_header():
    app = _build_test_app()
    with TestClient(app) as test_client:
        response = test_client.post("/mcp", json={})
    assert response.status_code == 401


def test_asgi_app_rejects_request_with_wrong_token():
    app = _build_test_app()
    with TestClient(app) as test_client:
        response = test_client.post(
            "/mcp", json={}, headers={"Authorization": "Bearer wrong"}
        )
    assert response.status_code == 401


def test_asgi_app_passes_through_request_with_correct_token():
    app = _build_test_app()
    with TestClient(app) as test_client:
        response = test_client.post(
            "/mcp",
            json={},
            headers={
                "Authorization": "Bearer secret123",
                "Accept": "application/json, text/event-stream",
            },
        )
    # The auth middleware must not 401 a correctly-authenticated request. Whatever
    # status the underlying MCP session manager returns past that point (e.g. it may
    # reject malformed MCP protocol bodies) is out of scope here -- a non-401 response
    # is sufficient proof the middleware let the request through.
    assert response.status_code != 401