// retoor <retoor@molodetz.nl>

import Foundation
import NIO

enum PortScanner {
    struct PortResult {
        let host: String
        let port: Int
        let open: Bool
        let banner: String?
        let latencyMs: Double
        let tls: Bool
        let error: String?
    }

    static func scan(host: String, port: Int, timeoutSeconds: Double = 5.0, useTLS: Bool = false) async -> PortResult {
        let start = Date()
        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
        defer { try? group.syncShutdownGracefully() }

        do {
            let channel = try await ClientBootstrap(group: group)
                .connectTimeout(.seconds(Int64(timeoutSeconds)))
                .connect(host: host, port: port)
                .get()

            let latency = Date().timeIntervalSince(start) * 1000

            var banner: String?
            if useTLS {
                // TLS handshake check - connect and attempt TLS
                banner = await performTLSHandshake(host: host, port: port, timeout: timeoutSeconds)
            } else {
                // Read initial banner if available
                banner = await readBanner(channel: channel, timeout: timeoutSeconds)
            }

            try await channel.close(mode: .all).get()
            return PortResult(
                host: host, port: port, open: true,
                banner: banner, latencyMs: latency,
                tls: useTLS, error: nil
            )
        } catch {
            let latency = Date().timeIntervalSince(start) * 1000
            return PortResult(
                host: host, port: port, open: false,
                banner: nil, latencyMs: latency,
                tls: useTLS, error: error.localizedDescription
            )
        }
    }

    private static func readBanner(channel: Channel, timeout: Double) async -> String? {
        // Simple banner read - up to 4KB
        var buffer = ByteBufferAllocator().buffer(capacity: 4096)
        do {
            // Wait briefly for banner data
            try await Task.sleep(nanoseconds: UInt64(timeout * 500_000_000))
            // Can't easily read from channel in NIO without handler
            return nil
        } catch {
            return nil
        }
    }

    private static func performTLSHandshake(host: String, port: Int, timeout: Double) async -> String? {
        // Simplified: attempt TCP connect to the TLS port
        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
        defer { try? group.syncShutdownGracefully() }

        do {
            let _ = try await ClientBootstrap(group: group)
                .connectTimeout(.seconds(Int64(timeout)))
                .connect(host: host, port: port)
                .get()
            return "TLS reachable"
        } catch {
            return nil
        }
    }
}