import Foundation
import CryptoKit

public final class NodeIdentity: @unchecked Sendable {
    public static let currentVersion = 1
    public let signingKey: Curve25519.Signing.PrivateKey
    public let agreementKey: Curve25519.KeyAgreement.PrivateKey
    public let record: NodeRecord

    public init(signingKey: Curve25519.Signing.PrivateKey = .init(),
                agreementKey: Curve25519.KeyAgreement.PrivateKey = .init()) throws {
        self.signingKey = signingKey; self.agreementKey = agreementKey
        let signing = signingKey.publicKey.rawRepresentation
        self.record = NodeRecord(address: try MeshAddress.derive(from: signing),
                                 signingPublicKey: signing,
                                 agreementPublicKey: agreementKey.publicKey.rawRepresentation)
    }

    public static func loadOrCreate(at url: URL) throws -> NodeIdentity {
        struct Stored: Codable { let version: Int; let signing: Data; let agreement: Data }
        let decoder = PropertyListDecoder()
        if let data = try? Data(contentsOf: url) {
            let value = try decoder.decode(Stored.self, from: data)
            guard value.version == currentVersion else { throw UMNError.invalidIdentity }
            return try NodeIdentity(signingKey: .init(rawRepresentation: value.signing),
                                    agreementKey: .init(rawRepresentation: value.agreement))
        }
        let identity = try NodeIdentity()
        try FileManager.default.createDirectory(at: url.deletingLastPathComponent(), withIntermediateDirectories: true)
        let encoder = PropertyListEncoder(); encoder.outputFormat = .binary
        let data = try encoder.encode(Stored(version: currentVersion,
                                             signing: identity.signingKey.rawRepresentation,
                                             agreement: identity.agreementKey.rawRepresentation))
        try data.write(to: url, options: .atomic)
        try FileManager.default.setAttributes([.posixPermissions: 0o600], ofItemAtPath: url.path)
        return identity
    }

    public func sign(_ data: Data) throws -> Data { try signingKey.signature(for: data) }

    public func makeLinkState(sequence: UInt64, neighbors: [MeshAddress]) throws -> LinkState {
        let unsigned = LinkState(origin: record, sequence: sequence, neighbors: neighbors, signature: Data())
        return LinkState(origin: record, sequence: sequence, neighbors: neighbors,
                         signature: try sign(unsigned.signingBytes()))
    }

    public func seal(_ inner: InnerFrame, to destination: NodeRecord) throws -> SealedPayload {
        guard destination.validate() else { throw UMNError.invalidIdentity }
        let encoder = PropertyListEncoder(); encoder.outputFormat = .binary
        let innerData = try encoder.encode(inner)
        struct Signed: Codable { let inner: Data; let signature: Data }
        let signed = try encoder.encode(Signed(inner: innerData, signature: sign(innerData)))
        let ephemeral = Curve25519.KeyAgreement.PrivateKey()
        let remote = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: destination.agreementPublicKey)
        let secret = try ephemeral.sharedSecretFromKeyAgreement(with: remote)
        let key = secret.hkdfDerivedSymmetricKey(using: SHA256.self, salt: Data("umn-e2e-v1".utf8),
                                                  sharedInfo: destination.address.bytes, outputByteCount: 32)
        let box = try ChaChaPoly.seal(signed, using: key)
        return SealedPayload(ephemeralPublicKey: ephemeral.publicKey.rawRepresentation,
                             combinedCiphertext: box.combined)
    }

    public func open(_ payload: SealedPayload, expectedSource: MeshAddress) throws -> InnerFrame {
        struct Signed: Codable { let inner: Data; let signature: Data }
        let ephemeral = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: payload.ephemeralPublicKey)
        let secret = try agreementKey.sharedSecretFromKeyAgreement(with: ephemeral)
        let key = secret.hkdfDerivedSymmetricKey(using: SHA256.self, salt: Data("umn-e2e-v1".utf8),
                                                  sharedInfo: record.address.bytes, outputByteCount: 32)
        let box = try ChaChaPoly.SealedBox(combined: payload.combinedCiphertext)
        let data = try ChaChaPoly.open(box, using: key)
        let decoder = PropertyListDecoder()
        let signed = try decoder.decode(Signed.self, from: data)
        let inner = try decoder.decode(InnerFrame.self, from: signed.inner)
        guard inner.sourceRecord.address == expectedSource, inner.sourceRecord.validate() else { throw UMNError.invalidIdentity }
        let keyVerify = try Curve25519.Signing.PublicKey(rawRepresentation: inner.sourceRecord.signingPublicKey)
        guard keyVerify.isValidSignature(signed.signature, for: signed.inner) else { throw UMNError.invalidIdentity }
        return inner
    }
}