#!/usr/bin/env python3
"""Only SSH/Mailpit availability and action spacing. Run in OrionToolbar project root."""
from pathlib import Path
from datetime import datetime
import shutil

root=Path.cwd()
cpath=root/'OrionToolbar/OrionVMController.swift'
vpath=root/'OrionToolbar/OrionMenuView.swift'
if not all(p.exists() for p in (cpath,vpath)):
    raise SystemExit('Run from OrionToolbar project root. No changes made.')
c=cpath.read_text(); v=vpath.read_text()
if 'func refreshServiceAvailability() async' in c:
    raise SystemExit('Service availability patch is already installed. No changes made.')

def replace_once(src,old,new,label):
    count=src.count(old)
    if count != 1: raise SystemExit(f'{label}: expected one match, found {count}. No changes made.')
    return src.replace(old,new,1)

c=replace_once(c,'    @Published var busyServers: Set<UUID> = []', '''    @Published var busyServers: Set<UUID> = []
    @Published private(set) var reachableSSH: Set<UUID> = []
    @Published private(set) var reachableMailpit: Set<UUID> = []''','published availability')
anchor='    func openSSH(_ s: ServerDefinition) {'
addition='''    // Availability is probed off the main actor. An open SSH port means
    // the SSH action can be offered; credentials are checked by ssh itself.
    func refreshServiceAvailability() async {
        let candidates = servers.filter { server in
            guard !server.host.isEmpty else { return false }
            return discoveredVMs.contains {
                ($0.id == server.vmUUID || $0.name == server.vmName) &&
                $0.status.lowercased() == "running"
            }
        }
        let results = await withTaskGroup(of: (UUID, Bool, Bool).self,
                                          returning: [(UUID, Bool, Bool)].self) { group in
            for server in candidates {
                group.addTask {
                    async let ssh = Self.tcpPortOpen(host: server.host, port: server.sshPort)
                    async let mailpit = Self.mailpitResponding(host: server.host, port: server.mailpitPort)
                    return (server.id, await ssh, await mailpit)
                }
            }
            var values: [(UUID, Bool, Bool)] = []
            for await result in group { values.append(result) }
            return values
        }
        reachableSSH = Set(results.filter { $0.1 }.map { $0.0 })
        reachableMailpit = Set(results.filter { $0.2 }.map { $0.0 })
    }

    private nonisolated static func tcpPortOpen(host: String, port: Int) async -> Bool {
        guard (1...65535).contains(port), !host.isEmpty else { return false }
        return await Task.detached(priority: .utility) {
            let process = Process()
            process.executableURL = URL(fileURLWithPath: "/usr/bin/nc")
            process.arguments = ["-G", "2", "-z", host, String(port)]
            process.standardOutput = FileHandle.nullDevice
            process.standardError = FileHandle.nullDevice
            do {
                try process.run()
                process.waitUntilExit()
                return process.terminationStatus == 0
            } catch { return false }
        }.value
    }

    private nonisolated static func mailpitResponding(host: String, port: Int?) async -> Bool {
        guard let port, (1...65535).contains(port), !host.isEmpty,
              let url = URL(string: "http://\\(host):\\(port)/") else { return false }
        var request = URLRequest(url: url)
        request.timeoutInterval = 3
        request.httpMethod = "GET"
        do {
            let (_, response) = try await URLSession.shared.data(for: request)
            guard let http = response as? HTTPURLResponse else { return false }
            return (200..<500).contains(http.statusCode)
        } catch { return false }
    }

'''
c=replace_once(c,anchor,addition+anchor,'service helper insertion')

v=replace_once(v,'        .task { await vm.discoverVMs(); await vm.refreshAll() }','''        .task {
            await vm.discoverVMs()
            await vm.refreshAll()
            while !Task.isCancelled {
                await vm.refreshServiceAvailability()
                try? await Task.sleep(for: .seconds(10))
            }
        }''','periodic probe')
v=replace_once(v,'''                action("SSH", "terminal", tint: OrionStyle.blue) { vm.openSSH(server) }
                    .disabled(server.host.isEmpty || server.sshUser.isEmpty)
                action("Mailpit", "globe", tint: OrionStyle.blue) { vm.openMailpit(server) }
                    .disabled(server.host.isEmpty || server.mailpitPort == nil)''','''                if !server.sshUser.isEmpty && vm.reachableSSH.contains(server.id) {
                    action("SSH", "terminal", tint: OrionStyle.blue) { vm.openSSH(server) }
                }
                if vm.reachableMailpit.contains(server.id) {
                    action("Mailpit", "globe", tint: OrionStyle.blue) { vm.openMailpit(server) }
                }''','running VM service buttons')
v=replace_once(v,'''            action("SSH", "terminal", tint: OrionStyle.card) { vm.openSSH(server) }
                .disabled(status.lowercased() != "running" || server.host.isEmpty || server.sshUser.isEmpty)
            action("Mailpit", "globe", tint: OrionStyle.card) { vm.openMailpit(server) }
                .disabled(status.lowercased() != "running" || server.host.isEmpty || server.mailpitPort == nil)''','''            if !server.sshUser.isEmpty && vm.reachableSSH.contains(server.id) {
                action("SSH", "terminal", tint: OrionStyle.card) { vm.openSSH(server) }
            }
            if vm.reachableMailpit.contains(server.id) {
                action("Mailpit", "globe", tint: OrionStyle.card) { vm.openMailpit(server) }
            }''','integrated service buttons')
# Space the right-side More menu away from Stop without altering other layout.
v=replace_once(v,'''} label: {
                Image(systemName: "ellipsis").frame(width: 16)
            }.menuStyle(.borderlessButton).frame(width: 22)''','''} label: {
                Image(systemName: "ellipsis").frame(width: 22)
            }.menuStyle(.borderlessButton).frame(width: 30)
            .padding(.leading, 12)''','More button spacing')
# Add a little breathing room between the two power buttons.
v=replace_once(v,'''                action("Stop", "stop.fill", tint: OrionStyle.red) { vm.operate(machine, action: "stop") }
            }
            Menu {''','''                action("Stop", "stop.fill", tint: OrionStyle.red) { vm.operate(machine, action: "stop") }
                    .padding(.leading, 6)
            }
            Menu {''','Stop button spacing')

backup=root/'backups'/datetime.now().strftime('%Y%m%d-%H%M%S-%f')
for p in (cpath,vpath):
    target=backup/p.relative_to(root)
    target.parent.mkdir(parents=True,exist_ok=True)
    shutil.copy2(p,target)
cpath.write_text(c);vpath.write_text(v)
print('Updated:',cpath.relative_to(root),vpath.relative_to(root))
print('Backups:',backup)