import CryptoKit
import XCTest

@testable import MultiSimAdmin

/// The phone's half of the protocol, checked against a peer that behaves like the Go agent.
///
/// The peer here is built from the same primitives the agent uses, in the opposite direction: it
/// verifies the request signature with the phone's public key, decrypts the payload with the agent's
/// private X25519 key, recomputes the confirmation intent, and answers with a signature and a
/// ciphertext the phone has to accept. A test that only stubbed the transport would prove the
/// plumbing; this proves the protocol, which is the part that has three implementations and no
/// shared code.
final class RelayTransportTests: XCTestCase {
    // MARK: - Agent-shaped peer

    /// Stands in for the node's agent and for the relay at once: the relay is a dumb pipe, so the
    /// only interesting behaviour is what the agent does with the envelope.
    private struct AgentPeer {
        let nodeID: String
        let deviceID: String
        let signingKey: Curve25519.Signing.PrivateKey
        let agreementKey: Curve25519.KeyAgreement.PrivateKey

        init(nodeID: String = "nd_7f3a9c2e", deviceID: String = "ap_5b1d0000") {
            self.nodeID = nodeID
            self.deviceID = deviceID
            self.signingKey = Curve25519.Signing.PrivateKey()
            self.agreementKey = Curve25519.KeyAgreement.PrivateKey()
        }

        /// What the phone needs in order to reach this peer.
        func endpoint(relayBaseURL: URL = URL(string: "https://relay.example")!) -> RelayEndpoint {
            RelayEndpoint(
                relayBaseURL: relayBaseURL,
                nodeID: nodeID,
                agentSigningKey: signingKey.publicKey.rawRepresentation,
                agentKxKey: agreementKey.publicKey.rawRepresentation,
                deviceToken: "device-token-abcdef",
                deviceID: deviceID,
                nodeName: "Lab Alpha"
            )
        }

        struct Received {
            var inner: RelayWire.InnerRequest
            var envelope: RelayWire.Envelope
        }

        /// Verifies and decrypts one request the way the agent does, in the agent's order.
        func receive(_ raw: Data, phonePublicKey: Curve25519.Signing.PublicKey) throws -> Received {
            let envelope = try JSONDecoder().decode(RelayWire.Envelope.self, from: raw)
            try RelayWire.verifyRequest(envelope, publicKey: phonePublicKey)
            let shared = try RelayWire.sharedSecret(privateKey: agreementKey, peerPublic: envelope.ek)
            let (key, iv) = try RelayWire.deriveKey(
                shared: shared,
                nodeID: nodeID,
                requestID: envelope.req,
                response: false
            )
            let plaintext = try RelayWire.open(envelope.p, key: key, iv: iv)
            let inner = try JSONDecoder().decode(RelayWire.InnerRequest.self, from: plaintext)
            return Received(inner: inner, envelope: envelope)
        }

        /// Answers the way the agent does: a fresh ephemeral key, encrypted to the phone's request
        /// ephemeral, signed with the agent's long-term key, and with `app` omitted from the signed
        /// canonical form.
        func answer(
            to request: RelayWire.Envelope,
            status: Int,
            outcome: String?,
            body: Data?
        ) throws -> Data {
            let ephemeral = Curve25519.KeyAgreement.PrivateKey()
            let shared = try RelayWire.sharedSecret(
                privateKey: ephemeral,
                peerPublic: request.ek
            )
            let (key, iv) = try RelayWire.deriveKey(
                shared: shared,
                nodeID: nodeID,
                requestID: request.req,
                response: true
            )
            let inner = RelayWire.InnerResponse(status: status, outcome: outcome, body: body)
            let plaintext = try JSONEncoder().encode(inner)
            var envelope = RelayWire.Envelope(
                node: nodeID,
                req: request.req,
                app: nil,
                ts: request.ts + 1,
                ek: ephemeral.publicKey.rawRepresentation,
                p: try RelayWire.seal(plaintext, key: key, iv: iv)
            )
            envelope.s = try signingKey.signature(for: RelayWire.canonical(envelope, response: true))
            return try JSONEncoder().encode(envelope)
        }
    }

