from __future__ import annotations

import gzip
import json
import socket
import ssl
import threading
from contextlib import contextmanager
from datetime import UTC, datetime, timedelta
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import httpx
import pytest
from conftest import make_runtime
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID

from scopewrit.errors import GatewayError
from scopewrit.gateway import GuardedExecutor, GuardedHTTPClient
from scopewrit.models import ActionProposal, Effect, ScopeRule


@contextmanager
def server(address="127.0.0.1", port=0, tls=None):
    hits = []

    class Handler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"

        def do_GET(self):
            value = {
                "address": self.server.server_address[0],
                "host": self.headers.get("Host"),
                "cookie": self.headers.get("Cookie"),
                "path": self.path,
            }
            hits.append(value)
            body = json.dumps(value).encode()
            self.send_response(302 if self.path == "/redirect" else 200)
            self.send_header("Set-Cookie", "server-secret=one; Path=/")
            if self.path == "/redirect":
                self.send_header("Location", "/outside")
            if self.path == "/gzip":
                body = gzip.compress(body)
                self.send_header("Content-Encoding", "gzip")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def log_message(self, *args):
            pass

    instance = ThreadingHTTPServer((address, port), Handler)
    instance.daemon_threads = True
    if tls:
        instance.socket = tls.wrap_socket(instance.socket, server_side=True)
    thread = threading.Thread(target=instance.serve_forever, daemon=True)
    thread.start()
    try:
        yield instance.server_port, hits
    finally:
        instance.shutdown()
        instance.server_close()
        thread.join(5)


def runtime_for(tmp_path, port, *, scheme="http", host="allowed.scopewrit.test"):
    return make_runtime(
        tmp_path,
        rules=[
            ScopeRule(
                rule_id="pinned-target",
                effect=Effect.ALLOW,
                action="network.http",
                resource=f"{scheme}://{host}:{port}/**",
                methods=("GET", "HEAD"),
                allow_private_network=True,
                resolved_ip_cidrs=("127.0.0.1/32",),
                max_uses=8,
            )
        ],
    )


def test_real_connection_survives_dns_rebinding_with_matched_legacy_control(tmp_path, monkeypatch):
    with server() as (port, allowed), server("127.0.0.2", port) as (_, outside):
        original = socket.getaddrinfo
        queries = []

        def rebinding(host, port, *args, **kwargs):
            if host == "allowed.scopewrit.test":
                queries.append(host)
                host = "127.0.0.1" if len(queries) == 1 else "127.0.0.2"
            return original(host, port, *args, **kwargs)

        monkeypatch.setattr(socket, "getaddrinfo", rebinding)
        url = f"http://allowed.scopewrit.test:{port}/value"
        legacy = runtime_for(tmp_path / "legacy", port)
        checked = GuardedHTTPClient.resolve(url)
        proposal = ActionProposal(
            request_id="legacy-control",
            run_id=legacy.manifest.run_id,
            action="network.http",
            resource=url,
            method="GET",
            attributes={"resolved_ips": list(checked)},
        )
        with httpx.Client(trust_env=False, timeout=2) as raw:
            GuardedExecutor(legacy.engine, adapter="http").execute(
                attestation=legacy.attestation, proposal=proposal, operation=lambda: raw.get(url)
            )
        assert len(outside) == 1 and not allowed, (
            "legacy check-then-resolve control must expose the gap"
        )
        queries.clear()
        allowed.clear()
        outside.clear()
        fixed = runtime_for(tmp_path / "fixed", port)
        with GuardedHTTPClient(fixed.engine, timeout_seconds=2) as client:
            result = client.request(
                attestation=fixed.attestation,
                run_id=fixed.manifest.run_id,
                request_id="pinned-candidate",
                method="GET",
                url=url,
            )
        assert len(queries) == 1
        assert len(allowed) == 1 and outside == []
        assert result.value.json()["host"] == f"allowed.scopewrit.test:{port}"
        assert result.value.request.url.host == "allowed.scopewrit.test"
        assert result.value.extensions["scopewrit_connection_ip"] == "127.0.0.1"
        ledger = fixed.engine.data_dir / "runs" / fixed.manifest.run_id / "flight-recorder.jsonl"
        execution = json.loads(ledger.read_text().splitlines()[-1])
        assert execution["payload"]["outcome"]["connection_ip"] == "127.0.0.1"


