// SPDX-License-Identifier: MIT
// Copyright (c) 2026 Thibault Ducray
//
// This file is part of MyPwdTool's open-source sync/encryption core — see
// LICENSE-SYNC-CRYPTO.md at the repo root and https://tducray.fr/mypwdtool/open-source/.
// The rest of this application is proprietary and NOT covered by this license.

import Foundation
import CryptoKit

// MARK: - Base64url helpers

func base64urlEncode(_ data: Data) -> String {
    data.base64EncodedString()
        .replacingOccurrences(of: "+", with: "-")
        .replacingOccurrences(of: "/", with: "_")
        .replacingOccurrences(of: "=", with: "")
}

func base64urlDecode(_ s: String) -> Data? {
    var str = s
        .replacingOccurrences(of: "-", with: "+")
        .replacingOccurrences(of: "_", with: "/")
    // Re-add padding
    let remainder = str.count % 4
    if remainder > 0 {
        str += String(repeating: "=", count: 4 - remainder)
    }
    return Data(base64Encoded: str)
}

// MARK: - Relay compile-time constants

private enum RelayConstants {
    /// Shared secret used for the registration challenge HMAC (§3.3 of the API spec).
    /// Must match the server's RELAY_MASTER_KEY environment variable.
    /// Base64url without padding — same encoding as the server.
    /// The actual value lives in AppSecrets/RelayMasterKey.swift, deliberately outside this
    /// directory — see that file's header comment for why.
    static var masterKeyB64URL: String { RelayMasterKey.base64URL }

    /// Sent verbatim in the HMAC message and stored by the server for diagnostics.
    static let clientVersion: String = {
        Bundle.main.infoDictionary?["CFBundleShortVersionString"] as? String ?? "1.0.0"
    }()

    #if os(iOS)
    static let platform = "iOS"
    #else
    static let platform = "MacOS"
    #endif
}

// MARK: - RelayError

enum RelayError: LocalizedError {
    case rateLimited(retryAfterSeconds: Int)
    case unauthorized
    case tokenExpired
    case challengeFailed
    case notFound
    case serverError(Int, String)
    case invalidResponse
    case networkError(Error)

    var errorDescription: String? {
        switch self {
        case .rateLimited(let s):       return "Rate limited — retry after \(s)s"
        case .unauthorized:             return "Unauthorized — invalid token"
        case .tokenExpired:             return "Token expired — please re-register"
        case .challengeFailed:          return "Registration challenge failed"
        case .notFound:                 return "Resource not found"
        case .serverError(let code, let msg): return "Server error \(code): \(msg)"
        case .invalidResponse:          return "Invalid server response"
        case .networkError(let e):      return "Network error: \(e.localizedDescription)"
        }
    }
}

// MARK: - Supporting Codable types

/// Response from GET /registration/challenge
struct RegistrationChallenge: Codable {
    var nonce: String
    var version: String
    var serverDate: String
    var expiresAt: String

    enum CodingKeys: String, CodingKey {
        case nonce
        case version
        case serverDate = "server_date"
        case expiresAt  = "expires_at"
    }
}

struct RegisterResponse: Codable {
    var deviceId: String
    var inboxId: String
    var sendToken: String
    var recvToken: String
    var limits: ServerLimits

    enum CodingKeys: String, CodingKey {
        case deviceId   = "device_id"
        case inboxId    = "inbox_id"
        case sendToken  = "send_token"
        case recvToken  = "recv_token"
        case limits
    }
}

struct OutboundRelayMessage {
    var messageId: String           // UUID, for idempotency
    var streamId: String
    var senderDeviceId: String
    var recipientInboxIds: [String]
    var ciphertext: String          // base64url
    var encryptedIdentifier: String? // optional base64url
    var ttlSeconds: Int?
    var cipherVersion: Int          // always 1
}

struct SendResponse: Codable {
    var accepted: Bool
    var messageId: String

    enum CodingKeys: String, CodingKey {
        case accepted
        case messageId = "message_id"
    }
}

struct PollResponse: Codable {
    var inboxId: String
    var nextCursor: String
    var messages: [InboundRelayMessage]

    enum CodingKeys: String, CodingKey {
        case inboxId    = "inbox_id"
        case nextCursor = "next_cursor"
        case messages
    }
}

