import AppKit
import Darwin
import Foundation
import NucleicCore
import Virtualization

// Standalone de-risking spike for the macOS-VM rework (sibling of `container-spike`). NOT part of the
// app. Drives the REAL `MacVMEngine` end-to-end to prove the Virtualization-framework path on a signed
// build:
//   • default mode — clone the golden base → boot the guest → `ssh` a command in it and capture stdout
//     (the `mac_vm_exec` agent path) → teardown.
//   • `NUCLEIC_MACVM_BUILD_BASE=1` — run the one-time base install instead (download/point-to the macOS
//     restore image and install it into the base bundle). Long-running; then provision separately
//     (scripts/provision-macos-guest.sh) before the default mode can SSH in.
//
// Build + run (must be signed with the virtualization entitlement, exactly like the app):
//   swift build --product macvm-spike
//   codesign --force --sign - --entitlements signing/spike.entitlements .build/.../macvm-spike
//   # 1) build the base once (≈14 GB download + multi-minute install), then provision it:
//   NUCLEIC_MACVM_BUILD_BASE=1 NUCLEIC_MACVM_IPSW=/path/to/UniversalMac.ipsw .build/.../macvm-spike
//   # 2) with a provisioned base present, prove the clone/boot/exec/teardown path:
//   .build/.../macvm-spike

func log(_ s: String) { FileHandle.standardError.write(Data((s + "\n").utf8)) }
func die(_ s: String) -> Never { log("✗ SPIKE FAIL: \(s)"); exit(1) }

func buildBase(_ engine: MacVMEngine) async throws {
    log("[base] MacVMEngine.buildBaseImage — resolve restore image + install macOS …")
    let ipsw = ProcessInfo.processInfo.environment["NUCLEIC_MACVM_IPSW"]
    if let ipsw { log("      using local restore image: \(ipsw)") } else {
        log("      no NUCLEIC_MACVM_IPSW set — fetching the latest supported restore image (~14 GB)")
    }
    try await engine.buildBaseImage(localRestoreImagePath: ipsw)
    log("✓ base install complete. Now provision it (Setup Assistant + scripts/provision-macos-guest.sh)")
    log("  before the default spike mode can SSH into a clone.")
}

func runRoundTrip(_ engine: MacVMEngine) async throws {
    let spec = MacVMSpec(name: "nucleic-macvm-spike", mounts: [], workdir: nil)

    log("[1/3] MacVMEngine.ensureRunning — clone base + boot + wait for SSH …")
    let (name, ip) = try await engine.ensureRunning(spec)
    log("      up: name=\(name) ip=\(ip)")

    log("[2/3] MacVMEngine.run — ssh a command in the guest and capture stdout …")
    let (code, out, err) = try await engine.run(
        name: name, command: "echo macvm-stdio-ok; sw_vers -productVersion", workdir: nil)
    log("      exit=\(code)")
    log("      stdout=\(out)")
    if !err.isEmpty { log("      stderr=\(err)") }

    // Computer-use path: capture the VM screen (needs a computer-use-provisioned base — auto-login +
    // Screen Recording TCC). Best-effort: reports the screenshot size, or why it came back empty.
    log("[2b/3] MacVMEngine.performComputerAction(screenshot) — capture the VM screen …")
    let (shot, shotSummary) = try await engine.performComputerAction(
        name: name, action: "screenshot", x: nil, y: nil, text: nil,
        scrollDirection: nil, scrollAmount: nil, durationMs: nil)
    if let shot {
        log("      screenshot ok: \(shot.count) base64 chars — \(shotSummary)")
    } else {
        log("      screenshot empty (expected unless the base is computer-use-provisioned): \(shotSummary)")
    }

    log("[3/3] MacVMEngine.remove (teardown) …")
    _ = await engine.remove(name: name)

    let ok = code == 0 && out.contains("macvm-stdio-ok")
    log("")
    log("─── RESULTS ───")
    log("  clone + boot + ssh exec + stdout (full MacVMEngine path) : \(ok ? "PASS ✓" : "FAIL ✗")")
    if ok {
        log("✓ SPIKE PASS — the app's macOS-VM clone-and-run path works end-to-end on this machine.")
        exit(0)
    }
    die("macOS-VM path check failed (see above)")
}

// ─────────────────────────────────────────────────────────────────────────────────────────────────
// Linux-guest spike (NUCLEIC_LINUXVM=1)
//
// Proves the Linux VM path on the SAME MacVMEngine: build the Linux base (kernel + rootfs, assembled
// in-guest), then clone → boot → exec (`uname -a`) → readiness → computer-use screenshot → teardown.
//   NUCLEIC_LINUXVM=1 NUCLEIC_LINUXVM_BUILD_BASE=1 .build/.../macvm-spike   # one-time base build
//   NUCLEIC_LINUXVM=1 .build/.../macvm-spike                               # clone/boot/exec/teardown
// ─────────────────────────────────────────────────────────────────────────────────────────────────

