import CryptoKit
import Foundation

/// What the app needs in order to reach a node, in either mode.
///
/// The abstraction exists so services stop taking a `baseURL` and a `cookie`: in relay mode there is
/// no base URL to dial and no cookie, and a service that kept asking for one would be asking a
/// question with no answer.
///
/// `instanceID` is part of the signature even though the design sketch showed only the endpoint,
/// because the direct path scopes its cookie by `(instanceID, origin)` and losing that would let two
/// instances pointed at the same node share a session. It is the one piece of the sketch that could
/// not survive contact with the existing session store.
public protocol NodeTransport: Sendable {
    func send(
        _ route: UpstreamRoute,
        instanceID: UUID,
        endpoint: InstanceEndpoint
    ) async throws -> UpstreamResponse
}

/// Picks the transport for an endpoint.
///
/// Services depend on this rather than on either concrete transport, so adding a mode later is a
/// change in one place instead of in every feature that talks to a node.
public struct NodeTransportRouter: NodeTransport {
    private let direct: any NodeTransport
    private let relay: any NodeTransport

    public init(direct: any NodeTransport, relay: any NodeTransport) {
        self.direct = direct
        self.relay = relay
    }

    public func send(
        _ route: UpstreamRoute,
        instanceID: UUID,
        endpoint: InstanceEndpoint
    ) async throws -> UpstreamResponse {
        switch endpoint {
        case .direct:
            return try await direct.send(route, instanceID: instanceID, endpoint: endpoint)
        case .relay:
            return try await relay.send(route, instanceID: instanceID, endpoint: endpoint)
        }
    }
}

/// The existing direct path, unchanged.
///
/// It is a thin adapter rather than a rewrite: the code path that reaches a node on the local network
/// today is the one users already depend on, and a refactor that changed its behaviour would be a
/// regression disguised as tidying.
public struct DirectTransport: NodeTransport {
    private let client: any SimAdminClientProtocol
    private let sessions: InstanceSessionStore
    private let credentials: (any CredentialStore)?

    public init(
        client: any SimAdminClientProtocol,
        sessions: InstanceSessionStore,
        credentials: (any CredentialStore)? = nil
    ) {
        self.client = client
        self.sessions = sessions
        self.credentials = credentials
    }

    public func send(
        _ route: UpstreamRoute,
        instanceID: UUID,
        endpoint: InstanceEndpoint
    ) async throws -> UpstreamResponse {
        guard let origin = endpoint.directOrigin else {
            // Reaching here means a relay instance was handed to the direct path, which is a wiring
            // bug rather than a network condition: fail loudly instead of dialing nothing.
            throw UpstreamError.invalidRequest
        }
        var cookie = await sessions.cookie(for: instanceID, origin: origin)
        do {
            return try await client.send(route, baseURL: origin, cookie: cookie)
        } catch UpstreamError.unauthorized, UpstreamError.forbidden {
            // A node whose login endpoints are unreachable (an old build, or one with protection
            // off) must not turn every read into a failed login attempt; the retry below is what
            // makes the second call happen only when credentials actually exist.
            guard let credentials, let password = try credentials.get(instanceID: instanceID),
                  !password.isEmpty
            else {
                throw UpstreamError.unauthorized
            }
            let service = InstanceSessionService(store: sessions, client: client)
            let instance = Instance(
                id: instanceID,
                name: "",
                endpoint: endpoint,
                createdAt: Date(),
                updatedAt: Date()
            )
            let result = try await service.login(instance, password: password)
            guard result.authenticated else { throw UpstreamError.unauthorized }
            cookie = await sessions.cookie(for: instanceID, origin: origin)
            return try await client.send(route, baseURL: origin, cookie: cookie)
        }
    }
}

/// The relay path: build an envelope, sign it, encrypt it, post it, then verify and decrypt the
/// answer.
///
/// It never touches a cookie and never sees a node password. The node's agent holds that credential
/// and logs in to its own loopback SimAdmin, which is what makes "the phone has no node password"
/// true rather than merely stated.
public struct RelayTransport: NodeTransport {
    /// Resolved per send rather than held, so a transport built before pairing still uses the keys
    /// the node actually paired with. Holding a copy would mean a pairing done after the transport
    /// was constructed signs with keys the node has never seen.
    private let identity: @Sendable () throws -> DeviceIdentity
    private let session: URLSession
    private let maximumBytes: Int
    private let now: @Sendable () -> Date

    public init(
        identity: @escaping @Sendable () throws -> DeviceIdentity,
        session: URLSession = .shared,
        maximumBytes: Int = 262_144,
        now: @escaping @Sendable () -> Date = { Date() }
    ) {
        self.identity = identity
        self.session = session
        self.maximumBytes = maximumBytes
        self.now = now
    }

