import Foundation
import Testing
import CryptoKit
@testable import SSHProtocolKit
@testable import SecretKit
@testable import SecretAgentKit

@Suite struct AgentTests {

    // MARK: Identity Listing

    @Test func emptyStores() async throws {
        let agent = Agent(storeList: SecretStoreList())
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestIdentities)
        let response = await agent.handle(request: request, provenance: .test)
        #expect(response == Constants.Responses.requestIdentitiesEmpty)
    }

    @Test func identitiesList() async throws {
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
        let agent = Agent(storeList: list)
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestIdentities)
        let response = await agent.handle(request: request, provenance: .test)

        let actual = OpenSSHReader(data: response)
        let expected = OpenSSHReader(data: Constants.Responses.requestIdentitiesMultiple)
        print(actual, expected)
        #expect(response == Constants.Responses.requestIdentitiesMultiple)
    }

    // MARK: Signatures

    @Test func noMatchingIdentities() async throws {
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
        let agent = Agent(storeList: list)
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestSignatureWithNoneMatching)
        let response = await agent.handle(request: request, provenance: .test)
        #expect(response == Constants.Responses.requestFailure)
    }

    @Test func ecdsaSignature() async throws {
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestSignature)
        guard case SSHAgent.Request.signRequest(let context) = request else { return }
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
        let agent = Agent(storeList: list)
        let response = await agent.handle(request: request, provenance: .test)
        let responseReader = OpenSSHReader(data: response)
        let length = try responseReader.readNextBytes(as: UInt32.self)
        let type = try responseReader.readNextBytes(as: UInt8.self)
        #expect(length == response.count - MemoryLayout<UInt32>.size)
        #expect(type == SSHAgent.Response.agentSignResponse.rawValue)
        let outer = OpenSSHReader(data: responseReader.remaining)
        let inner = try outer.readNextChunkAsSubReader()
        _ = try inner.readNextChunk()
        let rsData = try inner.readNextChunkAsSubReader()
        var r = try rsData.readNextChunk()
        var s = try rsData.readNextChunk()
        // This is fine IRL, but it freaks out CryptoKit
        if r[0] == 0 {
            r.removeFirst()
        }
        if s[0] == 0 {
            s.removeFirst()
        }
        var rs = r
        rs.append(s)
        let signature = try P256.Signing.ECDSASignature(rawRepresentation: rs)
        // Correct signature
        #expect(try P256.Signing.PublicKey(x963Representation: Constants.Secrets.ecdsa256Secret.publicKey)
            .isValidSignature(signature, for: context.dataToSign))
    }

    // MARK: Witness protocol

    @Test func witnessObjectionStopsRequest() async throws {
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret])
        let witness = StubWitness(speakNow: { _,_  in
            return true
        }, witness: { _, _ in })
        let agent = Agent(storeList: list, witness: witness)
        let response = await agent.handle(request: .signRequest(.empty), provenance: .test)
        #expect(response == Constants.Responses.requestFailure)
    }

    @Test func witnessSignature() async throws {
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret])
        nonisolated(unsafe) var witnessed = false
        let witness = StubWitness(speakNow: { _, trace  in
            return false
        }, witness: { _, trace in
            witnessed = true
        })
        let agent = Agent(storeList: list, witness: witness)
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestSignature)
        _ = await agent.handle(request: request, provenance: .test)
        #expect(witnessed)
    }

    @Test func requestTracing() async throws {
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret])
        nonisolated(unsafe) var speakNowTrace: SigningRequestProvenance?
        nonisolated(unsafe) var witnessTrace: SigningRequestProvenance?
        let witness = StubWitness(speakNow: { _, trace  in
            speakNowTrace = trace
            return false
        }, witness: { _, trace in
            witnessTrace = trace
        })
        let agent = Agent(storeList: list, witness: witness)
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestSignature)
        _ = await agent.handle(request: request, provenance: .test)
        #expect(witnessTrace == speakNowTrace)
        #expect(witnessTrace == .test)
    }

    // MARK: Exception Handling

    @Test func signatureException() async throws {
        let list = await storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
        let store = await list.stores.first?.base as! Stub.Store
        store.shouldThrow = true
        let agent = Agent(storeList: list)
        let request = try SSHAgentInputParser().parse(data: Constants.Requests.requestSignature)
        let response = await agent.handle(request: request, provenance: .test)
        #expect(response == Constants.Responses.requestFailure)
    }

    // MARK: Unsupported

    @Test func unhandledAdd() async throws {
        let agent = Agent(storeList: SecretStoreList())
        let response = await agent.handle(request: .addIdentity, provenance: .test)
        #expect(response == Constants.Responses.requestFailure)
    }

}

extension SigningRequestProvenance {

    static let test = SigningRequestProvenance(root: .init(pid: 0, processName: "test", appName: nil, iconURL: nil, path: "/", validSignature: true, parentPID: 0))

}

extension AgentTests {

    @MainActor func storeList(with secrets: [Stub.Secret]) async -> SecretStoreList {
        let store = Stub.Store()
        store.secrets.append(contentsOf: secrets)
        let storeList = SecretStoreList()
        storeList.add(store: store)
        return storeList
    }

    enum Constants {

