import Foundation
import Testing

@testable import NucleicProtocol

@Suite struct NoiseTests {
    // Run a handshake to completion by ping-ponging messages between two drivers.
    private func runHandshake(_ a: HandshakeDriver, _ b: HandshakeDriver) throws {
        var initiatorsTurn = a  // a is the initiator
        var other = b
        while !a.isComplete || !b.isComplete {
            let msg = try initiatorsTurn.write()
            _ = try other.read(msg)
            swap(&initiatorsTurn, &other)
        }
    }

    @Test func xxPairingHandshakeEstablishesAndExchanges() throws {
        let host = DeviceIdentity()
        let phone = DeviceIdentity()
        let psk = PairingPayload.freshSecret()

        let phoneDriver = NoiseSession.makeInitiator(
            identity: phone, hostStaticKey: host.staticPublicKey, psk: psk)
        let hostDriver = NoiseSession.makeResponder(identity: host, psk: psk)

        try runHandshake(phoneDriver, hostDriver)

        let phoneSession = try phoneDriver.session()
        let hostSession = try hostDriver.session()

        // Each side learned the other's real static key.
        #expect(phoneSession.remoteStaticKey == host.staticPublicKey)
        #expect(hostSession.remoteStaticKey == phone.staticPublicKey)
        // Channel-binding hash matches on both ends.
        #expect(phoneSession.handshakeHash == hostSession.handshakeHash)

        // Phone → Host application message survives the encrypted channel.
        let hello = ClientMsg.hello(Hello(deviceID: "dev", deviceLabel: "iPhone"))
        let frame = try phoneSession.seal(hello)
        let decoded = try hostSession.open(ClientMsg.self, from: frame)
        #expect(decoded == hello)

        // Host → Phone, the other direction (host seals with its send key, phone opens
        // with its matching recv key).
        let welcome = HostMsg.welcome(Welcome(
            grantedScope: .approve, host: HostInfo(hostID: "h", hostName: "Mac"),
            capabilities: WireCapabilities(canModifyToolInput: true, allowAlwaysScopes: [])))
        let phoneGot = try phoneSession.open(HostMsg.self, from: hostSession.seal(welcome))
        #expect(phoneGot == welcome)

        // Nonce counters advance per direction: a second phone→host message still opens.
        let ping = try hostSession.open(ClientMsg.self, from: phoneSession.seal(ClientMsg.ping))
        #expect(ping == .ping)
    }

    @Test func ikReconnectHandshakeEstablishes() throws {
        let host = DeviceIdentity()
        let phone = DeviceIdentity()

        // No PSK → IK pattern; the phone pins the host static from a prior pairing.
        let phoneDriver = NoiseSession.makeInitiator(
            identity: phone, hostStaticKey: host.staticPublicKey, psk: nil)
        let hostDriver = NoiseSession.makeResponder(identity: host, psk: nil)

        try runHandshake(phoneDriver, hostDriver)

        let phoneSession = try phoneDriver.session()
        let hostSession = try hostDriver.session()
        #expect(hostSession.remoteStaticKey == phone.staticPublicKey)

        let msg = ClientMsg.subscribe(Subscribe(
            sessionID: SessionID(rawValue: "s1"), sinceSeq: 10, verbosity: .full))
        let got = try hostSession.open(ClientMsg.self, from: phoneSession.seal(msg))
        #expect(got == msg)
    }

    @Test func mismatchedPairingSecretFailsHandshake() throws {
        let host = DeviceIdentity()
        let phone = DeviceIdentity()

        let phoneDriver = NoiseSession.makeInitiator(
            identity: phone, hostStaticKey: host.staticPublicKey, psk: PairingPayload.freshSecret())
        let hostDriver = NoiseSession.makeResponder(identity: host, psk: PairingPayload.freshSecret())

        // msg1 (psk,e) carries no decryptable payload, but msg2's `s` is AEAD-protected by a
        // key derived from the PSK — the responder's reply fails to open on the initiator.
        #expect(throws: (any Error).self) {
            let m1 = try phoneDriver.write()
            _ = try hostDriver.read(m1)
            let m2 = try hostDriver.write()
            _ = try phoneDriver.read(m2)   // expected to throw: wrong PSK
        }
    }

    @Test func tamperedFrameFailsToOpen() throws {
        let host = DeviceIdentity()
        let phone = DeviceIdentity()
        let psk = PairingPayload.freshSecret()
        let phoneDriver = NoiseSession.makeInitiator(
            identity: phone, hostStaticKey: host.staticPublicKey, psk: psk)
        let hostDriver = NoiseSession.makeResponder(identity: host, psk: psk)
        try runHandshake(phoneDriver, hostDriver)
        let phoneSession = try phoneDriver.session()
        let hostSession = try hostDriver.session()

        var frame = try phoneSession.seal(ClientMsg.ping)
        frame[frame.startIndex] ^= 0xff   // flip a byte
        #expect(throws: (any Error).self) {
            _ = try hostSession.open(ClientMsg.self, from: frame)
        }
    }

    @Test func identityRawRoundTrips() throws {
        let id = DeviceIdentity()
        let restored = try DeviceIdentity(importingRaw: id.exportRaw())
        #expect(restored.staticPublicKey == id.staticPublicKey)
        #expect(restored.signingPublicKey == id.signingPublicKey)
    }

    @Test func pairingQRRoundTrips() throws {
        let host = DeviceIdentity()
        let payload = PairingPayload(
            hostName: "Andrew's Mac", hostStaticKey: host.staticPublicKey,
            pairingSecret: PairingPayload.freshSecret(), lanHost: "192.168.1.5", lanPort: 51_900)
        let restored = try PairingPayload(qrString: payload.qrString())
        #expect(restored == payload)
    }

    @Test func ed25519SignVerify() throws {
        let id = DeviceIdentity()
        let data = Data("pair-this-qr".utf8)
        let sig = try id.sign(data)
        #expect(DeviceIdentity.verify(sig, of: data, publicKey: id.signingPublicKey))
        #expect(!DeviceIdentity.verify(sig, of: Data("other".utf8), publicKey: id.signingPublicKey))
    }
}
