import Foundation
import Security

/// Minimal PEM/DER helpers to turn the crypton-produced PEM strings into
/// `SecKey`s. iOS's `SecKeyCreateWithData` only accepts *raw PKCS#1* DER for
/// RSA, so we must strip the two wrapper formats crypton/Nextcloud use:
///
///  - Private key: crypton's `toPEM()` emits the `RSA PRIVATE KEY` label but the
///    DER body is actually **PKCS#8** `PrivateKeyInfo`
///    (SEQUENCE { INTEGER version, SEQUENCE algId, OCTET STRING pkcs1 }).
///    We extract the OCTET STRING content = the PKCS#1 `RSAPrivateKey`.
///
///  - Public key: `PUBLIC KEY` label, **SPKI** `SubjectPublicKeyInfo`
///    (SEQUENCE { SEQUENCE algId, BIT STRING spki }). We extract the BIT STRING
///    content and drop its leading 0x00 "unused bits" byte = the PKCS#1
///    `RSAPublicKey`.
///
/// Only DER lengths up to 4 bytes are handled — more than enough for RSA-2048.
enum PEM {

    /// Strips the PEM armor and returns the base64-decoded DER.
    static func der(fromPem pem: String) -> Data? {
        let body = pem
            .split(whereSeparator: { $0 == "\n" || $0 == "\r" })
            .filter { !$0.hasPrefix("-----") }
            .joined()
        return Data(base64Encoded: body)
    }

    /// PKCS#8 `PrivateKeyInfo` -> inner PKCS#1 `RSAPrivateKey`.
    static func pkcs1PrivateKey(fromPkcs8 der: Data) -> Data? {
        var reader = DERReader(der)
        guard let outer = reader.readTLV(), outer.tag == 0x30 else { return nil }
        var inner = DERReader(Data(outer.value))
        guard let version = inner.readTLV(), version.tag == 0x02 else { return nil } // INTEGER version
        guard let algId = inner.readTLV(), algId.tag == 0x30 else { return nil }      // SEQUENCE algId
        guard let octet = inner.readTLV(), octet.tag == 0x04 else { return nil }      // OCTET STRING
        return Data(octet.value)
    }

    /// SPKI `SubjectPublicKeyInfo` -> inner PKCS#1 `RSAPublicKey`.
    static func pkcs1PublicKey(fromSpki der: Data) -> Data? {
        var reader = DERReader(der)
        guard let outer = reader.readTLV(), outer.tag == 0x30 else { return nil }
        var inner = DERReader(Data(outer.value))
        guard let algId = inner.readTLV(), algId.tag == 0x30 else { return nil }   // SEQUENCE algId
        guard let bitString = inner.readTLV(), bitString.tag == 0x03 else { return nil } // BIT STRING
        var bytes = Array(bitString.value)
        guard let first = bytes.first, first == 0x00 else { return nil }           // unused-bits count
        bytes.removeFirst()
        return Data(bytes)
    }

    /// Builds an RSA `SecKey` from raw PKCS#1 DER.
    static func rsaKey(pkcs1 der: Data, isPrivate: Bool) -> SecKey? {
        let attributes: [CFString: Any] = [
            kSecAttrKeyType: kSecAttrKeyTypeRSA,
            kSecAttrKeyClass: isPrivate ? kSecAttrKeyClassPrivate : kSecAttrKeyClassPublic,
            kSecAttrKeySizeInBits: 2048,
        ]
        var error: Unmanaged<CFError>?
        let key = SecKeyCreateWithData(der as CFData, attributes as CFDictionary, &error)
        if key == nil {
            NSLog("[NSE] SecKeyCreateWithData failed: \(String(describing: error?.takeRetainedValue()))")
        }
        return key
    }
}

/// Tiny non-recursive DER TLV reader over a byte buffer.
private struct DERReader {
    private let bytes: [UInt8]
    private var pos = 0

    init(_ data: Data) { bytes = [UInt8](data) }

    /// Reads one tag-length-value triple, advancing past the value. Returns the
    /// tag byte and the value bytes, or nil on a malformed/truncated buffer.
    mutating func readTLV() -> (tag: UInt8, value: ArraySlice<UInt8>)? {
        guard pos < bytes.count else { return nil }
        let tag = bytes[pos]
        pos += 1

        guard pos < bytes.count else { return nil }
        var length = Int(bytes[pos])
        pos += 1

        if length & 0x80 != 0 {
            let numLengthBytes = length & 0x7F
            guard numLengthBytes > 0, numLengthBytes <= 4, pos + numLengthBytes <= bytes.count else {
                return nil
            }
            length = 0
            for _ in 0..<numLengthBytes {
                length = (length << 8) | Int(bytes[pos])
                pos += 1
            }
        }

        guard pos + length <= bytes.count else { return nil }
        let value = bytes[pos..<pos + length]
        pos += length
        return (tag, value)
    }
}