func buildLinuxBase(_ engine: MacVMEngine) async throws {
    log("[linux-base] MacVMEngine.buildLinuxBaseImage — download kernel/rootfs + assemble + provision …")
    try await engine.buildLinuxBaseImage()
    log("✓ Linux base build complete (desktop + toolchain provisioned).")
}

func runLinuxRoundTrip(_ engine: MacVMEngine) async throws {
    // A GUI (computer-use) Linux guest, so the virtio-gpu + USB-HID surface path is exercised too.
    let spec = MacVMSpec(name: "nucleic-linuxvm-spike", os: .linux, computerUse: true)

    log("[1/3] MacVMEngine.ensureRunning — clone Linux base + boot + wait for the vsock agent …")
    let (name, ip) = try await engine.ensureRunning(spec)
    log("      up: name=\(name) ip=\(ip)")

    log("[2/3] MacVMEngine.run — exec a command in the Linux guest over vsock …")
    let (code, out, err) = try await engine.run(
        name: name, command: "uname -a; . /etc/os-release 2>/dev/null; echo linuxvm-stdio-ok $PRETTY_NAME",
        workdir: nil)
    log("      exit=\(code)")
    log("      stdout=\(out)")
    if !err.isEmpty { log("      stderr=\(err)") }

    log("[2b/3] MacVMEngine.performComputerAction(screenshot) — capture the Linux desktop …")
    let (shot, shotSummary) = try await engine.performComputerAction(
        name: name, action: "screenshot", x: nil, y: nil, text: nil,
        scrollDirection: nil, scrollAmount: nil, durationMs: nil)
    if let shot {
        log("      screenshot ok: \(shot.count) base64 chars — \(shotSummary)")
    } else {
        log("      screenshot empty: \(shotSummary)")
    }

    log("[3/3] MacVMEngine.remove (teardown) …")
    _ = await engine.remove(name: name)

    let ok = code == 0 && out.contains("linuxvm-stdio-ok")
    log("")
    log("─── LINUX RESULTS ───")
    log("  clone + boot + vsock exec + stdout (full Linux path) : \(ok ? "PASS ✓" : "FAIL ✗")")
    if ok {
        log("✓ LINUX SPIKE PASS — the Linux VM clone-and-run path works end-to-end on this machine.")
        exit(0)
    }
    die("Linux-VM path check failed (see above)")
}

// ─────────────────────────────────────────────────────────────────────────────────────────────────
// Host-side virtual-IO spike (NUCLEIC_MACVM_SPIKE_HOSTIO=1)
//
// Proves the SIP-free computer-use path: capture the guest FRAMEBUFFER host-side via a
// `VZVirtualMachineView` and inject INPUT by delivering synthesized `NSEvent`s to that view — no
// in-guest agent, no TCC, no SIP. Boots a throwaway CLONE of the base (base stays pristine), lets it
// reach a visible screen, captures a PNG, injects a mouse move/click + keystrokes, captures again.
// Writes /tmp/nucleic-spike-hostio-{1,2}.png for inspection and logs whether the frames are non-blank.
// ─────────────────────────────────────────────────────────────────────────────────────────────────

/// The engine's storage layout, mirrored (the spike doesn't reach into the actor for paths).
func macvmStorageRoot() -> URL {
    let base = (try? FileManager.default.url(
        for: .applicationSupportDirectory, in: .userDomainMask, appropriateFor: nil, create: true))
        ?? URL(fileURLWithPath: NSTemporaryDirectory())
    return base.appendingPathComponent("Nucleic/macvms", isDirectory: true)
}

/// Clone the base bundle to a temp dir (copy identity + aux, clonefile the disk) so the spike boots a
/// disposable copy and never dirties the golden base. Returns the clone bundle.
func cloneBaseForSpike() throws -> MacVMBundle {
    let fm = FileManager.default
    let base = MacVMBundle(root: macvmStorageRoot().appendingPathComponent("base", isDirectory: true))
    guard base.isComplete else { die("no complete base bundle at \(base.root.path)") }
    let dest = MacVMBundle(
        root: fm.temporaryDirectory.appendingPathComponent(
            "nucleic-spike-clone-\(UUID().uuidString)", isDirectory: true))
    try fm.createDirectory(at: dest.root, withIntermediateDirectories: true)
    try fm.copyItem(at: base.hardwareModelURL, to: dest.hardwareModelURL)
    try fm.copyItem(at: base.machineIdentifierURL, to: dest.machineIdentifierURL)
    try fm.copyItem(at: base.auxiliaryStorageURL, to: dest.auxiliaryStorageURL)
    if clonefile(base.diskImageURL.path, dest.diskImageURL.path, 0) != 0 {
        try fm.copyItem(at: base.diskImageURL, to: dest.diskImageURL)  // non-APFS fallback
    }
    return dest
}