struct InboundRelayMessage: Codable {
    var messageId: String
    var streamId: String
    var senderDeviceId: String
    var recipientInboxIds: [String]
    var createdAt: String
    var cipherVersion: Int
    var encryptedIdentifier: String?
    var ciphertext: String

    enum CodingKeys: String, CodingKey {
        case messageId           = "message_id"
        case streamId            = "stream_id"
        case senderDeviceId      = "sender_device_id"
        case recipientInboxIds   = "recipient_inbox_ids"
        case createdAt           = "created_at"
        case cipherVersion       = "cipher_version"
        case encryptedIdentifier = "encrypted_identifier"
        case ciphertext
    }
}

// MARK: - RelayAPIClient

final class RelayAPIClient {
    let baseURL: URL
    private let session: URLSession

    init(baseURL: URL) {
        self.baseURL = baseURL
        let config = URLSessionConfiguration.default
        config.httpAdditionalHeaders = ["Accept": "application/json"]
        self.session = URLSession(configuration: config)
    }

    // MARK: - GET /registration/challenge

    func fetchChallenge() async throws -> RegistrationChallenge {
        let url = baseURL.appendingPathComponent("registration/challenge")
        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        let data = try await performRequest(request)
        return try JSONDecoder().decode(RegistrationChallenge.self, from: data)
    }

    // MARK: - POST /devices

    /// Registers a new device.  Fetches a challenge nonce first (§4.1), computes the
    /// two-level HMAC response (§3.3), then calls POST /devices (§4.2).
    func registerDevice(name: String) async throws -> RegisterResponse {
        let platform      = RelayConstants.platform
        let clientVersion = RelayConstants.clientVersion

        // Step 1 – get nonce from server
        let challenge = try await fetchChallenge()

        // Step 2 – compute challenge response
        let challengeResponse = try computeChallengeResponse(
            challenge: challenge,
            platform: platform,
            clientVersion: clientVersion
        )

        // Step 3 – register
        let url = baseURL.appendingPathComponent("devices")
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")

        let body: [String: String] = [
            "device_name":        name,
            "platform":           platform,
            "client_version":     clientVersion,
            "challenge_nonce":    challenge.nonce,
            "challenge_response": challengeResponse
        ]
        request.httpBody = try JSONEncoder().encode(body)

        let data = try await performRequest(request)
        return try JSONDecoder().decode(RegisterResponse.self, from: data)
    }

    // MARK: - Challenge HMAC (§3.3)

    /// daily_key = HMAC-SHA256(master_key,  version + ":" + server_date)
    /// response  = HMAC-SHA256(daily_key,   nonce + ":" + version + ":" + platform + ":" + client_version)
    /// Returns the response as base64url without padding.
    private func computeChallengeResponse(
        challenge: RegistrationChallenge,
        platform: String,
        clientVersion: String
    ) throws -> String {
        guard let masterKeyData = base64urlDecode(RelayConstants.masterKeyB64URL),
              !masterKeyData.isEmpty else {
            throw RelayError.challengeFailed
        }

        let masterKey = SymmetricKey(data: masterKeyData)

        // daily_key = HMAC-SHA256(master_key, "V1:2026-03-31")
        let dailyMessage = Data("\(challenge.version):\(challenge.serverDate)".utf8)
        let dailyKeyBytes = HMAC<SHA256>.authenticationCode(for: dailyMessage, using: masterKey)
        let dailyKey = SymmetricKey(data: dailyKeyBytes)

        // response = HMAC-SHA256(daily_key, "nonce:V1:iOS:3.0.0")
        let responseMessage = Data("\(challenge.nonce):\(challenge.version):\(platform):\(clientVersion)".utf8)
        let responseBytes = HMAC<SHA256>.authenticationCode(for: responseMessage, using: dailyKey)

        return base64urlEncode(Data(responseBytes))
    }

    // MARK: - DELETE /devices

    func deregisterDevice(sendToken: String) async throws {
        let url = baseURL.appendingPathComponent("devices")
        var request = URLRequest(url: url)
        request.httpMethod = "DELETE"
        request.setValue("Bearer \(sendToken)", forHTTPHeaderField: "Authorization")
        _ = try await performRequest(request)
    }

    // MARK: - POST /messages

