networkcity/Sources/NetworkCityApp/TrafficController.swift

import SwiftUI
import NetworkCityCore

/// Owns the data feed, enriches it (real process names + reverse-DNS org
/// grouping), and bridges per-district traffic to the scene on the main actor.
@MainActor
final class TrafficController: ObservableObject {
    let scene = CityScene(size: CGSize(width: 1200, height: 800))

    @Published var downKBs: Double = 0
    @Published var upKBs: Double = 0
    @Published var districtCount: Int = 0
    @Published var live = false
    @Published var selection: OrgInspection?

    private let source: ConnectionSource = NettopSource()
    private let differ = TrafficDiffer()
    private let namer = ProcessNamer()
    private let prober = LatencyProber()
    private let ownPID = Int32(ProcessInfo.processInfo.processIdentifier)

    private var orgForHost: [String: Org] = [:]   // resolved IP -> org
    private var pending: Set<String> = []         // resolutions in flight
    private var rttForHost: [String: Double] = [:] // measured RTT (ms)
    private var rttPending: Set<String> = []
    private var portForHost: [String: UInt16] = [:]
    private var task: Task<Void, Never>?

    // Inspector state: enough per-org history to render the panel live.
    private var selectedKey: String?
    private var orgTitles: [String: String] = [:]
    private var orgEndpoints: [String: Set<String>] = [:]       // all hosts ever seen per org
    private var orgProcesses: [String: Set<String>] = [:]
    private var lastEndpointBytes: [String: [String: [TrafficClass: (UInt64, UInt64)]]] = [:]
    private var lastInterval: TimeInterval = 2

    func start() {
        guard task == nil else { return }
        scene.onSelect = { [weak self] key in
            MainActor.assumeIsolated { self?.select(key) }
        }
        scene.onReap = { [weak self] key in
            MainActor.assumeIsolated { self?.prune(key) }
        }
        if ProcessInfo.processInfo.environment["NC_DEMO"] != nil {
            startDemo()
            return
        }
        task = Task { [weak self] in
            guard let self else { return }
            for await snap in source.snapshots(every: 2.0) {
                handle(snap)
            }
        }
    }

    /// Offline demo: feeds a synthetic multi-endpoint district through the real
    /// scene path so the hub-and-spoke fan-out can be exercised without nettop.
    private func startDemo() {
        task = Task { [weak self] in
            guard let self else { return }
            var tick = 0
            while !Task.isCancelled {
                tick += 1
                let classes: [TrafficClass] = [.https, .quic, .dns, .http, .other]
                let endpoints = (0..<8).map { i -> EndpointTraffic in
                    let active = (i + tick) % 3 == 0
                    let cls = classes[i % classes.count]
                    let heavy = (cls == .https || cls == .quic) && i % 4 == 0
                    return EndpointTraffic(host: "20.190.\(i).\(10 + i)", flows: [
                        ClassFlow(cls: cls,
                                  inBytes: active ? (heavy ? 400_000 : 36_000) : 0,
                                  outBytes: active ? 12_000 : 0),
                    ])
                }
                let inB = endpoints.reduce(0) { $0 + $1.inBytes }
                let outB = endpoints.reduce(0) { $0 + $1.outBytes }
                scene.apply([DistrictTraffic(key: "onedrive", title: "OneDrive", subtitle: "OneDrive",
                                             inBytes: inB, outBytes: outB, endpoints: endpoints)])

                // Keep the inspector maps live so the demo district is clickable.
                orgTitles["onedrive"] = "OneDrive"
                orgEndpoints["onedrive", default: []].formUnion(endpoints.map(\.host))
                orgProcesses["onedrive", default: []].insert("OneDrive")
                var bytesMap: [String: [TrafficClass: (UInt64, UInt64)]] = [:]
                for e in endpoints { for f in e.flows { bytesMap[e.host, default: [:]][f.cls] = (f.inBytes, f.outBytes) } }
                lastEndpointBytes["onedrive"] = bytesMap
                if let key = selectedKey { selection = buildInspection(key) }
                // Verification hook: auto-open the inspector in debug runs.
                if tick == 3, ProcessInfo.processInfo.environment["NC_DEBUG"] != nil { select("onedrive") }

                downKBs = Double(inB) / 1024 / 2
                upKBs = Double(outB) / 1024 / 2
                districtCount = scene.buildingCount
                live = true
                try? await Task.sleep(nanoseconds: 2_000_000_000)
            }
        }
    }

    func stop() { task?.cancel(); task = nil }

