import CryptoKit
import XCTest

@testable import MultiSimAdmin

/// The pairing exchange, checked against a node-shaped peer.
///
/// Pairing is the one moment the phone is inside the network and the only time the node's password
/// is involved at all, so the interesting properties are that the phone sends its real public keys,
/// that it records the id the node assigns, and that each failure mode is distinguishable — a wrong
/// password and a closed window call for different actions.
final class RelayPairingTests: XCTestCase {
    private final class PairingStub: URLProtocol, @unchecked Sendable {
        nonisolated(unsafe) static var handler: (@Sendable (URLRequest, Data) throws -> (Int, Data))?
        nonisolated(unsafe) static var lastBody: Data?
        nonisolated(unsafe) static var lastURL: URL?

        override class func canInit(with request: URLRequest) -> Bool { true }
        override class func canonicalRequest(for request: URLRequest) -> URLRequest { request }

        override func startLoading() {
            let body = request.httpBody ?? Self.drain(request.httpBodyStream) ?? Data()
            Self.lastBody = body
            Self.lastURL = request.url
            guard let handler = Self.handler else {
                client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse))
                return
            }
            do {
                let (status, data) = try handler(request, body)
                let response = HTTPURLResponse(
                    url: request.url!, statusCode: status, httpVersion: "HTTP/1.1", headerFields: nil
                )!
                client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
                client?.urlProtocol(self, didLoad: data)
                client?.urlProtocolDidFinishLoading(self)
            } catch {
                client?.urlProtocol(self, didFailWithError: error)
            }
        }

        override func stopLoading() {}

        private static func drain(_ stream: InputStream?) -> Data? {
            guard let stream else { return nil }
            stream.open()
            defer { stream.close() }
            var data = Data()
            var buffer = [UInt8](repeating: 0, count: 4096)
            while stream.hasBytesAvailable {
                let read = stream.read(&buffer, maxLength: buffer.count)
                if read <= 0 { break }
                data.append(buffer, count: read)
            }
            return data
        }
    }

    private func makePairing() -> RelayPairing {
        let configuration = URLSessionConfiguration.ephemeral
        configuration.protocolClasses = [PairingStub.self]
        return RelayPairing(session: URLSession(configuration: configuration))
    }

    private func makeIdentity() -> DeviceIdentity {
        DeviceIdentity(
            deviceID: "",
            signingKey: Curve25519.Signing.PrivateKey(),
            agreementKey: Curve25519.KeyAgreement.PrivateKey()
        )
    }

    private func successBody(agentSigning: Data, agentKx: Data) -> Data {
        Data("""
        {
          "nodeID": "nd_7f3a9c2e",
          "appID": "ap_5b1d0000",
          "agentSignPub": "\(agentSigning.base64EncodedString())",
          "agentKxPub": "\(agentKx.base64EncodedString())",
          "relayURL": "https://relay.example",
          "deviceToken": "0123456789abcdef0123456789abcdef"
        }
        """.utf8)
    }

    override func tearDown() {
        PairingStub.handler = nil
        PairingStub.lastBody = nil
        PairingStub.lastURL = nil
        super.tearDown()
    }

    func testPairingSendsTheRealPublicKeysAndReturnsTheAssignedIdentity() async throws {
        let identity = makeIdentity()
        let agentSigning = Curve25519.Signing.PrivateKey()
        let agentKx = Curve25519.KeyAgreement.PrivateKey()
        PairingStub.handler = { _, _ in
            (200, self.successBody(
                agentSigning: agentSigning.publicKey.rawRepresentation,
                agentKx: agentKx.publicKey.rawRepresentation
            ))
        }

        let result = try await makePairing().pair(
            host: "192.168.1.10",
            password: "hunter2",
            deviceName: "chick",
            identity: identity
        )

        // The keys the phone sent must be the ones it will sign and decrypt with; the node stores
        // them as its trust anchor, and a mismatch here is an unverifiable request later.
        let sent = try XCTUnwrap(
            try JSONSerialization.jsonObject(with: try XCTUnwrap(PairingStub.lastBody)) as? [String: Any]
        )
        XCTAssertEqual(sent["appSignPub"] as? String, identity.signingPublicKey.base64EncodedString())
        XCTAssertEqual(sent["appKxPub"] as? String, identity.agreementPublicKey.base64EncodedString())
        XCTAssertEqual(sent["password"] as? String, "hunter2")
        XCTAssertEqual(sent["deviceName"] as? String, "chick")
        XCTAssertEqual(PairingStub.lastURL?.absoluteString, "http://192.168.1.10:3001/provision")

        // The id is the node's choice, and the returned keys become the phone's trust anchor.
        XCTAssertEqual(result.deviceID, "ap_5b1d0000")
        XCTAssertEqual(result.endpoint.nodeID, "nd_7f3a9c2e")
        XCTAssertEqual(result.endpoint.deviceID, "ap_5b1d0000")
        XCTAssertEqual(result.endpoint.agentSigningKey, agentSigning.publicKey.rawRepresentation)
        XCTAssertEqual(result.endpoint.agentKxKey, agentKx.publicKey.rawRepresentation)
        XCTAssertEqual(result.endpoint.deviceToken, "0123456789abcdef0123456789abcdef")
        XCTAssertEqual(result.endpoint.relayBaseURL.absoluteString, "https://relay.example")
    }

    /// Each failure mode must be its own error: "wrong password" is a retype, "window closed" is a
    /// walk to the node, and neither is fixed by the other's action.
    func testFailuresAreDistinguishable() async {
        let identity = makeIdentity()
        let cases: [(Int, Data, RelayPairing.Failure)] = [
            (401, Data(#"{"error":"no"}"#.utf8), .rejected),
            (410, Data(#"{"error":"closed"}"#.utf8), .windowClosed),
            (500, Data(#"{"error":"boom"}"#.utf8), .malformedResponse),
        ]
        for (status, body, expected) in cases {
            PairingStub.handler = { _, _ in (status, body) }
            do {
                _ = try await makePairing().pair(
                    host: "192.168.1.10", password: "x", deviceName: "d", identity: identity
                )
                XCTFail("status \(status) was accepted")
            } catch let failure as RelayPairing.Failure {
                XCTAssertEqual(failure, expected, "status \(status)")
            } catch {
                XCTFail("status \(status) produced \(error)")
            }
        }
    }

    /// A malformed or truncated response must be refused rather than half-accepted: an endpoint
    /// missing its keys would fail later as an unreadable response, which the user reads as a broken
    /// node rather than as a bad pairing.
    func testAMalformedResponseIsRefused() async {
        let identity = makeIdentity()
        let good = Curve25519.Signing.PrivateKey().publicKey.rawRepresentation
        let goodKx = Curve25519.KeyAgreement.PrivateKey().publicKey.rawRepresentation
        let bodies: [Data] = [
            Data("not json".utf8),
            Data(#"{"nodeID":"nd_7f3a9c2e"}"#.utf8),
            // A key that is not 32 bytes decodes fine as base64 and is wrong for the format.
            Data("""
            {"nodeID":"nd_7f3a9c2e","appID":"ap_5b1d0000","agentSignPub":"\(Data(repeating: 1, count: 31).base64EncodedString())",
             "agentKxPub":"\(goodKx.base64EncodedString())","relayURL":"https://relay.example","deviceToken":"t"}
            """.utf8),
            // A plain-HTTP relay would make the end-to-end encryption decorative.
            Data("""
            {"nodeID":"nd_7f3a9c2e","appID":"ap_5b1d0000","agentSignPub":"\(good.base64EncodedString())",
             "agentKxPub":"\(goodKx.base64EncodedString())","relayURL":"http://relay.example","deviceToken":"t"}
            """.utf8),
            // No token: the phone could never authenticate to the relay.
            Data("""
            {"nodeID":"nd_7f3a9c2e","appID":"ap_5b1d0000","agentSignPub":"\(good.base64EncodedString())",
             "agentKxPub":"\(goodKx.base64EncodedString())","relayURL":"https://relay.example","deviceToken":""}
            """.utf8),
        ]
        for body in bodies {
            PairingStub.handler = { _, _ in (200, body) }
            do {
                _ = try await makePairing().pair(
                    host: "192.168.1.10", password: "x", deviceName: "d", identity: identity
                )
                XCTFail("accepted \(String(decoding: body, as: UTF8.self))")
            } catch let failure as RelayPairing.Failure {
                XCTAssertEqual(failure, .malformedResponse)
            } catch {
                XCTFail("produced \(error)")
            }
        }
    }

    func testAnUnreachableNodeIsReportedAsSuch() async {
        PairingStub.handler = { _, _ in throw URLError(.cannotConnectToHost) }
        do {
            _ = try await makePairing().pair(
                host: "192.168.1.10", password: "x", deviceName: "d", identity: makeIdentity()
            )
            XCTFail("an unreachable node was accepted")
        } catch let failure as RelayPairing.Failure {
            XCTAssertEqual(failure, .unreachable)
        } catch {
            XCTFail("produced \(error)")
        }
    }
}

/// The identity file is the one thing a pairing cannot be redone without.
final class DeviceIdentityStoreTests: XCTestCase {
    private var location: URL!

    override func setUp() {
        super.setUp()
        location = FileManager.default.temporaryDirectory
            .appendingPathComponent("msa-identity-\(UUID().uuidString).json")
    }

    override func tearDown() {
        try? FileManager.default.removeItem(at: location)
        super.tearDown()
    }

    func testTheIdentityIsStableAcrossLoads() throws {
        let store = FileDeviceIdentityStore(location: location)
        let first = try store.loadOrCreate()
        let second = try store.loadOrCreate()
        XCTAssertEqual(first.signingPublicKey, second.signingPublicKey)
        XCTAssertEqual(first.agreementPublicKey, second.agreementPublicKey)
        XCTAssertTrue(first.deviceID.isEmpty, "the phone cannot know its id before the node assigns one")

        // A fresh store over the same file must find the same keys, or a restart would silently
        // invalidate every pairing.
        let reopened = FileDeviceIdentityStore(location: location)
        XCTAssertEqual(try reopened.loadOrCreate().signingPublicKey, first.signingPublicKey)
    }

    /// The node verified the public keys it was given, so recording its id must not regenerate them.
    func testAssigningTheDeviceIDKeepsTheKeys() throws {
        let store = FileDeviceIdentityStore(location: location)
        let original = try store.loadOrCreate()
        let assigned = try store.assignDeviceID("ap_5b1d0000")
        XCTAssertEqual(assigned.deviceID, "ap_5b1d0000")
        XCTAssertEqual(assigned.signingPublicKey, original.signingPublicKey)
        XCTAssertEqual(assigned.agreementPublicKey, original.agreementPublicKey)

        let reloaded = try FileDeviceIdentityStore(location: location).loadOrCreate()
        XCTAssertEqual(reloaded.deviceID, "ap_5b1d0000")
        XCTAssertEqual(reloaded.signingPublicKey, original.signingPublicKey)
    }

    func testTheFileIsNotWorldReadable() throws {
        _ = try FileDeviceIdentityStore(location: location).loadOrCreate()
        let mode = try FileManager.default
            .attributesOfItem(atPath: location.path)[.posixPermissions] as? NSNumber
        XCTAssertEqual(mode?.int16Value, 0o600)
    }

    func testResetForgetsTheIdentity() throws {
        let store = FileDeviceIdentityStore(location: location)
        let first = try store.loadOrCreate()
        try store.reset()
        let second = try store.loadOrCreate()
        XCTAssertNotEqual(first.signingPublicKey, second.signingPublicKey)
    }
}

/// The router is the only place that decides which path a request takes.
final class NodeTransportRouterTests: XCTestCase {
    private struct Recorder: NodeTransport {
        let label: String
        let recorder: @Sendable (String) -> Void
        func send(
            _ route: UpstreamRoute,
            instanceID: UUID,
            endpoint: InstanceEndpoint
        ) async throws -> UpstreamResponse {
            recorder(label)
            return UpstreamResponse(status: 200, data: Data())
        }
    }

    func testTheRouterPicksThePathFromTheEndpoint() async throws {
        let seen = LockedStrings()
        let router = NodeTransportRouter(
            direct: Recorder(label: "direct") { seen.append($0) },
            relay: Recorder(label: "relay") { seen.append($0) }
        )

        _ = try await router.send(
            .health,
            instanceID: UUID(),
            endpoint: .direct(URL(string: "http://192.168.1.20:3000")!)
        )
        _ = try await router.send(
            .health,
            instanceID: UUID(),
            endpoint: .relay(InstanceEndpointMigrationTests.relayEndpoint(nodeID: "nd_7f3a9c2e"))
        )
        XCTAssertEqual(seen.values, ["direct", "relay"])
    }

    /// Handing a relay instance to the direct path is a wiring bug, not a network condition, so it
    /// must fail rather than dialing something arbitrary.
    func testTheDirectPathRefusesARelayEndpoint() async {
        let transport = DirectTransport(client: NoopClient(), sessions: InstanceSessionStore())
        do {
            _ = try await transport.send(
                .health,
                instanceID: UUID(),
                endpoint: .relay(InstanceEndpointMigrationTests.relayEndpoint(nodeID: "nd_7f3a9c2e"))
            )
            XCTFail("the direct path accepted a relay endpoint")
        } catch let error as UpstreamError {
            XCTAssertEqual(error, .invalidRequest)
        } catch {
            XCTFail("produced \(error)")
        }
    }

    private struct NoopClient: SimAdminClientProtocol {
        func send(_ route: UpstreamRoute, baseURL: URL, cookie: String?) async throws -> UpstreamResponse {
            XCTFail("the direct client must not be dialed for a relay endpoint")
            return UpstreamResponse(status: 200, data: Data())
        }
    }

    private final class LockedStrings: @unchecked Sendable {
        private let lock = NSLock()
        private var storage: [String] = []

        func append(_ value: String) {
            lock.withLock { storage.append(value) }
        }

        var values: [String] {
            lock.withLock { storage }
        }
    }
}