def test_disallowed_resolved_address_is_denied_before_any_socket_effect(tmp_path):
    with server() as (port, hits):
        runtime = runtime_for(tmp_path, port)
        with GuardedHTTPClient(runtime.engine) as client:
            client.resolve = lambda _: ("127.0.0.2",)
            with pytest.raises(GatewayError, match="action deny"):
                client.request(
                    attestation=runtime.attestation,
                    run_id=runtime.manifest.run_id,
                    request_id="bad-address",
                    method="GET",
                    url=f"http://allowed.scopewrit.test:{port}/value",
                )
        assert hits == []


def test_headers_cookies_redirects_and_decoding_remain_explicit(tmp_path):
    with server() as (port, hits):
        runtime = runtime_for(tmp_path, port)
        with GuardedHTTPClient(runtime.engine) as client:
            client.resolve = lambda _: ("127.0.0.1",)
            for index, path in enumerate(["/one", "/two", "/gzip", "/redirect"]):
                result = client.request(
                    attestation=runtime.attestation,
                    run_id=runtime.manifest.run_id,
                    request_id=f"request-{index}",
                    method="GET",
                    url=f"http://allowed.scopewrit.test:{port}{path}",
                )
                if path == "/gzip":
                    assert result.value.json()["path"] == "/gzip"
                    assert "content-encoding" not in result.value.headers
                    assert int(result.value.headers["content-length"]) == len(result.value.content)
                if path == "/redirect":
                    assert result.value.status_code == 302
        assert len(hits) == 4 and all(hit["cookie"] is None for hit in hits)
        assert all(hit["path"] != "/outside" for hit in hits)


def test_tls_uses_logical_hostname_and_never_reuses_another_hosts_identity(tmp_path):
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "allowed.scopewrit.test")])
    now = datetime.now(UTC)
    cert = (
        x509.CertificateBuilder()
        .subject_name(name)
        .issuer_name(name)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(days=1))
        .not_valid_after(now + timedelta(days=1))
        .add_extension(
            x509.SubjectAlternativeName([x509.DNSName("allowed.scopewrit.test")]), critical=False
        )
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .sign(key, hashes.SHA256())
    )
    cert_path = tmp_path / "server.crt"
    key_path = tmp_path / "server.key"
    cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
    key_path.write_bytes(
        key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption(),
        )
    )
    server_tls = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    server_tls.load_cert_chain(cert_path, key_path)
    names = []
    server_tls.set_servername_callback(lambda sock, name, context: names.append(name))
    client_tls = ssl.create_default_context(cafile=str(cert_path))
    with server(tls=server_tls) as (port, hits):
        runtime = runtime_for(tmp_path / "runtime", port, scheme="https", host="*.scopewrit.test")
        with GuardedHTTPClient(runtime.engine, tls_context=client_tls, timeout_seconds=2) as client:
            client.resolve = lambda _: ("127.0.0.1",)
            response = client.request(
                attestation=runtime.attestation,
                run_id=runtime.manifest.run_id,
                request_id="tls-good",
                method="GET",
                url=f"https://allowed.scopewrit.test:{port}/first",
            )
            assert response.value.status_code == 200
            with pytest.raises(httpx.ConnectError, match="CERTIFICATE_VERIFY_FAILED"):
                client.request(
                    attestation=runtime.attestation,
                    run_id=runtime.manifest.run_id,
                    request_id="tls-wrong-host",
                    method="GET",
                    url=f"https://other.scopewrit.test:{port}/second",
                )
        assert names == ["allowed.scopewrit.test", "other.scopewrit.test"]
        assert len(hits) == 1


def test_tls_verification_cannot_be_disabled_by_context(runtime):
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE
    with pytest.raises(ValueError, match="hostname verification"):
        GuardedHTTPClient(runtime.engine, tls_context=context)


@pytest.mark.parametrize("method,status", [("HEAD", 200), ("GET", 304)])
def test_metadata_only_responses_preserve_representation_headers(tmp_path, method, status):
    runtime = runtime_for(tmp_path, 80)
    transport = httpx.MockTransport(
        lambda request: httpx.Response(
            status, headers={"content-length": "777", "content-encoding": "gzip"}, content=b""
        )
    )
    with GuardedHTTPClient(runtime.engine, transport=transport) as client:
        client.resolve = lambda _: ("127.0.0.1",)
        response = client.request(
            attestation=runtime.attestation,
            run_id=runtime.manifest.run_id,
            request_id="metadata-only",
            method=method,
            url="http://allowed.scopewrit.test:80/representation",
        ).value
    assert response.content == b""
    assert response.headers["content-length"] == "777"
    assert response.headers["content-encoding"] == "gzip"
