import Foundation
import Testing
@testable import LocalData

@Suite(.serialized)
struct EncryptionHelperTests {
    private let masterKeyService = "LocalData"
    private let keyName = "LocalDataTests.encryption"
    private let payload = Data("payload".utf8)
    private let keychain = MockKeychainHelper()
    private let encryption: EncryptionHelper
    
    init() {
        self.encryption = EncryptionHelper(keychain: keychain)
    }
    
    @Test func aesGCMWithPBKDF2RoundTrip() async throws {
        let policy: SecurityPolicy.EncryptionPolicy = .aes256(
            keyDerivation: .pbkdf2(iterations: 1_000)
        )
        
        let encrypted = try await encryption.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await encryption.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
    }
    
    @Test func chaChaPolyWithHKDFRoundTrip() async throws {
        let policy: SecurityPolicy.EncryptionPolicy = .chacha20Poly1305(
            keyDerivation: .hkdf()
        )
        
        let encrypted = try await encryption.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await encryption.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
    }
    
    @Test func customConfigurationRoundTrip() async throws {
        let customService = "Test.CustomService"
        let customAccount = "Test.CustomAccount"
        let config = EncryptionConfiguration(
            masterKeyService: customService,
            masterKeyAccount: customAccount
        )
        
        await encryption.updateConfiguration(config)
        
        let policy: SecurityPolicy.EncryptionPolicy = .chacha20Poly1305(
            keyDerivation: .hkdf()
        )
        
        let encrypted = try await encryption.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await encryption.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
        
        // Cleanup mock keychain
        try await keychain.deleteAll(service: customService)
    }
    
    @Test func externalProviderWithHKDFRoundTrip() async throws {
        let source = KeyMaterialSource(id: "test.external")
        let provider = StaticKeyMaterialProvider(material: Data(repeating: 7, count: 32))
        await encryption.registerKeyMaterialProvider(provider, for: source)
        
        let policy: SecurityPolicy.EncryptionPolicy = .external(
            source: source,
            keyDerivation: .hkdf()
        )
        
        let encrypted = try await encryption.encrypt(
            payload,
            keyName: keyName,
            policy: policy
        )
        let decrypted = try await encryption.decrypt(
            encrypted,
            keyName: keyName,
            policy: policy
        )
        
        #expect(decrypted == payload)
    }
}

private struct StaticKeyMaterialProvider: KeyMaterialProviding {
    let material: Data

    func keyMaterial(for keyName: String) async throws -> Data {
        material
    }
}