import CommonCrypto
import CryptoKit
import Foundation
import Martin

/// Pure SCRAM-SHA-512 math per RFC 5802 (SHA-512 / HMAC-SHA-512), kept
/// separate from the Martin mechanism so the proof computation can be
/// unit-tested against reference vectors.
enum SCRAMSHA512 {
    struct ServerFirst {
        let nonce: String
        let salt: Data
        let iterations: Int
    }

    private static let nonceAlphabet = Array(
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    )

    static func makeNonce(length: Int = 24) -> String {
        String((0..<length).map { _ in nonceAlphabet.randomElement()! })
    }

    static func parseServerFirst(
        _ message: String,
        expectedNoncePrefix: String
    ) throws -> ServerFirst {
        let pattern = #"^(?:m=[^\000=]+,)?r=([\x21-\x2B\x2D-\x7E]+),s=([a-zA-Z0-9/+=]+),i=(\d+)(?:,.*)?$"#
        guard let regex = try? NSRegularExpression(pattern: pattern) else {
            throw ClientSaslException.badChallenge(msg: "Failed to parse challenge")
        }
        let fullRange = NSRange(message.startIndex..., in: message)
        guard let match = regex.firstMatch(in: message, range: fullRange),
              match.numberOfRanges >= 4,
              let nonceRange = Range(match.range(at: 1), in: message),
              let saltRange = Range(match.range(at: 2), in: message),
              let iterationsRange = Range(match.range(at: 3), in: message) else {
            throw ClientSaslException.badChallenge(msg: "Failed to parse challenge")
        }
        let nonce = String(message[nonceRange])
        let iterations = Int(message[iterationsRange]) ?? 0
        guard nonce.hasPrefix(expectedNoncePrefix),
              let salt = Data(base64Encoded: String(message[saltRange])),
              iterations > 0 else {
            throw ClientSaslException.badChallenge(msg: "Invalid challenge")
        }
        return ServerFirst(nonce: nonce, salt: salt, iterations: iterations)
    }

    /// PBKDF2-HMAC-SHA-512 (RFC 5802 \"Hi\" function).
    static func saltedPassword(password: String, salt: Data, iterations: Int) -> [UInt8] {
        let passwordBytes = Array(password.utf8)
        let saltBytes = [UInt8](salt)
        var output = [UInt8](repeating: 0, count: Int(CC_SHA512_DIGEST_LENGTH))
        CCKeyDerivationPBKDF(
            CCPBKDFAlgorithm(kCCPBKDF2),
            passwordBytes, passwordBytes.count,
            saltBytes, saltBytes.count,
            CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA512),
            UInt32(iterations),
            &output, output.count
        )
        return output
    }

    static func clientProof(saltedPassword: [UInt8], authMessage: String) -> [UInt8] {
        let clientKey = hmac(saltedPassword, Array("Client Key".utf8))
        let storedKey = digest(clientKey)
        let clientSignature = hmac(storedKey, Array(authMessage.utf8))
        return zip(clientKey, clientSignature).map(^)
    }

    static func serverSignature(saltedPassword: [UInt8], authMessage: String) -> [UInt8] {
        let serverKey = hmac(saltedPassword, Array("Server Key".utf8))
        return hmac(serverKey, Array(authMessage.utf8))
    }

    static func verifyServerSignature(
        saltedPassword: [UInt8],
        authMessage: String,
        finalMessage: String
    ) -> Bool {
        let pattern = #"^(?:e=([^,]+)|v=([a-zA-Z0-9/+=]+)(?:,.*)?)$"#
        guard let regex = try? NSRegularExpression(pattern: pattern),
              let match = regex.firstMatch(
                  in: finalMessage,
                  range: NSRange(finalMessage.startIndex..., in: finalMessage)
              ),
              match.numberOfRanges >= 3,
              let vRange = Range(match.range(at: 2), in: finalMessage),
              let value = Data(base64Encoded: String(finalMessage[vRange])) else {
            return false
        }
        return value == Data(serverSignature(saltedPassword: saltedPassword, authMessage: authMessage))
    }

    private static func hmac(_ key: [UInt8], _ data: [UInt8]) -> [UInt8] {
        let code = HMAC<SHA512>.authenticationCode(
            for: Data(data),
            using: SymmetricKey(data: Data(key))
        )
        return Array(code)
    }

    private static func digest(_ data: [UInt8]) -> [UInt8] {
        Array(SHA512.hash(data: Data(data)))
    }
}

