import Foundation
import Security

/// RSA-4096 operations — used for login challenge-response ONLY
enum RSACrypto {

    // MARK: - Key Generation

    /// Generate RSA-4096 keypair
    static func generateKeypair() throws -> (privateKey: SecKey, publicKey: SecKey) {
        let attributes: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrKeySizeInBits as String: 4096,
        ]

        var error: Unmanaged<CFError>?
        guard let privateKey = SecKeyCreateRandomKey(attributes as CFDictionary, &error) else {
            throw CryptoError.rsaKeyGenerationFailed
        }
        guard let publicKey = SecKeyCopyPublicKey(privateKey) else {
            throw CryptoError.rsaKeyGenerationFailed
        }
        return (privateKey, publicKey)
    }

    // MARK: - Serialization

    /// Serialize RSA private key. With password: DER → ECP1. Without: PEM PKCS#8.
    static func serializePrivateKey(_ key: SecKey, password: Data? = nil) throws -> Data {
        var error: Unmanaged<CFError>?
        guard let derData = SecKeyCopyExternalRepresentation(key, &error) as Data? else {
            throw CryptoError.rsaOperationFailed("Failed to export private key")
        }

        // SecKey exports in PKCS#1 format on iOS — wrap in PKCS#8 for Python compat
        let pkcs8 = wrapRSAPrivateKeyPKCS8(derData)

        if let password = password {
            return try KeyEncryption.encrypt(pkcs8, password: password)
        }

        // PEM encode for Python compatibility
        return pemEncode(pkcs8, label: "PRIVATE KEY")
    }

    /// Serialize RSA public key as PEM SubjectPublicKeyInfo (Python-compatible)
    static func serializePublicKey(_ key: SecKey) throws -> Data {
        var error: Unmanaged<CFError>?
        guard let derData = SecKeyCopyExternalRepresentation(key, &error) as Data? else {
            throw CryptoError.rsaOperationFailed("Failed to export public key")
        }

        // SecKey exports PKCS#1 on iOS — wrap in SubjectPublicKeyInfo
        let spki = wrapRSAPublicKeySPKI(derData)
        return pemEncode(spki, label: "PUBLIC KEY")
    }

    /// Load RSA private key. Auto-detects ECP1 vs PEM format.
    static func loadPrivateKey(_ data: Data, password: Data? = nil) throws -> SecKey {
        let derData: Data

        if KeyEncryption.isECP1Format(data) {
            guard let pwd = password else {
                throw CryptoError.invalidKeyData("ECP1 key requires password")
            }
            let raw = try KeyEncryption.decrypt(data, password: pwd)
            derData = unwrapPKCS8ToRSAPrivateKey(raw)
        } else {
            // PEM format
            let pem = String(data: data, encoding: .utf8) ?? ""
            derData = try pemDecode(pem, label: "PRIVATE KEY")
                .flatMap { unwrapPKCS8ToRSAPrivateKey($0) }
                ?? pemDecode(pem, label: "RSA PRIVATE KEY")
                ?? { throw CryptoError.invalidKeyData("Cannot parse RSA private key PEM") }()
        }

        let attributes: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrKeyClass as String: kSecAttrKeyClassPrivate,
        ]

        var error: Unmanaged<CFError>?
        guard let key = SecKeyCreateWithData(derData as CFData, attributes as CFDictionary, &error) else {
            throw CryptoError.invalidKeyData("Failed to create RSA private key from DER")
        }
        return key
    }

    /// Load RSA public key from PEM
    static func loadPublicKey(_ pemData: Data) throws -> SecKey {
        let pem = String(data: pemData, encoding: .utf8) ?? ""

        // Try SubjectPublicKeyInfo (PUBLIC KEY), unwrap to PKCS#1
        let derData: Data
        if let spki = pemDecode(pem, label: "PUBLIC KEY") {
            derData = unwrapSPKIToRSAPublicKey(spki)
        } else if let pkcs1 = pemDecode(pem, label: "RSA PUBLIC KEY") {
            derData = pkcs1
        } else {
            throw CryptoError.invalidKeyData("Cannot parse RSA public key PEM")
        }

        let attributes: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
        ]

        var error: Unmanaged<CFError>?
        guard let key = SecKeyCreateWithData(derData as CFData, attributes as CFDictionary, &error) else {
            throw CryptoError.invalidKeyData("Failed to create RSA public key from DER")
        }
        return key
    }

    // MARK: - Sign / Verify

    /// Sign data with RSA-PSS SHA-256.
    /// Note: iOS uses salt_length = hash_length (32). Server must use PSS.AUTO to verify.
    static func sign(_ privateKey: SecKey, data: Data) throws -> Data {
        var error: Unmanaged<CFError>?
        guard let signature = SecKeyCreateSignature(
            privateKey,
            .rsaSignatureMessagePSSSHA256,
            data as CFData,
            &error
        ) as Data? else {
            throw CryptoError.rsaOperationFailed("RSA signing failed")
        }
        return signature
    }

    /// Verify RSA-PSS SHA-256 signature
    static func verify(_ publicKey: SecKey, signature: Data, data: Data) -> Bool {
        SecKeyVerifySignature(
            publicKey,
            .rsaSignatureMessagePSSSHA256,
            data as CFData,
            signature as CFData,
            nil
        )
    }

    // MARK: - PEM Helpers

    private static func pemEncode(_ der: Data, label: String) -> Data {
        let base64 = der.base64EncodedString(options: .lineLength64Characters)
        let pem = "-----BEGIN \(label)-----\n\(base64)\n-----END \(label)-----\n"
        return Data(pem.utf8)
    }

    private static func pemDecode(_ pem: String, label: String) -> Data? {
        let beginMarker = "-----BEGIN \(label)-----"
        let endMarker = "-----END \(label)-----"

        guard let beginRange = pem.range(of: beginMarker),
              let endRange = pem.range(of: endMarker) else {
            return nil
        }

        let base64String = pem[beginRange.upperBound..<endRange.lowerBound]
            .replacingOccurrences(of: "\n", with: "")
            .replacingOccurrences(of: "\r", with: "")
            .replacingOccurrences(of: " ", with: "")

        return Data(base64Encoded: base64String)
    }

    // MARK: - ASN.1 PKCS#8 / SPKI Wrappers

    // SecKey on iOS exports RSA keys in PKCS#1 format, but Python expects PKCS#8 / SPKI.
    // These functions add/remove the ASN.1 wrapping.

    // RSA OID: 1.2.840.113549.1.1.1
    private static let rsaOID: [UInt8] = [0x06, 0x09, 0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x01]
    private static let nullParam: [UInt8] = [0x05, 0x00]

    /// Wrap PKCS#1 RSA private key in PKCS#8 PrivateKeyInfo envelope
    private static func wrapRSAPrivateKeyPKCS8(_ pkcs1: Data) -> Data {
        // PrivateKeyInfo ::= SEQUENCE {
        //   version INTEGER (0),
        //   algorithm AlgorithmIdentifier,
        //   privateKey OCTET STRING (containing PKCS#1 key)
        // }
        let version = Data([0x02, 0x01, 0x00]) // INTEGER 0
        let algorithmSeq = asn1Sequence(Data(rsaOID) + Data(nullParam))
        let privateKeyOctet = asn1OctetString(pkcs1)
        return asn1Sequence(version + algorithmSeq + privateKeyOctet)
    }

    /// Unwrap PKCS#8 to get PKCS#1 RSA private key
    private static func unwrapPKCS8ToRSAPrivateKey(_ pkcs8: Data) -> Data {
        // Parse SEQUENCE, skip version + algorithm, extract OCTET STRING
        guard pkcs8.count > 2 else { return pkcs8 }

        var offset = 0
        // Outer SEQUENCE
        guard pkcs8[offset] == 0x30 else { return pkcs8 }
        offset += 1
        offset = skipASN1Length(pkcs8, offset: offset)

        // Version INTEGER
        guard offset < pkcs8.count, pkcs8[offset] == 0x02 else { return pkcs8 }
        offset += 1
        let versionLen = readASN1Length(pkcs8, offset: &offset)
        offset += versionLen

        // Algorithm SEQUENCE
        guard offset < pkcs8.count, pkcs8[offset] == 0x30 else { return pkcs8 }
        offset += 1
        let algoLen = readASN1Length(pkcs8, offset: &offset)
        offset += algoLen

        // Private key OCTET STRING
        guard offset < pkcs8.count, pkcs8[offset] == 0x04 else { return pkcs8 }
        offset += 1
        let keyLen = readASN1Length(pkcs8, offset: &offset)
        guard offset + keyLen <= pkcs8.count else { return pkcs8 }
        return Data(pkcs8[offset..<(offset + keyLen)])
    }

    /// Wrap PKCS#1 RSA public key in SubjectPublicKeyInfo
    private static func wrapRSAPublicKeySPKI(_ pkcs1: Data) -> Data {
        // SubjectPublicKeyInfo ::= SEQUENCE {
        //   algorithm AlgorithmIdentifier,
        //   subjectPublicKey BIT STRING (containing PKCS#1 key)
        // }
        let algorithmSeq = asn1Sequence(Data(rsaOID) + Data(nullParam))
        let bitString = asn1BitString(pkcs1)
        return asn1Sequence(algorithmSeq + bitString)
    }

    /// Unwrap SubjectPublicKeyInfo to get PKCS#1 RSA public key
    private static func unwrapSPKIToRSAPublicKey(_ spki: Data) -> Data {
        guard spki.count > 2 else { return spki }

        var offset = 0
        // Outer SEQUENCE
        guard spki[offset] == 0x30 else { return spki }
        offset += 1
        offset = skipASN1Length(spki, offset: offset)

        // Algorithm SEQUENCE
        guard offset < spki.count, spki[offset] == 0x30 else { return spki }
        offset += 1
        let algoLen = readASN1Length(spki, offset: &offset)
        offset += algoLen

        // BIT STRING
        guard offset < spki.count, spki[offset] == 0x03 else { return spki }
        offset += 1
        let bitLen = readASN1Length(spki, offset: &offset)
        // Skip the unused bits byte
        guard offset < spki.count, spki[offset] == 0x00 else { return spki }
        offset += 1
        let keyLen = bitLen - 1
        guard offset + keyLen <= spki.count else { return spki }
        return Data(spki[offset..<(offset + keyLen)])
    }

    // MARK: - ASN.1 Primitives

    private static func asn1Length(_ length: Int) -> Data {
        if length < 0x80 {
            return Data([UInt8(length)])
        } else if length <= 0xFF {
            return Data([0x81, UInt8(length)])
        } else if length <= 0xFFFF {
            return Data([0x82, UInt8(length >> 8), UInt8(length & 0xFF)])
        } else {
            return Data([0x83, UInt8(length >> 16), UInt8((length >> 8) & 0xFF), UInt8(length & 0xFF)])
        }
    }

    private static func asn1Sequence(_ content: Data) -> Data {
        Data([0x30]) + asn1Length(content.count) + content
    }

    private static func asn1OctetString(_ content: Data) -> Data {
        Data([0x04]) + asn1Length(content.count) + content
    }

    private static func asn1BitString(_ content: Data) -> Data {
        // BIT STRING: tag + length + unused_bits(0) + content
        Data([0x03]) + asn1Length(content.count + 1) + Data([0x00]) + content
    }

    private static func readASN1Length(_ data: Data, offset: inout Int) -> Int {
        guard offset < data.count else { return 0 }
        let first = data[offset]
        offset += 1
        if first < 0x80 {
            return Int(first)
        }
        let numBytes = Int(first & 0x7F)
        var length = 0
        for _ in 0..<numBytes {
            guard offset < data.count else { return length }
            length = (length << 8) | Int(data[offset])
            offset += 1
        }
        return length
    }

    private static func skipASN1Length(_ data: Data, offset: Int) -> Int {
        var off = offset
        _ = readASN1Length(data, offset: &off)
        return off
    }
}