    public init(
        identity: DeviceIdentity,
        session: URLSession = .shared,
        maximumBytes: Int = 262_144,
        now: @escaping @Sendable () -> Date = { Date() }
    ) {
        self.init(identity: { identity }, session: session, maximumBytes: maximumBytes, now: now)
    }

    public func send(
        _ route: UpstreamRoute,
        instanceID: UUID,
        endpoint: InstanceEndpoint
    ) async throws -> UpstreamResponse {
        guard let relay = endpoint.relay else { throw UpstreamError.invalidRequest }
        guard relay.relayBaseURL.scheme == "https" else {
            // Plain HTTP to a relay would make the endpoint-to-endpoint encryption decorative: a
            // third party could substitute a different ciphertext, and nothing above would notice.
            throw UpstreamError.invalidRequest
        }
        guard let requestBody = try? Self.encode(route: route) else {
            throw UpstreamError.invalidRequest
        }
        _ = instanceID
        guard let identity = try? identity() else {
            throw UpstreamError.invalidRequest
        }

        // The ephemeral key pair is per request and its private half is kept only until the answer
        // is decrypted. The response is encrypted to this key, so the two directions never share a
        // secret and the phone's long-term X25519 key is not needed to read an answer.
        let ephemeral = Curve25519.KeyAgreement.PrivateKey()
        let requestID = UUID().uuidString.lowercased()
        let timestamp = UInt64(now().timeIntervalSince1970)

        let shared = try RelayWire.sharedSecret(
            privateKey: ephemeral,
            peerPublic: relay.agentKxKey
        )
        let (key, iv) = try RelayWire.deriveKey(
            shared: shared,
            nodeID: relay.nodeID,
            requestID: requestID,
            response: false
        )
        let sealed = try RelayWire.seal(requestBody, key: key, iv: iv)

        var envelope = RelayWire.Envelope(
            node: relay.nodeID,
            req: requestID,
            app: relay.deviceID,
            ts: timestamp,
            ek: ephemeral.publicKey.rawRepresentation,
            p: sealed
        )
        try RelayWire.signRequest(&envelope, privateKey: identity.signingKey)

        let answer = try await post(envelope, relay: relay)
        return try decrypt(answer, relay: relay, requestID: requestID, ephemeral: ephemeral)
    }

    // MARK: - Inner request

