import Foundation
import Security
import CryptoKit

/// How to evaluate the server's TLS certificate.
///
/// The appliance must present a CA-trusted certificate — a real cert, or an
/// internal CA installed on the device via MDM/configuration profile. There is
/// deliberately NO "accept self-signed" escape hatch: the app never bypasses
/// trust evaluation (parity with the Android client).
enum ServerTrustPolicy: Equatable, Codable {
    /// Default OS trust evaluation. The default policy and the only one the UI
    /// produces.
    case system

    /// Certificate/public-key pinning: trust only if the leaf's SHA-256
    /// SubjectPublicKeyInfo hash (base64) matches one of these. Code-only (no UI);
    /// the strictest option for a fixed appliance — capture the hash once and pin it.
    case pinPublicKeySHA256([String], host: String)
}

/// `URLSessionDelegate` that applies a `ServerTrustPolicy`. Retained by the
/// `URLSession` it's attached to.
final class ServerTrustEvaluator: NSObject, URLSessionDelegate {
    private let policy: ServerTrustPolicy

    init(policy: ServerTrustPolicy) {
        self.policy = policy
    }

    func urlSession(
        _ session: URLSession,
        didReceive challenge: URLAuthenticationChallenge,
        completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Void
    ) {
        guard challenge.protectionSpace.authenticationMethod == NSURLAuthenticationMethodServerTrust,
              let trust = challenge.protectionSpace.serverTrust else {
            completionHandler(.performDefaultHandling, nil)
            return
        }

        let host = challenge.protectionSpace.host

        switch policy {
        case .system:
            completionHandler(.performDefaultHandling, nil)

        case .pinPublicKeySHA256(let pins, let pinnedHost):
            guard host == pinnedHost,
                  let leaf = Self.leafCertificate(from: trust),
                  let hash = Self.publicKeySHA256Base64(for: leaf),
                  pins.contains(hash) else {
                completionHandler(.cancelAuthenticationChallenge, nil)
                return
            }
            completionHandler(.useCredential, URLCredential(trust: trust))
        }
    }

    // MARK: - Pinning helpers

    private static func leafCertificate(from trust: SecTrust) -> SecCertificate? {
        if #available(iOS 15.0, *) {
            return (SecTrustCopyCertificateChain(trust) as? [SecCertificate])?.first
        } else {
            return SecTrustGetCertificateAtIndex(trust, 0)
        }
    }

    /// Base64 of SHA-256 over the DER-encoded **SubjectPublicKeyInfo**, i.e. the
    /// standard public-key pin. `SecKeyCopyExternalRepresentation` returns only the
    /// raw key (PKCS#1 `RSAPublicKey` for RSA, the `04||X||Y` point for EC) — NOT
    /// the SPKI — so we prepend the algorithm-specific ASN.1 SPKI header before
    /// hashing. The result matches the canonical capture:
    ///
    ///   openssl s_client -connect host:443 </dev/null 2>/dev/null \
    ///     | openssl x509 -pubkey -noout \
    ///     | openssl pkey -pubin -outform der \
    ///     | openssl dgst -sha256 -binary | base64
    ///
    /// Returns nil for key types/sizes without a known header (caller fails closed).
    static func publicKeySHA256Base64(for certificate: SecCertificate) -> String? {
        guard let key = SecCertificateCopyKey(certificate),
              let raw = SecKeyCopyExternalRepresentation(key, nil) as Data?,
              let attrs = SecKeyCopyAttributes(key) as? [String: Any],
              let keyType = attrs[kSecAttrKeyType as String] as? String,
              let keySize = attrs[kSecAttrKeySizeInBits as String] as? Int,
              let header = spkiHeader(keyType: keyType, keySizeInBits: keySize)
        else { return nil }
        var spki = Data(header)
        spki.append(raw)
        return CryptoSHA256.base64(of: spki)
    }

    /// ASN.1 SubjectPublicKeyInfo prefix for the raw key returned by
    /// `SecKeyCopyExternalRepresentation`, keyed by (algorithm, key size). These
    /// are the fixed, well-known DER headers; the raw key bits follow.
    private static func spkiHeader(keyType: String, keySizeInBits: Int) -> [UInt8]? {
        if keyType == (kSecAttrKeyTypeRSA as String) {
            switch keySizeInBits {
            case 2048: return rsa2048SPKIHeader
            case 3072: return rsa3072SPKIHeader
            case 4096: return rsa4096SPKIHeader
            default: return nil
            }
        }
        if keyType == (kSecAttrKeyTypeECSECPrimeRandom as String) {
            switch keySizeInBits {
            case 256: return ecP256SPKIHeader
            case 384: return ecP384SPKIHeader
            default: return nil
            }
        }
        return nil
    }

    private static let rsa2048SPKIHeader: [UInt8] = [
        0x30, 0x82, 0x01, 0x22, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
        0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x01, 0x0f, 0x00]
    private static let rsa3072SPKIHeader: [UInt8] = [
        0x30, 0x82, 0x01, 0xa2, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
        0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x01, 0x8f, 0x00]
    private static let rsa4096SPKIHeader: [UInt8] = [
        0x30, 0x82, 0x02, 0x22, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
        0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x02, 0x0f, 0x00]
    private static let ecP256SPKIHeader: [UInt8] = [
        0x30, 0x59, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02,
        0x01, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x03,
        0x42, 0x00]
    private static let ecP384SPKIHeader: [UInt8] = [
        0x30, 0x76, 0x30, 0x10, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02,
        0x01, 0x06, 0x05, 0x2b, 0x81, 0x04, 0x00, 0x22, 0x03, 0x62, 0x00]
}

/// Minimal SHA-256 wrapper using CryptoKit.
enum CryptoSHA256 {
    static func base64(of data: Data) -> String {
        let digest = SHA256.hash(data: data)
        return Data(digest).base64EncodedString()
    }
}