/// Drives the host-IO spike on the main thread under an AppKit run loop (VZVirtualMachineView is
/// main-thread + main-queue-VM only).
@MainActor
final class HostIOSpike {
    let cloneBundle: MacVMBundle
    var vm: VZVirtualMachine!
    var vmView: VZVirtualMachineView!
    var window: NSWindow!

    init() {
        do { cloneBundle = try cloneBaseForSpike() } catch { die("clone failed: \(error)") }
    }

    /// Run `block` on the main actor after `seconds` (Task-based; the NSApplication run loop drives
    /// the main-actor executor).
    func after(_ seconds: Double, _ block: @escaping @MainActor () -> Void) {
        Task { @MainActor in
            try? await Task.sleep(for: .seconds(seconds))
            block()
        }
    }

    func start() {
        log("[hostio] cloned base → \(cloneBundle.root.lastPathComponent); building config …")
        let config: VZVirtualMachineConfiguration
        do {
            config = try MacVMEngine.makeConfiguration(
                bundle: cloneBundle, cpus: 4, memoryGiB: 8,
                mac: VZMACAddress.randomLocallyAdministered().string, mounts: [])
        } catch { die("makeConfiguration: \(error)") }

        // On-screen window + view so both cacheDisplay and (fallback) window capture have something to
        // render into. The VM is created on the MAIN queue — required for an interactive view.
        let frame = NSRect(x: 0, y: 0, width: 1920, height: 1200)
        vmView = VZVirtualMachineView(frame: frame)
        window = NSWindow(
            contentRect: frame, styleMask: [.titled, .closable, .resizable],
            backing: .buffered, defer: false)
        window.title = "macvm-spike host-IO"
        window.contentView = vmView
        window.setFrameOrigin(NSPoint(x: 80, y: 80))
        window.makeKeyAndOrderFront(nil)

        vm = VZVirtualMachine(configuration: config)  // main-queue (default) VM
        vmView.virtualMachine = vm
        vmView.capturesSystemKeys = true

        log("[hostio] starting guest … (giving it ~45s to reach a visible screen)")
        vm.start { result in
            if case .failure(let e) = result { die("vm.start: \(e)") }
        }
        window.makeFirstResponder(vmView)

        after(45) { self.step1() }
    }

    func step1() {
        log("[hostio] capture #1 (initial screen) …")
        capture(label: "1")
        log("[hostio] injecting input: move → click (center), then type 'hello' + return …")
        injectInput()
        after(4) { self.step2() }
    }

    func step2() {
        log("[hostio] capture #2 (after input) …")
        capture(label: "2")
        log("[hostio] stopping guest + cleaning up clone …")
        if vm.canStop {
            vm.stop { _ in Task { @MainActor in self.finish() } }
        } else {
            finish()
        }
    }

    func finish() {
        try? FileManager.default.removeItem(at: cloneBundle.root)
        log("")
        log("─── HOST-IO RESULTS ───")
        log("  PNGs: /tmp/nucleic-spike-hostio-1.png (initial), -2.png (after input)")
        log("  Inspect them: capture non-blank ⇒ host-side framebuffer capture works (no SIP).")
        log("               a change between #1 and #2 ⇒ host-side input reaches the guest.")
        exit(0)
    }

    /// Host-side capture via `cacheDisplay` (the ideal path — needs NO host permission, unlike SCK).
    /// Writes a PNG and reports whether it's blank so we learn if Metal content renders this way.
    func capture(label: String) {
        guard let rep = vmView.bitmapImageRepForCachingDisplay(in: vmView.bounds) else {
            log("      cacheDisplay: no bitmap rep"); return
        }
        vmView.cacheDisplay(in: vmView.bounds, to: rep)
        writeRep(rep, path: "/tmp/nucleic-spike-hostio-\(label).png", how: "cacheDisplay")
    }

    func writeRep(_ rep: NSBitmapImageRep, path: String, how: String) {
        guard let png = rep.representation(using: .png, properties: [:]) else {
            log("      [\(how)] no PNG data"); return
        }
        try? png.write(to: URL(fileURLWithPath: path))
        log("      [\(how)] wrote \(png.count) bytes → \(path)  (\(uniformity(rep)))")
    }