    /// Builds the inner payload. The path is the route's own path, query string included, so the
    /// agent never assembles one from parts and there is no traversal surface from the relay.
    ///
    /// A dangerous write carries an intent computed here, from the exact bytes about to be sent. The
    /// agent recomputes it and refuses a mismatch, so a relay cannot edit the recipient of an SMS in
    /// flight and keep the confirmation valid. The intent is produced at the moment the user's
    /// gesture is already behind us — this method is only called for a request the UI has decided to
    /// make — which is what gives it any meaning at all.
    static func encode(route: UpstreamRoute) throws -> Data {
        let body = route.body
        var inner = RelayWire.InnerRequest(m: route.method, p: route.path, body: body)
        if route.requiresConfirmation {
            inner.intent = try RelayWire.intentDigest(
                method: route.method,
                path: route.path,
                body: body
            )
        }
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.sortedKeys]
        return try encoder.encode(inner)
    }

    // MARK: - HTTP

    private func post(_ envelope: RelayWire.Envelope, relay: RelayEndpoint) async throws -> RelayWire.Envelope {
        guard let url = URL(string: relay.relayBaseURL.absoluteString + "/v1/request") else {
            throw UpstreamError.invalidRequest
        }
        var request = URLRequest(url: url, timeoutInterval: 45)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "content-type")
        request.setValue("application/json", forHTTPHeaderField: "accept")
        request.setValue("Bearer \(relay.deviceToken)", forHTTPHeaderField: "authorization")
        request.httpBody = try JSONEncoder().encode(envelope)

        let data: Data
        let response: URLResponse
        do {
            (data, response) = try await session.data(for: request)
        } catch let error as URLError where error.code == .timedOut {
            throw UpstreamError.relayTimeout
        } catch {
            throw UpstreamError.relayUnreachable
        }
        guard let http = response as? HTTPURLResponse, data.count <= maximumBytes else {
            throw UpstreamError.invalidResponse
        }
        if http.statusCode != 200 {
            throw Self.problem(from: data, status: http.statusCode)
        }
        guard let envelope = try? JSONDecoder().decode(RelayWire.Envelope.self, from: data) else {
            throw UpstreamError.invalidResponse
        }
        return envelope
    }

    /// Maps a relay ProblemDetails document to a typed error.
    ///
    /// The four classes the design calls for are kept apart on purpose: "the relay is unreachable",
    /// "the agent is offline", "the envelope was refused" and "the node returned an error" lead to
    /// four different user actions, and folding them into one "request failed" is what makes relay
    /// problems unactionable.
    static func problem(from data: Data, status: Int) -> UpstreamError {
        let code = (try? JSONSerialization.jsonObject(with: data) as? [String: Any])
            .flatMap { $0?["code"] as? String }
            .flatMap(RelayErrorCode.init(rawValue:))
        switch code {
        case .unauthorized:
            // The relay rejected this phone's bearer. Nothing the node can do about it, and retrying
            // with the same token will fail the same way.
            return .relayRejectedDevice
        case .nodeUnknown:
            return .relayNodeUnknown
        case .agentOffline:
            return .agentOffline
        case .mailboxFull:
            return .nodeBusy
        case .agentTimeout:
            return .agentTimeout
        case .badEnvelope:
            return .relayBadEnvelope
        case .payloadTooLarge:
            return .payloadTooLarge
        case .rateLimited:
            return .relayRateLimited
        case .sessionExpired, .badChallenge:
            return .agentOffline
        case .internalError, .none:
            return status >= 500 ? .relayUnreachable : .invalidResponse
        }
    }

    // MARK: - Response

    private func decrypt(
        _ envelope: RelayWire.Envelope,
        relay: RelayEndpoint,
        requestID: String,
        ephemeral: Curve25519.KeyAgreement.PrivateKey
    ) throws -> UpstreamResponse {
        guard envelope.req == requestID, envelope.node == relay.nodeID else {
            // The relay answered a different question. Since the signature covers both fields, a
            // mismatch here means the answer was not produced for this request at all.
            throw UpstreamError.relayResponseMismatch
        }
        let signingKey: Curve25519.Signing.PublicKey
        do {
            signingKey = try Curve25519.Signing.PublicKey(rawRepresentation: relay.agentSigningKey)
        } catch {
            throw UpstreamError.relayResponseUnverifiable
        }
        do {
            try RelayWire.verifyResponse(envelope, publicKey: signingKey)
        } catch {
            // An answer the node did not sign is not an answer. This is the check that makes the
            // relay's inability to forge responses a fact rather than a claim.
            throw UpstreamError.relayResponseUnverifiable
        }
        let shared = try RelayWire.sharedSecret(privateKey: ephemeral, peerPublic: envelope.ek)
        let (key, iv) = try RelayWire.deriveKey(
            shared: shared,
            nodeID: relay.nodeID,
            requestID: requestID,
            response: true
        )
        let plaintext: Data
        do {
            plaintext = try RelayWire.open(envelope.p, key: key, iv: iv)
        } catch {
            throw UpstreamError.relayResponseUnverifiable
        }
        guard let inner = try? JSONDecoder().decode(RelayWire.InnerResponse.self, from: plaintext) else {
            throw UpstreamError.invalidResponse
        }

        let outcome = inner.outcome.flatMap(AgentOutcome.init(rawValue:))
        switch outcome {
        case .ok, .none:
            return UpstreamResponse(status: inner.status, data: inner.body ?? Data())
        case .nodeAuthFailed:
            // Nothing the phone can do: the node's agent is misconfigured. Distinguished from an
            // upstream error because the remedy is on the node, not in the app.
            throw UpstreamError.nodeAuthFailed
        case .policyDenied, .unknownDevice:
            throw UpstreamError.relayPolicyDenied
        case .confirmationMissing, .confirmationMismatch:
            throw UpstreamError.confirmationRejected
        case .expired, .replay:
            // The envelope was stale or already used. Both are clock or retry problems, and both
            // resolve by re-issuing the request with a fresh id and timestamp.
            throw UpstreamError.envelopeExpired
        case .rateLimited:
            throw UpstreamError.nodeRateLimited
        case .payloadTooLarge:
            throw UpstreamError.payloadTooLarge
        case .badEnvelope, .badSignature:
            throw UpstreamError.envelopeRejected
        case .upstreamError:
            // The node answered, but its local SimAdmin did not. The status is still useful: it is
            // the only branch where a 4xx/5xx from the node reaches the caller.
            return UpstreamResponse(status: inner.status, data: inner.body ?? Data())
        }
    }
}

extension Data {
    /// Lowercase hex, which is the spelling the relay's bearer-token comparison expects.
    var hexString: String { map { String(format: "%02x", $0) }.joined() }
}