    private func handle(_ snapshot: ConnectionSnapshot) {
        // Exclude our own traffic so we don't render our latency probes.
        let deltas = differ.ingest(snapshot).filter { $0.connection.isExternal && $0.connection.pid != ownPID }

        // Group resolved hosts by org; kick off resolution for the rest. A
        // brand-new host shows up one tick late, which is fine — traffic is
        // continuous, and this avoids spawning an IP building we'd have to merge.
        var districts: [String: DistrictTraffic] = [:]
        // orgKey -> host -> class -> (in, out)
        var endpoints: [String: [String: [TrafficClass: (UInt64, UInt64)]]] = [:]
        for d in deltas {
            let host = d.connection.remote.host
            guard let org = orgForHost[host] else { resolve(host); continue }

            let proc = namer.name(for: d.connection.pid) ?? d.connection.processName
            var district = districts[org.key]
                ?? DistrictTraffic(key: org.key, title: org.name, subtitle: proc,
                                   inBytes: 0, outBytes: 0, endpoints: [])
            district.inBytes += d.bytesInDelta
            district.outBytes += d.bytesOutDelta
            district.subtitle = proc
            districts[org.key] = district

            let cls = TrafficClass.classify(proto: d.connection.proto, port: d.connection.remote.port)
            var hosts = endpoints[org.key] ?? [:]
            var classes = hosts[host] ?? [:]
            let prior = classes[cls] ?? (0, 0)
            classes[cls] = (prior.0 + d.bytesInDelta, prior.1 + d.bytesOutDelta)
            hosts[host] = classes
            endpoints[org.key] = hosts

            orgEndpoints[org.key, default: []].insert(host)
            if let p = d.connection.remote.port, p > 0 { portForHost[host] = UInt16(p) }
            probeIfNeeded(host)
        }
        let rendered = districts.map { key, d -> DistrictTraffic in
            var d = d
            d.endpoints = (endpoints[key] ?? [:]).map { host, classMap in
                EndpointTraffic(host: host, flows: classMap.map {
                    ClassFlow(cls: $0.key, inBytes: $0.value.0, outBytes: $0.value.1)
                })
            }
            // Distance = closest measured edge of this org (the nearest server you're served from).
            d.rttMs = (orgEndpoints[key] ?? []).compactMap { rttForHost[$0] }.min()
            return d
        }
        scene.apply(rendered)

        // Retain per-org detail for the inspector panel.
        lastInterval = deltas.first?.interval ?? 2
        lastEndpointBytes = endpoints
        for d in rendered {
            orgTitles[d.key] = d.title
            orgProcesses[d.key, default: []].insert(d.subtitle)
        }
        if let key = selectedKey { selection = buildInspection(key) }

        // HUD totals reflect *all* external traffic, resolved or not, so the
        // throughput numbers stay honest while districts are still resolving.
        downKBs = deltas.reduce(0) { $0 + $1.inBytesPerSec } / 1024
        upKBs = deltas.reduce(0) { $0 + $1.outBytesPerSec } / 1024
        districtCount = scene.buildingCount
        live = true
    }

    // MARK: - Inspector

    func select(_ key: String?) {
        selectedKey = key
        scene.highlight(key: key)
        selection = key.flatMap(buildInspection)
        if let s = selection, ProcessInfo.processInfo.environment["NC_DEBUG"] != nil {
            FileHandle.standardError.write(Data("[inspect] \(s.title): \(s.endpointCount) pts ▼\(Int(s.downKBs)) ▲\(Int(s.upKBs)) KB/s top=\(s.endpoints.first?.host ?? "-")\n".utf8))
        }
    }

    func deselect() { select(nil) }

    /// A district was demolished — drop its rendering/inspection state. The
    /// resolution + RTT caches are kept (small, bounded by distinct hosts, and
    /// save rework if the district comes back).
    private func prune(_ key: String) {
        orgEndpoints[key] = nil
        orgProcesses[key] = nil
        orgTitles[key] = nil
        lastEndpointBytes[key] = nil
        if selectedKey == key { selectedKey = nil; selection = nil }
    }

    private func buildInspection(_ key: String) -> OrgInspection {
        let hosts = orgEndpoints[key] ?? []
        var infos: [EndpointInfo] = []
        var totalDown = 0.0, totalUp = 0.0

        for host in hosts {
            let classMap = lastEndpointBytes[key]?[host] ?? [:]
            var down = 0.0, up = 0.0
            var top: (cls: TrafficClass, bytes: UInt64)?
            for (cls, bytes) in classMap {
                down += Double(bytes.0) / lastInterval / 1024
                up += Double(bytes.1) / lastInterval / 1024
                let total = bytes.0 + bytes.1
                if top == nil || total > top!.bytes { top = (cls, total) }
            }
            totalDown += down
            totalUp += up
            infos.append(EndpointInfo(host: host, rdns: orgForHost[host]?.resolvedHost,
                                      downKBs: down, upKBs: up, topClass: top?.cls ?? .other,
                                      rttMs: rttForHost[host]))
        }
        infos.sort { ($0.downKBs + $0.upKBs) > ($1.downKBs + $1.upKBs) }

        return OrgInspection(
            key: key, title: orgTitles[key] ?? key,
            downKBs: totalDown, upKBs: totalUp,
            endpointCount: hosts.count, endpoints: infos,
            processes: (orgProcesses[key] ?? []).sorted()
        )
    }

    /// Measure RTT to a newly-seen endpoint once (cached), then store it so the
    /// next tick can glide its district to the right distance.
    private func probeIfNeeded(_ host: String) {
        guard rttForHost[host] == nil, !rttPending.contains(host),
              let port = portForHost[host] else { return }
        rttPending.insert(host)
        Task { [prober] in
            // Unreachable (firewalled) endpoints read as "far" rather than missing.
            let rtt = await prober.measure(host: host, port: port) ?? LatencyLayout.rttCeiling
            rttForHost[host] = rtt
            rttPending.remove(host)
            if ProcessInfo.processInfo.environment["NC_DEBUG"] != nil {
                FileHandle.standardError.write(Data("[rtt] \(host):\(port) → \(Int(rtt))ms\n".utf8))
            }
        }
    }

    /// Resolve a host's org off the main thread, then store it back on the main
    /// actor. Reverse DNS can block for seconds, so it runs detached.
    private func resolve(_ host: String) {
        guard orgForHost[host] == nil, pending.insert(host).inserted else { return }
        Task {
            let org = await Task.detached(priority: .utility) {
                let ptr = ReverseDNS.resolve(host)
                return OrgClassifier.classify(ip: host, ptr: ptr)
            }.value
            orgForHost[host] = org
            pending.remove(host)
            if ProcessInfo.processInfo.environment["NC_DEBUG"] != nil {
                FileHandle.standardError.write(Data("[org] \(host) → \(org.name) (\(org.resolvedHost ?? "no-ptr"))\n".utf8))
            }
        }
    }
}