import Foundation
import NeXTsoundCore
import USBShim

// A bench probe for one SoundBox: reads and writes its HID feature reports through libusb, as the app does, and
// prints what came back. It never resets, flashes or changes the host's audio defaults.

let vendorID: UInt16 = 0xCAFE
let productID: UInt16 = 0x4014
let hidInterface: UInt16 = 2
let reportValue: UInt16 = 0x0301

final class Board {
    private var context: UnsafeMutableRawPointer?
    private var handle: UnsafeMutableRawPointer?

    init() throws {
        let result = nxs_libusb_open(vendorID, productID, &context, &handle)
        guard result == 0 else {
            throw ProbeError.transport("open \(result)")
        }
    }

    deinit {
        nxs_libusb_close(context, handle)
    }

    /// Sends raw report bytes and returns the raw reply bytes.
    func transfer(_ report: Data) throws -> Data {
        var reply = [UInt8](repeating: 0, count: WireCodec.reportSize)
        let count = report.withUnsafeBytes { request in
            reply.withUnsafeMutableBufferPointer { buffer in
                nxs_libusb_exchange_feature_report(
                    context, handle, hidInterface, reportValue,
                    request.bindMemory(to: UInt8.self).baseAddress, buffer.baseAddress, buffer.count, 2000
                )
            }
        }
        guard count == WireCodec.reportSize else {
            throw ProbeError.transport("exchange \(count)")
        }
        return Data(reply)
    }

    /// Sends a request and polls with the identical bytes until its own reply arrives, as the app does for deferred
    /// replies.
    func exchange(_ request: WireMessage) throws -> WireMessage {
        let report = try WireCodec.encode(request)
        let deadline = Date().addingTimeInterval(10)
        while true {
            let reply = try WireCodec.decode(transfer(report))
            if reply.matchesRequest(request) {
                return reply
            }
            guard Date() < deadline else {
                throw ProbeError.transport("no reply to request \(request.requestID)")
            }
            usleep(50000)
        }
    }

    func request(_ operation: WireOperation, minor: UInt8, generation: UInt32 = 0, payload: Data = Data()) throws -> WireMessage {
        let message = WireMessage(
            operation: operation,
            requestID: .random(in: 1 ... .max),
            generation: generation,
            payload: payload,
            protocolMinor: minor
        )
        return try exchange(message)
    }
}

enum ProbeError: Error {
    case transport(String)
    case failed(String)
}

func hex(_ data: Data) -> String {
    data.map { String(format: "%02x", $0) }.joined()
}

func data(fromHex text: String) -> Data {
    var bytes = Data()
    var index = text.startIndex
    while index < text.endIndex {
        let next = text.index(index, offsetBy: 2)
        bytes.append(UInt8(text[index ..< next], radix: 16) ?? 0)
        index = next
    }
    return bytes
}

func result(_ reply: WireMessage) -> String {
    reply.resultCode == 0 ? "ok" : "result \(reply.resultCode)"
}

/// Prints the board's state as key=value lines: capabilities at both protocol minors, status, the active snapshot and,
/// where the firmware offers them, the diagnostics.
func dump(_ board: Board, newestMinor: UInt8) throws {
    for minor in Set([0, newestMinor]).sorted() {
        let reply = try board.request(.capabilities, minor: minor)
        guard reply.resultCode == 0 else {
            print("capabilities.minor\(minor)=\(result(reply))")
            continue
        }
        let capabilities = try WireCodec.decodeCapabilities(reply.payload)
        print(
            "capabilities.minor\(minor) firmware=\(capabilities.firmwareVersion) serial=\(capabilities.serial)",
            "name=\"\(capabilities.deviceName)\" flags=\(capabilities.featureFlags) flash=\(capabilities.flashSizeBytes)",
            "diagnostics=\(capabilities.supportsDiagnostics) reply.minor=\(reply.protocolMinor) generation=\(reply.generation)"
        )
    }
    let status = try board.request(.status, minor: 0)
    let runtime = try WireCodec.decodeRuntimeStatus(status.payload)
    print(
        "status generation=\(status.generation) persisted=\(runtime.persistedGeneration) storage=\(runtime.storageIsAvailable)",
        "streaming=\(runtime.isStreaming) hostMuted=\(runtime.hostIsMuted) hostVolume=\(runtime.hostVolumeDb)",
        "underruns=\(runtime.underrunCount) overruns=\(runtime.overrunCount) maxUs=\(runtime.maxProcessingUs)",
        "deadlineUs=\(runtime.blockDeadlineUs)"
    )
    let state = try board.request(.readState, minor: 0)
    let snapshot = try WireCodec.decodeSnapshot(state.payload)
    print(
        "snapshot generation=\(state.generation) muted=\(snapshot.settings.isMuted) attenuation=\(snapshot.settings.attenuationDb)",
        "revision=\(snapshot.profileRevision) profile=\(snapshot.profileID)"
    )
    print("snapshot.hex=\(hex(state.payload))")
    if newestMinor > 0 {
        let reply = try board.request(.diagnostics, minor: newestMinor)
        guard reply.resultCode == 0 else {
            print("diagnostics=\(result(reply))")
            return
        }
        let diagnostics = try WireCodec.decodeRuntimeDiagnostics(reply.payload)
        print(
            "diagnostics ringUnderruns=\(diagnostics.ringUnderrunCount) i2sStalls=\(diagnostics.i2sStallCount)",
            "rejected=\(diagnostics.rejectedPacketCount) dspFailures=\(diagnostics.dspFailureCount)",
            "deadlineMisses=\(diagnostics.deadlineMissCount) suspends=\(diagnostics.usbSuspendCount) resumes=\(diagnostics.usbResumeCount)"
        )
    }
}