    /// Cheap non-blank check: sample a grid and report how many distinct colors appear.
    func uniformity(_ rep: NSBitmapImageRep) -> String {
        var colors = Set<UInt32>()
        let w = rep.pixelsWide, h = rep.pixelsHigh
        guard w > 0, h > 0 else { return "empty" }
        for gy in 0..<16 {
            for gx in 0..<16 {
                let x = gx * w / 16, y = gy * h / 16
                if let c = rep.colorAt(x: x, y: y) {
                    let r = UInt32(c.redComponent * 255), g = UInt32(c.greenComponent * 255)
                    let b = UInt32(c.blueComponent * 255)
                    colors.insert((r << 16) | (g << 8) | b)
                }
            }
        }
        return colors.count <= 1 ? "BLANK (1 color)" : "non-blank (\(colors.count) sampled colors)"
    }

    func injectInput() {
        let center = NSPoint(x: vmView.bounds.midX, y: vmView.bounds.midY)
        postMouse(.mouseMoved, at: center)
        postMouse(.leftMouseDown, at: center)
        postMouse(.leftMouseUp, at: center)
        for ch in "hello" { postKey(String(ch)) }
        postKeyCode(0x24)  // return
    }

    func postMouse(_ type: NSEvent.EventType, at p: NSPoint) {
        guard let e = NSEvent.mouseEvent(
            with: type, location: p, modifierFlags: [], timestamp: ProcessInfo.processInfo.systemUptime,
            windowNumber: window.windowNumber, context: nil, eventNumber: 0, clickCount: 1, pressure: 1)
        else { return }
        window.sendEvent(e)
    }

    // Minimal ANSI keycodes for the spike's "hello" + return (the full table lives in
    // VMAgentCore.KeyMap, a different package the spike can't import).
    static let spikeKeyCodes: [Character: CGKeyCode] = [
        "h": 0x04, "e": 0x0E, "l": 0x25, "o": 0x1F,
    ]

    func postKey(_ ch: String) {
        let code = ch.first.flatMap { Self.spikeKeyCodes[$0] } ?? 0
        postKeyCode(code, chars: ch)
    }

    func postKeyCode(_ code: CGKeyCode, chars: String = "") {
        for type in [NSEvent.EventType.keyDown, .keyUp] {
            if let e = NSEvent.keyEvent(
                with: type, location: .zero, modifierFlags: [],
                timestamp: ProcessInfo.processInfo.systemUptime, windowNumber: window.windowNumber,
                context: nil, characters: chars, charactersIgnoringModifiers: chars,
                isARepeat: false, keyCode: UInt16(code))
            {
                window.sendEvent(e)
            }
        }
    }
}

func run() async throws {
    guard MacVMEngine.isSupported else {
        die(MacVMEngine.unsupportedReason ?? "MacVMEngine.isSupported == false")
    }
    // Uses the app's default storage root, so a base built/provisioned by an earlier run (or the
    // Settings action) is found and reused.
    let engine = MacVMEngine()
    if ProcessInfo.processInfo.environment["NUCLEIC_LINUXVM"] == "1" {
        if ProcessInfo.processInfo.environment["NUCLEIC_LINUXVM_BUILD_BASE"] == "1" {
            try await buildLinuxBase(engine)
            exit(0)
        }
        try await runLinuxRoundTrip(engine)
        return
    }
    if ProcessInfo.processInfo.environment["NUCLEIC_MACVM_BUILD_BASE"] == "1" {
        try await buildBase(engine)
        exit(0)
    }
    try await runRoundTrip(engine)
}

// Virtualization.framework delivers VM callbacks on dispatch queues; keep the process alive with a
// run loop while the async work drives the VM (run() calls exit()).
setbuf(stderr, nil)
log("macvm-spike: starting")

if ProcessInfo.processInfo.environment["NUCLEIC_MACVM_SPIKE_HOSTIO"] == "1" {
    // Host-side virtual-IO spike: needs an AppKit run loop (VZVirtualMachineView) rather than
    // dispatchMain. Set up on the main thread and hand control to NSApplication.
    guard MacVMEngine.isSupported else { die(MacVMEngine.unsupportedReason ?? "unsupported") }
    let app = NSApplication.shared
    app.setActivationPolicy(.accessory)
    Task { @MainActor in HostIOSpike().start() }
    app.run()
} else if ProcessInfo.processInfo.environment["NUCLEIC_MACVM_SPIKE_SURFACE"] == "1" {
    // Surface spike: drive the REAL MacVMSurfaceHost protocol + MacVMKeyMap + JPEG capture (the
    // shipping computer-use pipeline) against a live clone. Also an AppKit run loop.
    guard MacVMEngine.isSupported else { die(MacVMEngine.unsupportedReason ?? "unsupported") }
    let app = NSApplication.shared
    app.setActivationPolicy(.accessory)
    Task { @MainActor in SurfaceSpikeDriver().run() }
    app.run()
} else {
    Task {
        do { try await run() } catch { die("threw: \(error)") }
    }
    dispatchMain()
}
