import Foundation
import XCTest

@testable import NucleicMDM

final class MDMProtocolTests: XCTestCase {
    private func plist(_ dict: [String: Any]) throws -> Data {
        try MDMPlist.serialize(dict)
    }

    func testParseAuthenticate() throws {
        let body = try plist(["MessageType": "Authenticate", "UDID": "ABC-123", "Topic": "com.apple.mgmt.External.x"])
        let msg = try MDMCheckInMessage.parse(body)
        XCTAssertEqual(msg, .authenticate(udid: "ABC-123", topic: "com.apple.mgmt.External.x"))
        XCTAssertEqual(msg.udid, "ABC-123")
    }

    func testParseTokenUpdateHexEncodesToken() throws {
        let body = try plist([
            "MessageType": "TokenUpdate", "UDID": "ABC-123",
            "PushMagic": "MAGIC", "Token": Data([0xDE, 0xAD, 0xBE, 0xEF]),
        ])
        let msg = try MDMCheckInMessage.parse(body)
        XCTAssertEqual(msg, .tokenUpdate(udid: "ABC-123", pushMagic: "MAGIC", tokenHex: "deadbeef"))
    }

    func testParseCheckOut() throws {
        let body = try plist(["MessageType": "CheckOut", "UDID": "ABC-123"])
        XCTAssertEqual(try MDMCheckInMessage.parse(body), .checkOut(udid: "ABC-123"))
    }

    func testUnmodeledMessageTypeParsesAsOtherNotThrows() throws {
        // A 400 on the check-in channel aborts enrollment, and modern mdmclient sends types we don't
        // model (SetBootstrapToken/GetBootstrapToken/…) during a normal handshake. They must parse as
        // `.other` (the server acks them with an empty 200), not throw (finding 8 / Workstream 2.2).
        let body = try plist(["MessageType": "GetBootstrapToken", "UDID": "x"])
        XCTAssertEqual(try MDMCheckInMessage.parse(body), .other(type: "GetBootstrapToken", udid: "x"))
    }

    func testCheckInWithoutMessageTypeStillThrows() throws {
        // A body that is not a plist with a MessageType is still a hard error (the 400 case).
        let body = try plist(["UDID": "x"])
        XCTAssertThrowsError(try MDMCheckInMessage.parse(body))
    }

    func testInstallProfileCommandPlistShape() throws {
        let profile = Data("<plist/>".utf8)
        let cmd = MDMCommand.installProfile(profile, uuid: "CMD-1")
        let dict = try MDMPlist.parse(try cmd.plist())
        XCTAssertEqual(dict["CommandUUID"] as? String, "CMD-1")
        let command = try XCTUnwrap(dict["Command"] as? [String: Any])
        XCTAssertEqual(command["RequestType"] as? String, "InstallProfile")
        XCTAssertEqual(command["Payload"] as? Data, profile)
    }

    func testRemoveProfileCommandPlistShape() throws {
        let cmd = MDMCommand.removeProfile(identifier: "xyz.nucleic.fda", uuid: "CMD-2")
        let command = try XCTUnwrap(try MDMPlist.parse(try cmd.plist())["Command"] as? [String: Any])
        XCTAssertEqual(command["RequestType"] as? String, "RemoveProfile")
        XCTAssertEqual(command["Identifier"] as? String, "xyz.nucleic.fda")
    }
}

final class MDMCommandQueueTests: XCTestCase {
    private func idle(_ udid: String = "DEV") -> MDMCommandStatusReport {
        MDMCommandStatusReport(status: "Idle", commandUUID: nil, udid: udid)
    }
    private func ack(_ uuid: String, _ udid: String = "DEV") -> MDMCommandStatusReport {
        MDMCommandStatusReport(status: "Acknowledged", commandUUID: uuid, udid: udid)
    }

    func testDrainDeliversCommandsOneByOneInOrder() async throws {
        let q = MDMCommandQueue()
        await q.enqueue([
            MDMCommand.installProfile(Data("a".utf8), uuid: "A"),
            MDMCommand.installProfile(Data("b".utf8), uuid: "B"),
        ])

        // Initial Idle → first command.
        let first = await q.next(afterReport: idle())
        XCTAssertEqual(first?.uuid, "A")
        // Before ack, the queue must NOT advance even on another poll — it re-offers the in-flight.
        let stillFirst = await q.next(afterReport:
            MDMCommandStatusReport(status: "NotNow", commandUUID: "A", udid: "DEV"))
        XCTAssertEqual(stillFirst?.uuid, "A")
        // Ack A → B.
        let second = await q.next(afterReport: ack("A"))
        XCTAssertEqual(second?.uuid, "B")
        // Ack B → empty (queue drained).
        let none = await q.next(afterReport: ack("B"))
        XCTAssertNil(none)

        let drained = await q.isDrained
        let acked = await q.acknowledgedCount
        XCTAssertTrue(drained)
        XCTAssertEqual(acked, 2)
    }

    func testErrorReportStillAdvancesAndIsRecorded() async throws {
        let q = MDMCommandQueue()
        await q.enqueue(MDMCommand.installProfile(Data("x".utf8), uuid: "X"))
        _ = await q.next(afterReport: idle())
        let next = await q.next(afterReport:
            MDMCommandStatusReport(status: "Error", commandUUID: "X", udid: "DEV"))
        XCTAssertNil(next)
        let completed = await q.completedCommands
        XCTAssertEqual(completed.count, 1)
        XCTAssertEqual(completed.first?.status, "Error")
        let acked = await q.acknowledgedCount
        XCTAssertEqual(acked, 0)
    }

    func testReportForWrongUUIDIsIgnored() async throws {
        let q = MDMCommandQueue()
        await q.enqueue(MDMCommand.installProfile(Data("x".utf8), uuid: "X"))
        _ = await q.next(afterReport: idle())  // X in flight
        let matched = await q.report(status: "Acknowledged", commandUUID: "WRONG")
        XCTAssertFalse(matched)
        let flight = await q.inFlightUUID
        XCTAssertEqual(flight, "X")  // still in flight
    }

    func testCompletionStatusesAreScopedToTheRequestedBatch() async throws {
        let q = MDMCommandQueue()
        await q.enqueue([
            MDMCommand.installProfile(Data("old".utf8), uuid: "OLD"),
            MDMCommand.installProfile(Data("ok".utf8), uuid: "OK"),
            MDMCommand.installProfile(Data("bad".utf8), uuid: "BAD"),
        ])
        _ = await q.next(afterReport: idle())
        _ = await q.next(afterReport: ack("OLD"))
        _ = await q.next(afterReport: ack("OK"))
        _ = await q.next(afterReport:
            MDMCommandStatusReport(status: "Error", commandUUID: "BAD", udid: "DEV"))

        let batch = await q.completionStatuses(for: ["OK", "BAD", "MISSING"])
        XCTAssertEqual(batch, ["OK": "Acknowledged", "BAD": "Error"])
        XCTAssertNil(batch["OLD"])
        XCTAssertNil(batch["MISSING"])
    }
}