/// Applies a snapshot's raw payload at the board's current generation and returns the reply.
func apply(_ board: Board, payload: Data) throws -> WireMessage {
    let generation = try board.request(.capabilities, minor: 0).generation
    return try board.request(.applyState, minor: 0, generation: generation, payload: payload)
}

/// The active snapshot's payload with profile mute set, everything else unchanged.
func mutedPayload(_ board: Board) throws -> Data {
    var snapshot = try WireCodec.decodeSnapshot(board.request(.readState, minor: 0).payload)
    snapshot.settings.isMuted = true
    return try WireCodec.encodeSnapshot(snapshot)
}

/// Sends ten malformed requests and checks that each is refused while the snapshot and the generation stay unchanged.
func malformed(_ board: Board) throws {
    let before = try board.request(.readState, minor: 0)
    let valid = try WireCodec.encode(WireMessage(operation: .readState, requestID: 0x0BAD_0001, protocolMinor: 0))
    func mutated(_ change: (inout Data) -> Void) -> Data {
        var report = valid
        change(&report)
        return report
    }
    let staleSnapshot = try WireCodec.encode(WireMessage(
        operation: .applyState, requestID: 0x0BAD_0101, generation: before.generation &- 1, payload: before.payload, protocolMinor: 0
    ))
    var invalidValues = before.payload
    invalidValues[24 ..< 28] = Data([0, 0, 0xC0, 0x7F])
    let cases: [(String, Data)] = [
        ("wrong report id", mutated { $0[0] = 2 }),
        ("wrong magic", mutated { $0[4] = 0x58 }),
        ("major 2", mutated { $0[1] = 2 }),
        ("minor 2", mutated { $0[2] = 2 }),
        ("unknown operation 99", mutated { $0[3] = 99 }),
        ("diagnostics at minor 0", mutated { $0[3] = 7 }),
        ("response flag set", mutated { $0[8] = 1 }),
        ("dirty padding", mutated { $0[$0.count - 1] = 1 }),
        ("stale generation apply", staleSnapshot),
        ("invalid snapshot value", try WireCodec.encode(WireMessage(
            operation: .applyState, requestID: 0x0BAD_0102, generation: before.generation, payload: invalidValues, protocolMinor: 0
        ))),
    ]
    for (name, report) in cases {
        let reply = try board.transfer(report)
        let decoded = try? WireCodec.decode(reply)
        let requestID = report.withUnsafeBytes { $0.loadUnaligned(fromByteOffset: 12, as: UInt32.self) }
        let answered = decoded.map { $0.requestID == requestID ? "reply result \($0.resultCode)" : "no reply of its own" }
        print("malformed \(name): \(answered ?? "undecodable reply")")
    }
    let after = try board.request(.readState, minor: 0)
    print("malformed state unchanged=\(after.payload == before.payload) generation \(before.generation)->\(after.generation)")
}

/// Polls status and, where offered, diagnostics every 100 ms for the given seconds and prints each change of the
/// streaming state, of the aggregate counters and of the diagnostic counters.
func watch(_ board: Board, seconds: Double, minor: UInt8) throws {
    let end = Date().addingTimeInterval(seconds)
    var last = ""
    var samples = 0
    while Date() < end {
        let status = try WireCodec.decodeRuntimeStatus(board.request(.status, minor: 0).payload)
        var line = "streaming=\(status.isStreaming) underruns=\(status.underrunCount) overruns=\(status.overrunCount)"
        if minor > 0 {
            let diagnostics = try WireCodec.decodeRuntimeDiagnostics(board.request(.diagnostics, minor: minor).payload)
            line += " ring=\(diagnostics.ringUnderrunCount) i2s=\(diagnostics.i2sStallCount) rejected=\(diagnostics.rejectedPacketCount)"
            line += " dsp=\(diagnostics.dspFailureCount) deadline=\(diagnostics.deadlineMissCount)"
        }
        samples += 1
        if line != last {
            print(String(format: "%.1fs %@ maxUs=%u", seconds - end.timeIntervalSinceNow, line, status.maxProcessingUs))
            last = line
        }
        usleep(100_000)
    }
    print("watch samples=\(samples)")
}

let arguments = Array(CommandLine.arguments.dropFirst())
let board = try Board()
let minor = UInt8(ProcessInfo.processInfo.environment["NXS_MINOR"] ?? "1") ?? 1
switch arguments.first {
case "dump":
    try dump(board, newestMinor: minor)
case "apply-muted":
    print("apply muted: \(result(try apply(board, payload: mutedPayload(board))))")
case "apply-hex":
    print("apply: \(result(try apply(board, payload: data(fromHex: arguments[1]))))")
case "persist":
    let generation = try board.request(.capabilities, minor: 0).generation
    print("persist at \(generation): \(result(try board.request(.persistState, minor: 0, generation: generation)))")
case "malformed":
    try malformed(board)
case "watch":
    try watch(board, seconds: Double(arguments[1]) ?? 5, minor: minor)
default:
    print("usage: BoardProbe dump | apply-muted | apply-hex <hex> | persist | malformed | watch <seconds>")
}