    func sendMessage(_ msg: OutboundRelayMessage, sendToken: String) async throws -> SendResponse {
        let url = baseURL.appendingPathComponent("messages")
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.setValue("Bearer \(sendToken)", forHTTPHeaderField: "Authorization")

        var body: [String: Any] = [
            "message_id": msg.messageId,
            "stream_id": msg.streamId,
            "sender_device_id": msg.senderDeviceId,
            "recipient_inbox_ids": msg.recipientInboxIds,
            "ciphertext": msg.ciphertext,
            "cipher_version": msg.cipherVersion
        ]
        if let ei = msg.encryptedIdentifier { body["encrypted_identifier"] = ei }
        if let ttl = msg.ttlSeconds { body["ttl_seconds"] = ttl }

        request.httpBody = try JSONSerialization.data(withJSONObject: body)

        let data = try await performRequest(request)
        return try JSONDecoder().decode(SendResponse.self, from: data)
    }

    // MARK: - GET /messages

    func pollMessages(recvToken: String, inboxId: String, after: String?, limit: Int, waitMs: Int) async throws -> PollResponse {
        var components = URLComponents(url: baseURL.appendingPathComponent("messages"), resolvingAgainstBaseURL: false)
        var queryItems = [
            URLQueryItem(name: "inbox_id", value: inboxId),
            URLQueryItem(name: "limit", value: "\(limit)"),
            URLQueryItem(name: "wait_ms", value: "\(waitMs)")
        ]
        if let cursor = after { queryItems.append(URLQueryItem(name: "after", value: cursor)) }
        components?.queryItems = queryItems

        guard let url = components?.url else { throw RelayError.invalidResponse }

        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        request.setValue("Bearer \(recvToken)", forHTTPHeaderField: "Authorization")
        // Long-poll timeout: waitMs/1000 + 15 seconds
        request.timeoutInterval = Double(waitMs) / 1000.0 + 15.0

        let data = try await performRequest(request)
        return try JSONDecoder().decode(PollResponse.self, from: data)
    }

    // MARK: - POST /messages/ack

    func ackMessages(recvToken: String, inboxId: String, messageIds: [String]) async throws {
        let url = baseURL.appendingPathComponent("messages/ack")
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.setValue("Bearer \(recvToken)", forHTTPHeaderField: "Authorization")

        let body: [String: Any] = [
            "inbox_id": inboxId,
            "message_ids": messageIds
        ]
        request.httpBody = try JSONSerialization.data(withJSONObject: body)
        _ = try await performRequest(request)
    }

    // MARK: - GET /limits

    func fetchLimits() async throws -> ServerLimits {
        let url = baseURL.appendingPathComponent("limits")
        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        let data = try await performRequest(request)
        return try JSONDecoder().decode(ServerLimits.self, from: data)
    }

    // MARK: - GET /health

    func healthCheck() async throws -> Bool {
        let url = baseURL.appendingPathComponent("health")
        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        request.timeoutInterval = 10
        do {
            _ = try await performRequest(request)
            return true
        } catch RelayError.serverError {
            return false
        }
    }

    // MARK: - Internal request performer

    private func performRequest(_ request: URLRequest) async throws -> Data {
        let (data, response): (Data, URLResponse)
        do {
            (data, response) = try await session.data(for: request)
        } catch {
            throw RelayError.networkError(error)
        }

        guard let httpResponse = response as? HTTPURLResponse else {
            throw RelayError.invalidResponse
        }

        switch httpResponse.statusCode {
        case 200...299:
            return data
        case 401:
            // Distinguish token_expired and challenge_failed from generic unauthorized
            if let errorCode = (try? JSONDecoder().decode(RelayAPIError.self, from: data))?.error.code {
                switch errorCode {
                case "token_expired":    throw RelayError.tokenExpired
                case "challenge_failed", "challenge_required": throw RelayError.challengeFailed
                default: break
                }
            }
            throw RelayError.unauthorized
        case 404:
            throw RelayError.notFound
        case 429:
            let retryAfter: Int
            if let retryHeader = httpResponse.value(forHTTPHeaderField: "Retry-After"),
               let seconds = Int(retryHeader) {
                retryAfter = seconds
            } else {
                retryAfter = 60
            }
            throw RelayError.rateLimited(retryAfterSeconds: retryAfter)
        default:
            let msg = String(data: data, encoding: .utf8) ?? "Unknown error"
            throw RelayError.serverError(httpResponse.statusCode, msg)
        }
    }
}

// MARK: - Server error envelope (for error code discrimination)

private struct RelayAPIError: Decodable {
    struct ErrorBody: Decodable { let code: String }
    let error: ErrorBody
}