#!/usr/bin/env python3
"""Compile and exercise strict generated Swift and Rust gateway contracts."""

from __future__ import annotations

import pathlib
import shutil
import subprocess
import tempfile


ROOT = pathlib.Path(__file__).resolve().parents[1]


def run(command: list[str], directory: pathlib.Path) -> None:
    result = subprocess.run(command, cwd=directory, text=True, capture_output=True, check=False)
    if result.returncode != 0:
        raise RuntimeError("%s\n%s%s" % (" ".join(command), result.stdout, result.stderr))


def run_failure(command: list[str], directory: pathlib.Path, expected: str) -> None:
    result = subprocess.run(command, cwd=directory, text=True, capture_output=True, check=False)
    if result.returncode == 0 or expected not in result.stdout + result.stderr:
        raise RuntimeError("expected failure: %s\n%s%s" % (" ".join(command), result.stdout, result.stderr))


def main() -> int:
    with tempfile.TemporaryDirectory(prefix="versevdi-generated-contracts-") as temporary:
        workspace = pathlib.Path(temporary)
        swift = workspace / "main.swift"
        swift.write_text(
            """import Foundation

let capability = try CapabilityProfile(
    transport: "quic-tls13", framing: "datagram-v1", media: "encoded",
    audio: "encoded", sourceRateControl: "server", clientDecode: ["h264-opus"]
)
guard currentWireVersion == "2", nMinus1WireVersion == "1", nMinus2WireVersion == "0" else {
    fatalError("unexpected control wire compatibility declaration")
}
_ = try CapabilityProfile(
    transport: "quic-tls13", framing: "datagram-v2", media: "encoded",
    audio: "encoded", sourceRateControl: "server", clientDecode: ["h264-opus"]
)
do {
    _ = try CapabilityProfile(
        transport: "quic-tls13", framing: "datagram-v3", media: "encoded",
        audio: "encoded", sourceRateControl: "server", clientDecode: ["h264-opus"]
    )
    fatalError("unregistered framing was accepted")
} catch { }
let request = try TunnelAdmissionRequest(
    version: "1", sessionId: "session", gatewayId: "gateway", audience: "audience",
    grant: String(repeating: "g", count: 43), reconnectSequence: 0,
    clientNonce: String(repeating: "n", count: 16), deviceSignature: String(repeating: "s", count: 86), capabilities: capability
)
_ = request
let transcript = "versevdi/tunnel-admission/v17:session7:gateway8:audience43:" + String(repeating: "g", count: 43) + "1:016:" + String(repeating: "n", count: 16) + "10:quic-tls1311:datagram-v17:encoded7:encoded6:server1:19:h264-opus"
guard String(data: request.deviceAdmissionTranscript(), encoding: .utf8) == transcript else {
    fatalError("unexpected device admission transcript")
}
let proofServerID = Data(repeating: 1, count: 16)
let proofPrincipalID = Data(repeating: 2, count: 16)
let proofDeviceID = Data(repeating: 3, count: 16)
let proofChallenge = Data(repeating: 4, count: 32)
let proofTranscript = try deviceRegistrationProofTranscript(
    serverID: proofServerID, principalID: proofPrincipalID, deviceID: proofDeviceID,
    challenge: proofChallenge, expiryUnixMilliseconds: 1
)
guard proofTranscript.count == 112,
      String(data: proofTranscript.prefix(24), encoding: .utf8) == "versevdi-device-proof-v1",
      Array(proofTranscript.suffix(8)) == [0, 0, 0, 0, 0, 0, 0, 1] else {
    fatalError("unexpected device registration proof transcript")
}
let invalidProofInputs: [(String, String, Data, Data, Data, Data, Int64)] = [
    ("server-short", "server_id", Data(repeating: 0, count: 15), proofPrincipalID, proofDeviceID, proofChallenge, 0),
    ("server-long", "server_id", Data(repeating: 0, count: 17), proofPrincipalID, proofDeviceID, proofChallenge, 0),
    ("principal-short", "principal_id", proofServerID, Data(repeating: 0, count: 15), proofDeviceID, proofChallenge, 0),
    ("principal-long", "principal_id", proofServerID, Data(repeating: 0, count: 17), proofDeviceID, proofChallenge, 0),
    ("device-short", "device_id", proofServerID, proofPrincipalID, Data(repeating: 0, count: 15), proofChallenge, 0),
    ("device-long", "device_id", proofServerID, proofPrincipalID, Data(repeating: 0, count: 17), proofChallenge, 0),
    ("challenge-short", "challenge", proofServerID, proofPrincipalID, proofDeviceID, Data(repeating: 0, count: 31), 0),
    ("challenge-long", "challenge", proofServerID, proofPrincipalID, proofDeviceID, Data(repeating: 0, count: 33), 0),
    ("negative-expiry", "expiry_unix_milliseconds", proofServerID, proofPrincipalID, proofDeviceID, proofChallenge, -1),
]
for (name, field, serverID, principalID, deviceID, challenge, expiry) in invalidProofInputs {
    do {
        _ = try deviceRegistrationProofTranscript(
            serverID: serverID, principalID: principalID, deviceID: deviceID,
            challenge: challenge, expiryUnixMilliseconds: expiry
        )
        fatalError("\(name) was accepted")
    } catch let error as ContractValidationError {
        guard error.field == field,
              error.code == (field == "expiry_unix_milliseconds" ? "minimum" : "invalid_length") else {
            fatalError("\(name) returned the wrong validation error")
        }
    }
}
let incompatible = try CapabilityProfile(
    transport: "quic-tls13", framing: "datagram-v1", media: "encoded",
    audio: "encoded", sourceRateControl: "server", clientDecode: ["hevc-opus"]
)
let gatewayCapability = try CapabilityProfile(
    transport: "quic-tls13", framing: "datagram-v1", media: "encoded",
    audio: "encoded", sourceRateControl: "server", clientDecode: ["hevc-opus", "h264-opus"]
)
do {
    guard try CapabilityProfile.intersection([capability, capability]) == capability else {
        fatalError("matching capability profiles did not intersect")
    }
} catch { fatalError("matching capability profiles did not intersect") }
guard try CapabilityProfile.intersection([gatewayCapability, capability]).clientDecode == ["h264-opus"] else {
    fatalError("ordered registered profile intersection changed")
}
do {
    _ = try CapabilityProfile.intersection([capability, incompatible])
    fatalError("profiles without overlap were accepted")
} catch { }
let valid = try request.encodeJSON()
var unsupported = try JSONSerialization.jsonObject(with: valid) as! [String: Any]
unsupported["version"] = "2"
var downgrade = try JSONSerialization.jsonObject(with: valid) as! [String: Any]
downgrade["version"] = "0"
var unknown = try JSONSerialization.jsonObject(with: valid) as! [String: Any]
unknown["unknown"] = true
for invalid in [
    try JSONSerialization.data(withJSONObject: unsupported),
    try JSONSerialization.data(withJSONObject: downgrade),
    try JSONSerialization.data(withJSONObject: unknown),
    Data("{".utf8),
    valid + Data(" {}".utf8),
] {
    do {
        _ = try TunnelAdmissionRequest.decodeJSON(invalid)
        fatalError("invalid tunnel admission request was accepted")
    } catch { }
}
do {
    _ = try AllocationPolicy(
        minimumKbps: 100, targetKbps: 50, maximumKbps: 25, tier: "standard",
        audience: "audience", protocolValue: "verse", protocolVersion: 1,
        grantTtlSeconds: 60, reservationLeaseSeconds: 300
    )
    fatalError("invalid allocation bounds were accepted")
} catch { }
let displayMode = try DisplayMode(resolutionWidth: 2560, resolutionHeight: 1440, fps: 120)
for invalid in [
    { try DisplayMode(resolutionWidth: 319, resolutionHeight: 1440, fps: 120) },
    { try DisplayMode(resolutionWidth: 2560, resolutionHeight: 199, fps: 120) },
    { try DisplayMode(resolutionWidth: 2560, resolutionHeight: 1440, fps: 241) },
] {
    do {
        _ = try invalid()
        fatalError("invalid display mode was accepted")
    } catch { }
}
let policyFreeV2Request = try SessionRequest(
    clientDeviceId: "device-1", deviceKeyId: "key-1", poolId: "pool-1",
    idempotencyKey: "request-1", requestedDisplayMode: nil
).encodeJSON()
guard !String(data: policyFreeV2Request, encoding: .utf8)!.contains("requested_display_mode") else {
    fatalError("wire-v2 request encoded an absent display mode")
}
let displayRequest = try SessionRequest(
    clientDeviceId: "device-1", deviceKeyId: "key-1", poolId: "pool-1",
    idempotencyKey: "request-1", requestedDisplayMode: displayMode
)
guard try SessionRequest.decodeJSON(displayRequest.encodeJSON()).requestedDisplayMode == displayMode else {
    fatalError("display mode did not round-trip")
}
var nullDisplayRequest = try JSONSerialization.jsonObject(with: displayRequest.encodeJSON()) as! [String: Any]
nullDisplayRequest["requested_display_mode"] = NSNull()
do {
    _ = try SessionRequest.decodeJSON(try JSONSerialization.data(withJSONObject: nullDisplayRequest))
    fatalError("explicit null display mode was accepted")
} catch { }
let nativeIdentity = try NativeSessionIdentity(clientDeviceId: "device-1", deviceKeyId: "key-1")
let browserSession = try BrowserAuthenticatedSession(
    username: "alice", provider: "local", roles: ["user"], role: "user"
)
guard !String(data: try browserSession.encodeJSON(), encoding: .utf8)!.contains("native_identity") else {
    fatalError("browser session encoded native identity")
}
let nativeSession = try NativeAuthenticatedSession(
    username: "alice", provider: "local", roles: ["user"], role: "user", nativeIdentity: nativeIdentity
)
guard try NativeAuthenticatedSession.decodeJSON(nativeSession.encodeJSON()).nativeIdentity == nativeIdentity else {
    fatalError("native session identity did not round-trip")
}
do {
    _ = try BrowserAuthenticatedSession.decodeJSON(nativeSession.encodeJSON())
    fatalError("browser session accepted native identity")
} catch { }
do {
    _ = try NativeAuthenticatedSession.decodeJSON(browserSession.encodeJSON())
    fatalError("native session accepted missing identity")
} catch { }
var partialNativeSession = try JSONSerialization.jsonObject(with: nativeSession.encodeJSON()) as! [String: Any]
partialNativeSession["native_identity"] = ["client_device_id": "device-1"]
do {
    _ = try NativeAuthenticatedSession.decodeJSON(try JSONSerialization.data(withJSONObject: partialNativeSession))
    fatalError("partial native identity was accepted")
} catch { }
for roles in [[""], [String(repeating: "r", count: 65)]] {
    do {
        _ = try BrowserAuthenticatedSession(username: "alice", provider: "local", roles: roles, role: "user")
        fatalError("invalid role item length was accepted")
    } catch { }
}
_ = try NativeTunnelCredential(
    clientDeviceId: "device-1", deviceKeyId: "key-1", certificateChainPem: "certificate",
    trustBundlePem: "trust", expiresAt: "2099-01-01T00:00:00Z"
)
for expiresAt in ["2099-01-01T00:00:00+00:00", "2099-01-01T00:00:00.100Z"] {
    do {
        _ = try NativeTunnelCredential(
            clientDeviceId: "device-1", deviceKeyId: "key-1", certificateChainPem: "certificate",
            trustBundlePem: "trust", expiresAt: expiresAt
        )
        fatalError("noncanonical RFC3339 UTC timestamp was accepted")
    } catch { }
}
let streamPolicy = try ProviderStreamPolicy(
    resolutionWidth: 2560, resolutionHeight: 1440, fps: 120,
    codec: "HEVC", bitrateKbps: 40000, audioEnabled: true
)
guard streamPolicy.codec == "HEVC" else { fatalError("stream policy changed") }
for invalid in [
    { try ProviderStreamPolicy(resolutionWidth: 319, resolutionHeight: 1440, fps: 120, codec: "HEVC", bitrateKbps: 40000, audioEnabled: true) },
    { try ProviderStreamPolicy(resolutionWidth: 2560, resolutionHeight: 1440, fps: 241, codec: "HEVC", bitrateKbps: 40000, audioEnabled: true) },
    { try ProviderStreamPolicy(resolutionWidth: 2560, resolutionHeight: 1440, fps: 120, codec: "VP9", bitrateKbps: 40000, audioEnabled: true) },
] {
    do {
        _ = try invalid()
        fatalError("invalid stream policy was accepted")
    } catch { }
}
let telemetry = try GatewayTelemetry(
    admittedSessions: 1, admissionRejects: 2, reconnects: 3, drainTransitions: 4,
    mediaDrops: 5, mediaPackets: 6, mediaBytes: 7, queueDelayMicros: 8,
    processingDelayMicros: 9, processingSamples: 10, pacingDelayMicros: 11,
    providerErrors: 12, inputRejected: 13, controlRttMicros: 14,
    controlJitterMicros: 15, controlLossPpm: 16, pendingReliable: 17,
    providerState: "ready"
)
guard telemetry.mediaBytes == 7 else { fatalError("gateway telemetry changed") }
do {
    _ = try GatewayTelemetry(
        admittedSessions: 1, admissionRejects: 2, reconnects: 3, drainTransitions: 4,
        mediaDrops: 5, mediaPackets: 6, mediaBytes: 7, queueDelayMicros: 8,
        processingDelayMicros: 9, processingSamples: 10, pacingDelayMicros: 11,
        providerErrors: 12, inputRejected: 13, controlRttMicros: 14,
        controlJitterMicros: 15, controlLossPpm: 1000001, pendingReliable: 17,
        providerState: "ready"
    )
    fatalError("invalid gateway telemetry was accepted")
} catch { }
for text in [
    String(repeating: "a", count: 65536),
    String(repeating: "é", count: 32768),
    String(repeating: "\\\"", count: 32768),
] {
    let clipboard = try GatewayClipboardText(
        direction: "client_to_provider", text: text, encoding: "utf-8",
        loopToken: "abcdefghijklmnop"
    )
    let decoded = try GatewayClipboardText.decodeJSON(clipboard.encodeJSON())
    guard decoded.text == text else { fatalError("clipboard text changed during round-trip") }
}
do {
    _ = try GatewayClipboardText(
        direction: "client_to_provider", text: String(repeating: "a", count: 65537),
        encoding: "utf-8", loopToken: "abcdefghijklmnop"
    )
    fatalError("oversized clipboard text was accepted")
} catch { }
""",
            encoding="utf-8",
        )
        run(["swiftc", str(ROOT / "gen/swift/Protocol.swift"), str(swift), "-o", str(workspace / "swift-contracts")], ROOT)
        run([str(workspace / "swift-contracts")], ROOT)

        rust = workspace / "protocol.rs"
        shutil.copyfile(ROOT / "gen/rust/protocol.rs", rust)
        with rust.open("a", encoding="utf-8") as output:
            output.write(
                """
fn main() {
    assert_eq!(CURRENT_WIRE_VERSION, "2");
    assert_eq!(N_MINUS_1_WIRE_VERSION, "1");
    assert_eq!(N_MINUS_2_WIRE_VERSION, "0");
    let capabilities = CapabilityProfile::new(
        "quic-tls13".into(), "datagram-v1".into(), "encoded".into(),
        "encoded".into(), "server".into(), vec!["h264-opus".into()],
    ).unwrap();
    assert!(CapabilityProfile::new(
        "quic-tls13".into(), "datagram-v2".into(), "encoded".into(),
        "encoded".into(), "server".into(), vec!["h264-opus".into()],
    ).is_ok());
    assert!(CapabilityProfile::new(
        "quic-tls13".into(), "datagram-v3".into(), "encoded".into(),
        "encoded".into(), "server".into(), vec!["h264-opus".into()],
    ).is_err());
    let request = TunnelAdmissionRequest::new(
        "1".into(), "session".into(), "gateway".into(), "audience".into(),
        "g".repeat(43), 0, "n".repeat(16), "s".repeat(86), capabilities.clone(),
    ).unwrap();
    let transcript = "versevdi/tunnel-admission/v17:session7:gateway8:audience43:".to_string()
        + &"g".repeat(43) + "1:016:" + &"n".repeat(16)
        + "10:quic-tls1311:datagram-v17:encoded7:encoded6:server1:19:h264-opus";
    assert_eq!(request.device_admission_transcript(), transcript.into_bytes());
    let proof_server_id = vec![1u8; 16];
    let proof_principal_id = vec![2u8; 16];
    let proof_device_id = vec![3u8; 16];
    let proof_challenge = vec![4u8; 32];
    let proof = device_registration_proof_transcript(
        &proof_server_id, &proof_principal_id, &proof_device_id, &proof_challenge, 1,
    ).unwrap();
    assert_eq!(proof.len(), 112);
    assert_eq!(&proof[..24], b"versevdi-device-proof-v1");
    assert_eq!(&proof[104..], &[0, 0, 0, 0, 0, 0, 0, 1]);
    for (server_id, principal_id, device_id, challenge, expiry, field, code) in [
        (vec![0; 15], proof_principal_id.clone(), proof_device_id.clone(), proof_challenge.clone(), 0, "server_id", "invalid_length"),
        (vec![0; 17], proof_principal_id.clone(), proof_device_id.clone(), proof_challenge.clone(), 0, "server_id", "invalid_length"),
        (proof_server_id.clone(), vec![0; 15], proof_device_id.clone(), proof_challenge.clone(), 0, "principal_id", "invalid_length"),
        (proof_server_id.clone(), vec![0; 17], proof_device_id.clone(), proof_challenge.clone(), 0, "principal_id", "invalid_length"),
        (proof_server_id.clone(), proof_principal_id.clone(), vec![0; 15], proof_challenge.clone(), 0, "device_id", "invalid_length"),
        (proof_server_id.clone(), proof_principal_id.clone(), vec![0; 17], proof_challenge.clone(), 0, "device_id", "invalid_length"),
        (proof_server_id.clone(), proof_principal_id.clone(), proof_device_id.clone(), vec![0; 31], 0, "challenge", "invalid_length"),
        (proof_server_id.clone(), proof_principal_id.clone(), proof_device_id.clone(), vec![0; 33], 0, "challenge", "invalid_length"),
        (proof_server_id.clone(), proof_principal_id.clone(), proof_device_id.clone(), proof_challenge.clone(), -1, "expiry_unix_milliseconds", "minimum"),
    ] {
        assert_eq!(
            device_registration_proof_transcript(&server_id, &principal_id, &device_id, &challenge, expiry),
            Err(ValidationError::new(field, code)),
        );
    }
    assert!(TunnelAdmissionRequest::new(
        "2".into(), "session".into(), "gateway".into(), "audience".into(),
        "g".repeat(43), 0, "n".repeat(16), "s".repeat(86), capabilities.clone(),
    ).is_err());
    assert!(TunnelAdmissionRequest::new(
        "0".into(), "session".into(), "gateway".into(), "audience".into(),
        "g".repeat(43), 0, "n".repeat(16), "s".repeat(86), capabilities.clone(),
    ).is_err());
    assert!(TunnelAdmissionRequest::new(
        "1".into(), "session".into(), "gateway".into(), "audience".into(),
        "g".repeat(43), 0, "short".into(), "s".repeat(86), capabilities.clone(),
    ).is_err());
    assert!(intersect_capability_profiles(&[capabilities.clone(), capabilities.clone()]).is_ok());
    let incompatible = CapabilityProfile::new(
        "quic-tls13".into(), "datagram-v1".into(), "encoded".into(),
        "encoded".into(), "server".into(), vec!["hevc-opus".into()],
    ).unwrap();
    let gateway_capability = CapabilityProfile::new(
        "quic-tls13".into(), "datagram-v1".into(), "encoded".into(),
        "encoded".into(), "server".into(), vec!["hevc-opus".into(), "h264-opus".into()],
    ).unwrap();
    assert_eq!(
        intersect_capability_profiles(&[gateway_capability, capabilities.clone()]).unwrap().clientDecode(),
        &vec!["h264-opus".to_string()],
    );
    assert!(intersect_capability_profiles(&[capabilities, incompatible]).is_err());
    assert!(AllocationPolicy::new(
        100, 50, 25, "standard".into(), "audience".into(), "verse".into(), 1, 60, 300,
    ).is_err());
    let display_mode = DisplayMode::new(2560, 1440, 120).unwrap();
    assert!(DisplayMode::new(319, 1440, 120).is_err());
    assert!(DisplayMode::new(2560, 199, 120).is_err());
    assert!(DisplayMode::new(2560, 1440, 241).is_err());
    let policy_free_v2_request = SessionRequest::new(
        "device-1".into(), "key-1".into(), "pool-1".into(), "request-1".into(),
        None,
    ).unwrap();
    assert!(policy_free_v2_request.requestedDisplayMode().is_none());
    let display_request = SessionRequest::new(
        "device-1".into(), "key-1".into(), "pool-1".into(), "request-1".into(),
        Some(display_mode.clone()),
    ).unwrap();
    assert_eq!(display_request.requestedDisplayMode(), &Some(display_mode));
    let native_identity = NativeSessionIdentity::new("device-1".into(), "key-1".into()).unwrap();
    assert!(BrowserAuthenticatedSession::new(
        "alice".into(), "local".into(), vec!["user".into()], "user".into(),
    ).is_ok());
    assert!(NativeAuthenticatedSession::new(
        "alice".into(), "local".into(), vec!["user".into()], "user".into(), native_identity,
    ).is_ok());
    assert!(BrowserAuthenticatedSession::new(
        "alice".into(), "local".into(), vec![String::new()], "user".into(),
    ).is_err());
    assert!(BrowserAuthenticatedSession::new(
        "alice".into(), "local".into(), vec!["r".repeat(65)], "user".into(),
    ).is_err());
    assert!(NativeTunnelCredential::new(
        "device-1".into(), "key-1".into(), "certificate".into(), "trust".into(),
        "2099-01-01T00:00:00Z".into(),
    ).is_ok());
    for expires_at in ["2099-01-01T00:00:00+00:00", "2099-01-01T00:00:00.100Z"] {
        assert!(NativeTunnelCredential::new(
            "device-1".into(), "key-1".into(), "certificate".into(), "trust".into(),
            expires_at.into(),
        ).is_err());
    }
    assert!(ProviderStreamPolicy::new(
        2560, 1440, 120, "HEVC".into(), 40000, true,
    ).is_ok());
    assert!(ProviderStreamPolicy::new(
        319, 1440, 120, "HEVC".into(), 40000, true,
    ).is_err());
    assert!(ProviderStreamPolicy::new(
        2560, 1440, 241, "HEVC".into(), 40000, true,
    ).is_err());
    assert!(ProviderStreamPolicy::new(
        2560, 1440, 120, "VP9".into(), 40000, true,
    ).is_err());
    assert!(GatewayTelemetry::new(
        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, "ready".into(),
    ).is_ok());
    assert!(GatewayTelemetry::new(
        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 1000001, 17, "ready".into(),
    ).is_err());
    assert!(GatewayTelemetry::new(
        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, "provider.example:47984".into(),
    ).is_err());
    for text in [
        "a".repeat(65536),
        "é".repeat(32768),
        "\\\"".repeat(32768),
    ] {
        assert!(GatewayClipboardText::new(
            "client_to_provider".into(), text, "utf-8".into(), "abcdefghijklmnop".into(),
        ).is_ok());
    }
    assert!(GatewayClipboardText::new(
        "client_to_provider".into(), "a".repeat(65537), "utf-8".into(),
        "abcdefghijklmnop".into(),
    ).is_err());
}
"""
            )
        run(["rustc", str(rust), "-o", str(workspace / "rust-contracts")], ROOT)
        run([str(workspace / "rust-contracts")], ROOT)
        rust_unknown = workspace / "unknown.rs"
        shutil.copyfile(ROOT / "gen/rust/protocol.rs", rust_unknown)
        with rust_unknown.open("a", encoding="utf-8") as output:
            output.write("\nfn main() { let _ = CapabilityProfile { unknown: String::new() }; }\n")
        run_failure(["rustc", str(rust_unknown), "-o", str(workspace / "rust-unknown")], ROOT, "no field named `unknown`")
    print("Generated strict contract checks passed")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())