import Foundation
import CryptoKit

// A focused implementation of the Noise Protocol Framework (revision 34) for the two
// patterns the sync layer needs (SYNC_PROTOCOL §4): XXpsk0 for QR pairing and IK for
// reconnect, both over 25519 / ChaChaPoly / SHA256. Only the primitives those patterns
// touch are implemented; anything else throws. Crypto is CryptoKit (X25519, ChaChaPoly,
// SHA256, HMAC) — no third-party crypto.

enum NoiseError: Error, Equatable {
    case decryptFailed
    case handshakeNotFinished
    case unexpectedToken(String)
    case messageTooShort
    case nonceExhausted
}

/// AEAD keyed by a 32-byte symmetric key with a 64-bit incrementing nonce (Noise §5.1).
struct CipherState {
    private(set) var key: SymmetricKey?
    private(set) var nonce: UInt64 = 0

    static func empty() -> CipherState { CipherState(key: nil) }
    init(key: SymmetricKey?) { self.key = key }

    var hasKey: Bool { key != nil }

    private static func noiseNonce(_ n: UInt64) -> ChaChaPoly.Nonce {
        // 96-bit IV = 32 zero bits ++ 64-bit little-endian counter (Noise §12.3).
        var bytes = [UInt8](repeating: 0, count: 12)
        var le = n.littleEndian
        withUnsafeBytes(of: &le) { src in
            for i in 0..<8 { bytes[4 + i] = src[i] }
        }
        return try! ChaChaPoly.Nonce(data: Data(bytes))
    }

    mutating func encrypt(ad: Data, plaintext: Data) throws -> Data {
        guard let key else { return plaintext }
        guard nonce < .max else { throw NoiseError.nonceExhausted }
        let box = try ChaChaPoly.seal(
            plaintext, using: key, nonce: Self.noiseNonce(nonce), authenticating: ad)
        nonce += 1
        return box.ciphertext + box.tag
    }

    mutating func decrypt(ad: Data, ciphertext: Data) throws -> Data {
        guard let key else { return ciphertext }
        guard nonce < .max else { throw NoiseError.nonceExhausted }
        guard ciphertext.count >= 16 else { throw NoiseError.messageTooShort }
        let tag = ciphertext.suffix(16)
        let ct = ciphertext.prefix(ciphertext.count - 16)
        let box = try ChaChaPoly.SealedBox(
            nonce: Self.noiseNonce(nonce), ciphertext: ct, tag: tag)
        do {
            let pt = try ChaChaPoly.open(box, using: key, authenticating: ad)
            nonce += 1
            return pt
        } catch {
            throw NoiseError.decryptFailed
        }
    }
}

/// Tracks the chaining key `ck` and handshake hash `h` (Noise §5.2).
struct SymmetricState {
    private(set) var cipher = CipherState.empty()
    private var ck: Data
    private(set) var h: Data

    init(protocolName: String) {
        let name = Data(protocolName.utf8)
        if name.count <= 32 {
            var padded = name
            padded.append(Data(repeating: 0, count: 32 - name.count))
            h = padded
        } else {
            h = Data(SHA256.hash(data: name))
        }
        ck = h
    }

    mutating func mixHash(_ data: Data) {
        h = Data(SHA256.hash(data: h + data))
    }

    mutating func mixKey(_ ikm: Data) {
        let out = Self.hkdf(ck: ck, ikm: ikm, outputs: 2)
        ck = out[0]
        cipher = CipherState(key: SymmetricKey(data: out[1]))
    }

    /// Mix a pre-shared key into both `ck` and `h` (Noise §9, PSK).
    mutating func mixKeyAndHash(_ ikm: Data) {
        let out = Self.hkdf(ck: ck, ikm: ikm, outputs: 3)
        ck = out[0]
        mixHash(out[1])
        cipher = CipherState(key: SymmetricKey(data: out[2]))
    }

    mutating func encryptAndHash(_ plaintext: Data) throws -> Data {
        let ct = try cipher.encrypt(ad: h, plaintext: plaintext)
        mixHash(ct)
        return ct
    }

    mutating func decryptAndHash(_ ciphertext: Data) throws -> Data {
        let pt = try cipher.decrypt(ad: h, ciphertext: ciphertext)
        mixHash(ciphertext)
        return pt
    }

    /// Final transport keys (Noise §5.2 Split): (send, recv) for the initiator;
    /// the responder swaps them.
    func split() -> (CipherState, CipherState) {
        let out = Self.hkdf(ck: ck, ikm: Data(), outputs: 2)
        return (CipherState(key: SymmetricKey(data: out[0])),
                CipherState(key: SymmetricKey(data: out[1])))
    }

    /// Noise's HKDF (HMAC-SHA256 chained), distinct from RFC 5869 Expand-with-info.
    static func hkdf(ck: Data, ikm: Data, outputs: Int) -> [Data] {
        let tempKey = hmac(key: ck, data: ikm)
        let o1 = hmac(key: tempKey, data: Data([0x01]))
        let o2 = hmac(key: tempKey, data: o1 + Data([0x02]))
        if outputs == 2 { return [o1, o2] }
        let o3 = hmac(key: tempKey, data: o2 + Data([0x03]))
        return [o1, o2, o3]
    }

    private static func hmac(key: Data, data: Data) -> Data {
        let mac = HMAC<SHA256>.authenticationCode(for: data, using: SymmetricKey(data: key))
        return Data(mac)
    }
}

// MARK: - DH helpers

enum NoiseDH {
    static func generate() -> Curve25519.KeyAgreement.PrivateKey {
        Curve25519.KeyAgreement.PrivateKey()
    }

    static func agree(_ priv: Curve25519.KeyAgreement.PrivateKey, _ pub: Data) throws -> Data {
        let pk = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: pub)
        let secret = try priv.sharedSecretFromKeyAgreement(with: pk)
        return secret.withUnsafeBytes { Data($0) }
    }
}