        enum Requests {
            static let requestIdentities = Data(base64Encoded: "AAAAAQs=")!
            static let requestSignatureWithNoneMatching = Data(base64Encoded: "AAABhA0AAACIAAAAE2VjZHNhLXNoYTItbmlzdHAzODQAAAAIbmlzdHAzODQAAABhBEqCbkJbOHy5S1wVCaJoKPmpS0egM4frMqllgnlRRQ/Uvnn6EVS8oV03cPA2Bz0EdESyRKA/sbmn0aBtgjIwGELxu45UXEW1TEz6TxyS0u3vuIqR3Wo1CrQWRDnkrG/pBQAAAO8AAAAgbqmrqPUtJ8mmrtaSVexjMYyXWNqjHSnoto7zgv86xvcyAAAAA2dpdAAAAA5zc2gtY29ubmVjdGlvbgAAAAlwdWJsaWNrZXkBAAAAE2VjZHNhLXNoYTItbmlzdHAzODQAAACIAAAAE2VjZHNhLXNoYTItbmlzdHAzODQAAAAIbmlzdHAzODQAAABhBEqCbkJbOHy5S1wVCaJoKPmpS0egM4frMqllgnlRRQ/Uvnn6EVS8oV03cPA2Bz0EdESyRKA/sbmn0aBtgjIwGELxu45UXEW1TEz6TxyS0u3vuIqR3Wo1CrQWRDnkrG/pBQAAAAA=")!
            static let requestSignature = Data(base64Encoded: "AAABRA0AAABoAAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBKzOkUiVJEcACMtAd9X7xalbc0FYZyhbmv2dsWl4IP2GWIi+RcsaHQNw+nAIQ8CKEYmLnl0VLDp5Ef8KMhgIy08AAADPAAAAIBIFsbCZ4/dhBmLNGHm0GKj7EJ4N8k/jXRxlyg+LFIYzMgAAAANnaXQAAAAOc3NoLWNvbm5lY3Rpb24AAAAJcHVibGlja2V5AQAAABNlY2RzYS1zaGEyLW5pc3RwMjU2AAAAaAAAABNlY2RzYS1zaGEyLW5pc3RwMjU2AAAACG5pc3RwMjU2AAAAQQSszpFIlSRHAAjLQHfV+8WpW3NBWGcoW5r9nbFpeCD9hliIvkXLGh0DcPpwCEPAihGJi55dFSw6eRH/CjIYCMtPAAAAAA==")!
        }

        enum Responses {
            static let requestIdentitiesEmpty = Data(base64Encoded: "AAAABQwAAAAA")!
            static let requestIdentitiesMultiple = Data(base64Encoded: "AAABLwwAAAACAAAAaAAAABNlY2RzYS1zaGEyLW5pc3RwMjU2AAAACG5pc3RwMjU2AAAAQQSszpFIlSRHAAjLQHfV+8WpW3NBWGcoW5r9nbFpeCD9hliIvkXLGh0DcPpwCEPAihGJi55dFSw6eRH/CjIYCMtPAAAAFWVjZHNhLTI1NkBleGFtcGxlLmNvbQAAAIgAAAATZWNkc2Etc2hhMi1uaXN0cDM4NAAAAAhuaXN0cDM4NAAAAGEEspLMDmreMJverQkqKC9zF9ZUasn5uSWkbRlz1jNTCtuyH1KKm+VImL6wdAj47SbzwM6lEEC24AdfrR64P9i/bnS2i83v/4wQVtcZn+Et13QGgWlZst8lxCPzTookaVwMAAAAFWVjZHNhLTM4NEBleGFtcGxlLmNvbQ==")!
            static let requestFailure = Data(base64Encoded: "AAAAAQU=")!
        }

        enum Secrets {
            static let ecdsa256Secret =  Stub.Secret(keySize: 256, publicKey: Data(base64Encoded: "BKzOkUiVJEcACMtAd9X7xalbc0FYZyhbmv2dsWl4IP2GWIi+RcsaHQNw+nAIQ8CKEYmLnl0VLDp5Ef8KMhgIy08=")!, privateKey: Data(base64Encoded: "BKzOkUiVJEcACMtAd9X7xalbc0FYZyhbmv2dsWl4IP2GWIi+RcsaHQNw+nAIQ8CKEYmLnl0VLDp5Ef8KMhgIy09nw780wy/TSfUmzj15iJkV234AaCLNl+H8qFL6qK8VIg==")!)
            static let ecdsa384Secret =  Stub.Secret(keySize: 384, publicKey: Data(base64Encoded: "BLKSzA5q3jCb3q0JKigvcxfWVGrJ+bklpG0Zc9YzUwrbsh9SipvlSJi+sHQI+O0m88DOpRBAtuAHX60euD/Yv250tovN7/+MEFbXGZ/hLdd0BoFpWbLfJcQj806KJGlcDA==")!, privateKey: Data(base64Encoded: "BLKSzA5q3jCb3q0JKigvcxfWVGrJ+bklpG0Zc9YzUwrbsh9SipvlSJi+sHQI+O0m88DOpRBAtuAHX60euD/Yv250tovN7/+MEFbXGZ/hLdd0BoFpWbLfJcQj806KJGlcDHNapAOzrt9E+9QC4/KYoXS7Uw4pmdAz53uIj02tttiq3c0ZyIQ7XoscWWRqRrz8Kw==")!)
        }

    }

}