/// SCRAM-SHA-512 SASL mechanism. Martin ships only SCRAM-SHA-1 and
/// SCRAM-SHA-256; modern servers prefer SHA-512, so Luma registers this
/// mechanism ahead of Martin's ones.
final class LumaScramSha512Mechanism: SaslMechanism {
    let name = "SCRAM-SHA-512"
    private(set) var status: SaslMechanismStatus = .new

    private var stage = 0
    private var clientNonce = ""
    private var clientFirstMessageBare = ""
    private var authMessage = ""
    private var saltedPassword: [UInt8] = []

    func reset(scopes: Set<ResetableScope>) {
        guard scopes.contains(.stream) else { return }
        status = .new
        stage = 0
        clientNonce = ""
        clientFirstMessageBare = ""
        authMessage = ""
        saltedPassword = []
    }

    func isAllowedToUse(_ context: Context) -> Bool {
        if case .password(_, _, _) = context.connectionConfiguration.credentials {
            return true
        }
        return false
    }

    func evaluateChallenge(_ input: String?, context: Context) throws -> String? {
        guard status != .completed else {
            guard input == nil else {
                throw ClientSaslException.genericError(msg: "Already authorized")
            }
            return nil
        }
        switch stage {
        case 0:
            guard case .password(_, _, _) = context.connectionConfiguration.credentials else {
                throw ClientSaslException.genericError(msg: "Invalid credentials type")
            }
            clientNonce = SCRAMSHA512.makeNonce()
            clientFirstMessageBare =
                "n=\(context.userBareJid.localPart ?? ""),r=\(clientNonce)"
            stage = 1
            status = .completedExpected
            let first = "n,," + clientFirstMessageBare
            return first.data(using: .utf8)?.base64EncodedString()

        case 1:
            guard case .password(let password, _, _) = context.connectionConfiguration.credentials,
                  let input,
                  let data = Data(base64Encoded: input),
                  let serverFirst = String(data: data, encoding: .utf8) else {
                throw ClientSaslException.badChallenge(msg: "Invalid challenge")
            }
            let parsed = try SCRAMSHA512.parseServerFirst(
                serverFirst,
                expectedNoncePrefix: clientNonce
            )
            let clientFinalWithoutProof = "c=biws,r=\(parsed.nonce)"
            authMessage = clientFirstMessageBare + "," + serverFirst + "," + clientFinalWithoutProof
            saltedPassword = SCRAMSHA512.saltedPassword(
                password: password,
                salt: parsed.salt,
                iterations: parsed.iterations
            )
            let proof = SCRAMSHA512.clientProof(
                saltedPassword: saltedPassword,
                authMessage: authMessage
            )
            stage = 2
            let final = clientFinalWithoutProof + ",p=" + Data(proof).base64EncodedString()
            return final.data(using: .utf8)?.base64EncodedString()

        case 2:
            guard let input,
                  let data = Data(base64Encoded: input),
                  let finalMessage = String(data: data, encoding: .utf8),
                  SCRAMSHA512.verifyServerSignature(
                      saltedPassword: saltedPassword,
                      authMessage: authMessage,
                      finalMessage: finalMessage
                  ) else {
                throw ClientSaslException.invalidServerSignature
            }
            status = .completed
            return nil

        default:
            throw ClientSaslException.genericError(msg: "Illegal state")
        }
    }
}