    /// A URLProtocol stub, so the transport under test runs its real code path: real envelope
    /// construction, real signing, real HTTP handling.
    private final class RelayStub: URLProtocol, @unchecked Sendable {
        nonisolated(unsafe) static var handler: (@Sendable (URLRequest, Data) throws -> (Int, Data))?
        nonisolated(unsafe) static var lastRequest: URLRequest?
        nonisolated(unsafe) static var lastBody: Data?

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

        override func startLoading() {
            // `URLProtocol` hands the body as a stream, so it has to be drained here to get bytes.
            let body = request.httpBody ?? Self.drain(request.httpBodyStream)
            Self.lastRequest = request
            Self.lastBody = body
            guard let handler = Self.handler else {
                client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse))
                return
            }
            do {
                let (status, data) = try handler(request, body ?? Data())
                let response = HTTPURLResponse(
                    url: request.url!, statusCode: status,
                    httpVersion: "HTTP/1.1",
                    headerFields: ["content-type": "application/problem+json"]
                )!
                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 makeTransport(
        identity: DeviceIdentity,
        now: @escaping @Sendable () -> Date = { Date() }
    ) -> RelayTransport {
        let configuration = URLSessionConfiguration.ephemeral
        configuration.protocolClasses = [RelayStub.self]
        return RelayTransport(
            identity: identity,
            session: URLSession(configuration: configuration),
            now: now
        )
    }

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

    override func tearDown() {
        RelayStub.handler = nil
        RelayStub.lastRequest = nil
        RelayStub.lastBody = nil
        super.tearDown()
    }

    // MARK: - Round trip

    func testAReadRoundTripsThroughAnAgentShapedPeer() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        var received: AgentPeer.Received?
        RelayStub.handler = { request, body in
            received = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            let answer = try peer.answer(
                to: received!.envelope,
                status: 200,
                outcome: AgentOutcome.ok.rawValue,
                body: Data(#"{"status":"up"}"#.utf8)
            )
            return (200, answer)
        }

        let transport = makeTransport(identity: identity)
        let response = try await transport.send(
            .health,
            instanceID: UUID(),
            endpoint: .relay(peer.endpoint())
        )

        XCTAssertEqual(response.status, 200)
        XCTAssertEqual(String(decoding: response.data, as: UTF8.self), #"{"status":"up"}"#)
        let seen = try XCTUnwrap(received)
        XCTAssertEqual(seen.inner.m, "GET")
        XCTAssertEqual(seen.inner.p, "/api/health")
        XCTAssertNil(seen.inner.intent, "a read must not carry a confirmation intent")
        XCTAssertEqual(seen.envelope.app, "ap_5b1d0000")
    }

    /// The path carries its query string, and the agent forwards it verbatim: the phone must not
    /// assemble one from parts, because a part the agent reassembled is a part a relay could steer.
    func testThePathCarriesItsQueryString() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        var path: String?
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            path = seen.inner.p
            return (200, try peer.answer(to: seen.envelope, status: 200, outcome: "ok", body: Data("{}".utf8)))
        }

        let transport = makeTransport(identity: identity)
        _ = try await transport.send(
            .smsList(limit: 25, offset: 50, direction: "incoming"),
            instanceID: UUID(),
            endpoint: .relay(peer.endpoint())
        )
        XCTAssertEqual(path, "/api/sms/list?limit=25&offset=50&direction=incoming")
    }

    /// A dangerous write carries an intent that the agent recomputes from the bytes it is about to
    /// forward. If the phone's digest were computed over different bytes, the agent would refuse it
    /// — which is the property that makes a relay unable to edit an SMS in flight.
    func testADangerousWriteCarriesAnIntentTheAgentRecomputes() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        var inner: RelayWire.InnerRequest?
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            inner = seen.inner
            return (200, try peer.answer(to: seen.envelope, status: 200, outcome: "ok", body: Data("{}".utf8)))
        }

        let transport = makeTransport(identity: identity)
        _ = try await transport.send(
            .smsSend(phoneNumber: "+8613800000000", content: "hello"),
            instanceID: UUID(),
            endpoint: .relay(peer.endpoint())
        )

        let seen = try XCTUnwrap(inner)
        let intent = try XCTUnwrap(seen.intent, "a dangerous write must carry an intent")
        let recomputed = try RelayWire.intentDigest(
            method: "POST",
            path: "/api/sms/send",
            body: seen.body
        )
        XCTAssertEqual(intent, recomputed, "the agent recomputes this digest and refuses a mismatch")

        // And the digest must not survive an edit to the body, which is the attack it exists for.
        let edited = try RelayWire.intentDigest(
            method: "POST",
            path: "/api/sms/send",
            body: Data(#"{"phone_number":"+8613800000001","content":"hello"}"#.utf8)
        )
        XCTAssertNotEqual(intent, edited)
    }

    /// The relay must not be able to answer without the agent's key. This is the check that makes
    /// "the relay cannot forge a response" a fact rather than an assumption.
    func testAnUnsignedAnswerIsRejected() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            let raw = try peer.answer(to: seen.envelope, status: 200, outcome: "ok", body: Data("{}".utf8))
            var envelope = try JSONDecoder().decode(RelayWire.Envelope.self, from: raw)
            // A forged answer: correctly shaped, wrong signer.
            let forger = Curve25519.Signing.PrivateKey()
            envelope.s = try forger.signature(for: RelayWire.canonical(envelope, response: true))
            return (200, try JSONEncoder().encode(envelope))
        }

        let transport = makeTransport(identity: identity)
        await assertThrows(.relayResponseUnverifiable) {
            try await transport.send(.health, instanceID: UUID(), endpoint: .relay(peer.endpoint()))
        }
    }

    /// A relay that swaps in a different ciphertext cannot produce a valid signature over it either,
    /// since the signature covers the ciphertext.
    func testATamperedCiphertextIsRejected() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            let raw = try peer.answer(to: seen.envelope, status: 200, outcome: "ok", body: Data("{}".utf8))
            var envelope = try JSONDecoder().decode(RelayWire.Envelope.self, from: raw)
            var ciphertext = envelope.p
            ciphertext[ciphertext.count - 1] ^= 0x01
            envelope.p = ciphertext
            return (200, try JSONEncoder().encode(envelope))
        }

        let transport = makeTransport(identity: identity)
        await assertThrows(.relayResponseUnverifiable) {
            try await transport.send(.health, instanceID: UUID(), endpoint: .relay(peer.endpoint()))
        }
    }

    /// An answer to a different request id must be discarded: the signature covers `req`, so this
    /// catches both a confused relay and a replay of an earlier answer.
    func testAnAnswerForADifferentRequestIsRejected() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            var misdirected = seen.envelope
            misdirected.req = UUID().uuidString.lowercased()
            // Re-signing with the real key would make it valid; the point is that the phone checks
            // the id before it checks anything else, so an unsigned mismatch is enough.
            misdirected.s = Data(repeating: 0, count: 64)
            let raw = try JSONEncoder().encode(misdirected)
            return (200, raw)
        }

        let transport = makeTransport(identity: identity)
        await assertThrows(.relayResponseMismatch) {
            try await transport.send(.health, instanceID: UUID(), endpoint: .relay(peer.endpoint()))
        }
    }

    // MARK: - Error mapping

    /// Every relay error code must arrive as its own typed error. Collapsing them would put the
    /// user back to one message for four different situations.
    func testRelayProblemCodesMapToDistinctErrors() {
        let cases: [(RelayErrorCode, UpstreamError)] = [
            (.unauthorized, .relayRejectedDevice),
            (.nodeUnknown, .relayNodeUnknown),
            (.agentOffline, .agentOffline),
            (.mailboxFull, .nodeBusy),
            (.agentTimeout, .agentTimeout),
            (.badEnvelope, .relayBadEnvelope),
            (.payloadTooLarge, .payloadTooLarge),
            (.rateLimited, .relayRateLimited),
            (.internalError, .relayUnreachable),
        ]
        for (code, expected) in cases {
            let body = Data(#"{"code":"\#(code.rawValue)","status":500}"#.utf8)
            XCTAssertEqual(RelayTransport.problem(from: body, status: 500), expected, code.rawValue)
        }
        // An unrecognised code on a 5xx is a relay fault; on a 4xx it is a protocol mismatch.
        XCTAssertEqual(
            RelayTransport.problem(from: Data(#"{"code":"something_new"}"#.utf8), status: 503),
            .relayUnreachable
        )
        XCTAssertEqual(
            RelayTransport.problem(from: Data(#"{"code":"something_new"}"#.utf8), status: 400),
            .invalidResponse
        )
    }

    func testAnAgentRefusalArrivesAsItsOwnError() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        let cases: [(String, UpstreamError)] = [
            (AgentOutcome.nodeAuthFailed.rawValue, .nodeAuthFailed),
            (AgentOutcome.policyDenied.rawValue, .relayPolicyDenied),
            (AgentOutcome.confirmationMismatch.rawValue, .confirmationRejected),
            (AgentOutcome.expired.rawValue, .envelopeExpired),
            (AgentOutcome.replay.rawValue, .envelopeExpired),
            (AgentOutcome.rateLimited.rawValue, .nodeRateLimited),
            (AgentOutcome.badSignature.rawValue, .envelopeRejected),
        ]
        for (outcome, expected) in cases {
            RelayStub.handler = { _, body in
                let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
                return (200, try peer.answer(to: seen.envelope, status: 403, outcome: outcome, body: nil))
            }
            let transport = makeTransport(identity: identity)
            await assertThrows(expected, "outcome \(outcome)") {
                try await transport.send(.health, instanceID: UUID(), endpoint: .relay(peer.endpoint()))
            }
        }
    }

    /// An upstream error from the node is the one branch where the node's own status reaches the
    /// caller: the agent did its job, SimAdmin answered with an error, and the status is meaningful.
    func testAnUpstreamErrorKeepsTheNodeStatus() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            return (200, try peer.answer(
                to: seen.envelope,
                status: 502,
                outcome: AgentOutcome.upstreamError.rawValue,
                body: Data(#"{"error":"upstream"}"#.utf8)
            ))
        }

        let transport = makeTransport(identity: identity)
        let response = try await transport.send(.health, instanceID: UUID(), endpoint: .relay(peer.endpoint()))
        XCTAssertEqual(response.status, 502)
        XCTAssertEqual(String(decoding: response.data, as: UTF8.self), #"{"error":"upstream"}"#)
    }

    /// The relay is a bearer-authenticated hop, and the token must travel as a bearer header rather
    /// than in the envelope: it is the relay's abuse control, not part of the signed material.
    func testTheRequestCarriesTheDeviceBearerAndNeverAPassword() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        RelayStub.handler = { _, body in
            let seen = try peer.receive(body, phonePublicKey: identity.signingKey.publicKey)
            return (200, try peer.answer(to: seen.envelope, status: 200, outcome: "ok", body: Data("{}".utf8)))
        }

        let transport = makeTransport(identity: identity)
        _ = try await transport.send(
            .smsSend(phoneNumber: "+8613800000000", content: "hello"),
            instanceID: UUID(),
            endpoint: .relay(peer.endpoint())
        )

        let request = try XCTUnwrap(RelayStub.lastRequest)
        XCTAssertEqual(request.url?.path, "/v1/request")
        XCTAssertEqual(request.httpMethod, "POST")
        XCTAssertEqual(
            request.value(forHTTPHeaderField: "authorization"),
            "Bearer device-token-abcdef"
        )
        // The node password never appears anywhere in what leaves the phone.
        let body = String(decoding: try XCTUnwrap(RelayStub.lastBody), as: UTF8.self)
        XCTAssertFalse(body.lowercased().contains("password"))
    }

    /// Plain HTTP to a relay would make the end-to-end encryption decorative, so it is refused
    /// before any traffic is sent.
    func testAPlainHTTPRelayIsRefused() async throws {
        let peer = AgentPeer()
        let identity = makeIdentity()
        RelayStub.handler = { _, _ in
            XCTFail("a plain-HTTP relay must be refused before any request is made")
            return (500, Data())
        }
        let transport = makeTransport(identity: identity)
        var endpoint = peer.endpoint(relayBaseURL: URL(string: "http://relay.example")!)
        endpoint.relayBaseURL = URL(string: "http://relay.example")!
        await assertThrows(.invalidRequest) {
            try await transport.send(.health, instanceID: UUID(), endpoint: .relay(endpoint))
        }
    }

    // MARK: - Freshness and status vocabulary

    func testProbeStatusMappingKeepsTheRelayCasesApart() {
        XCTAssertEqual(ConnectionProbe.status(for: .agentOffline, error: UpstreamError.agentOffline), .offline)
        XCTAssertEqual(ConnectionProbe.status(for: .relayUnreachable, error: UpstreamError.relayUnreachable), .offline)
        XCTAssertEqual(ConnectionProbe.status(for: .relayRejectedDevice, error: UpstreamError.relayRejectedDevice), .offline)
        XCTAssertEqual(ConnectionProbe.status(for: .nodeAuthFailed, error: UpstreamError.nodeAuthFailed), .degraded)
        XCTAssertEqual(ConnectionProbe.status(for: .confirmationRejected, error: UpstreamError.confirmationRejected), .degraded)
        XCTAssertEqual(ConnectionProbe.status(for: .unauthorized, error: UpstreamError.unauthorized), .authRequired)
        // Every relay-mode case must have its own message. A generic fallback is the failure this
        // whole vocabulary exists to prevent, so the check is that the text is specific rather than
        // merely present.
        let generic = "请求失败："
        var seen = Set<String>()
        for error in Self.allRelayErrors {
            let text = ErrorMessage.describe(error, nodeName: "Lab Alpha")
            XCTAssertFalse(text.isEmpty)
            XCTAssertFalse(text.hasPrefix(generic), "\(error) fell back to the generic message")
            XCTAssertTrue(seen.insert(text).inserted, "\(error) shares a message with another case")
        }
        XCTAssertEqual(seen.count, Self.allRelayErrors.count)
    }

    static let allRelayErrors: [UpstreamError] = [
        .relayUnreachable, .relayTimeout, .relayRejectedDevice, .relayNodeUnknown, .agentOffline,
        .nodeBusy, .agentTimeout, .relayBadEnvelope, .relayRateLimited, .relayPolicyDenied,
        .confirmationRejected, .envelopeExpired, .nodeRateLimited, .nodeAuthFailed, .payloadTooLarge,
        .relayResponseUnverifiable, .relayResponseMismatch, .envelopeRejected,
    ]

    // MARK: - Helpers

    private func assertThrows(
        _ expected: UpstreamError,
        _ message: String = "",
        _ body: () async throws -> UpstreamResponse,
        file: StaticString = #filePath,
        line: UInt = #line
    ) async {
        do {
            _ = try await body()
            XCTFail("expected \(expected) but the call succeeded. \(message)", file: file, line: line)
        } catch let error as UpstreamError {
            XCTAssertEqual(error, expected, message, file: file, line: line)
        } catch {
            XCTFail("expected \(expected), got \(error). \(message)", file: file, line: line)
        }
    }
}