import Foundation
import Network
import Observation

public struct DiscoveredHost: Identifiable, Hashable, Sendable {
    public let id: String
    public let displayName: String
    public let serviceType: String
    public let host: String?
    public let port: Int?
    fileprivate let endpoint: BonjourEndpointReference?

    public init(
        id: String,
        displayName: String,
        serviceType: String,
        host: String? = nil,
        port: Int? = nil
    ) {
        self.id = id
        self.displayName = displayName
        self.serviceType = serviceType
        self.host = host
        self.port = port
        self.endpoint = nil
    }

    fileprivate init(
        id: String,
        displayName: String,
        serviceType: String,
        endpoint: BonjourEndpointReference?
    ) {
        self.id = id
        self.displayName = displayName
        self.serviceType = serviceType
        self.host = nil
        self.port = nil
        self.endpoint = endpoint
    }
}

fileprivate struct BonjourEndpointReference: Hashable, Sendable {
    let endpoint: NWEndpoint

    static func == (lhs: BonjourEndpointReference, rhs: BonjourEndpointReference) -> Bool {
        lhs.endpoint.debugDescription == rhs.endpoint.debugDescription
    }

    func hash(into hasher: inout Hasher) {
        hasher.combine(endpoint.debugDescription)
    }
}

public struct ResolvedHost: Sendable, Hashable {
    public let host: String
    public let port: Int
}

public enum DiscoveryError: Error, Sendable {
    case unresolvable
    case timedOut
    case unsupported
}

@Observable
@MainActor
public final class DiscoveryService {
    public private(set) var hosts: [DiscoveredHost] = []
    public private(set) var isBrowsing = false

    private var browsers: [NWBrowser] = []

    public init() {}

    public func start() {
        guard !isBrowsing else { return }
        isBrowsing = true
        hosts = []

        for type in DiscoveryService.serviceTypes {
            let descriptor = NWBrowser.Descriptor.bonjour(type: type, domain: nil)
            let browser = NWBrowser(for: descriptor, using: .tcp)
            browser.browseResultsChangedHandler = { [weak self] results, _ in
                let snapshot = results
                Task { @MainActor in
                    self?.merge(results: snapshot, serviceType: type)
                }
            }
            browser.start(queue: .main)
            browsers.append(browser)
        }
    }

    public func stop() {
        browsers.forEach { $0.cancel() }
        browsers.removeAll()
        isBrowsing = false
    }

    public func resolve(_ host: DiscoveredHost) async throws -> ResolvedHost {
        guard let reference = host.endpoint else {
            if let h = host.host, let p = host.port {
                return ResolvedHost(host: h, port: p)
            }
            throw DiscoveryError.unsupported
        }
        return try await Self.resolve(endpoint: reference.endpoint,
                                      defaultPort: host.serviceType == "_rfb._tcp" ? 5900 : 5900)
    }

    nonisolated static let serviceTypes = ["_rfb._tcp", "_workstation._tcp"]

    private func merge(results: Set<NWBrowser.Result>, serviceType: String) {
        let newHosts = results.compactMap { result -> DiscoveredHost? in
            guard case let .service(name, type, _, _) = result.endpoint else { return nil }
            let id = "\(type)\(name)"
            return DiscoveredHost(
                id: id,
                displayName: name,
                serviceType: serviceType,
                endpoint: BonjourEndpointReference(endpoint: result.endpoint)
            )
        }
        var merged = hosts.filter { $0.serviceType != serviceType }
        merged.append(contentsOf: newHosts)
        hosts = merged.sorted { $0.displayName < $1.displayName }
    }

    nonisolated private static func resolve(endpoint: NWEndpoint,
                                            defaultPort: Int) async throws -> ResolvedHost {
        let parameters = NWParameters.tcp
        let connection = NWConnection(to: endpoint, using: parameters)
        let resumeBox = ResumeBox()

        return try await withCheckedThrowingContinuation { continuation in
            connection.stateUpdateHandler = { state in
                switch state {
                case .ready:
                    if let endpoint = connection.currentPath?.remoteEndpoint,
                       let resolvedHost = Self.parseEndpoint(endpoint, defaultPort: defaultPort) {
                        if resumeBox.tryClaim() {
                            connection.cancel()
                            continuation.resume(returning: resolvedHost)
                        }
                    } else {
                        if resumeBox.tryClaim() {
                            connection.cancel()
                            continuation.resume(throwing: DiscoveryError.unresolvable)
                        }
                    }
                case .failed(let error):
                    if resumeBox.tryClaim() {
                        connection.cancel()
                        continuation.resume(throwing: error)
                    }
                case .cancelled:
                    if resumeBox.tryClaim() {
                        continuation.resume(throwing: CancellationError())
                    }
                default:
                    break
                }
            }
            connection.start(queue: .global(qos: .userInitiated))
        }
    }

    private final class ResumeBox: @unchecked Sendable {
        private let lock = NSLock()
        private var resumed = false
        func tryClaim() -> Bool {
            lock.lock()
            defer { lock.unlock() }
            guard !resumed else { return false }
            resumed = true
            return true
        }
    }

    nonisolated private static func parseEndpoint(_ endpoint: NWEndpoint,
                                                  defaultPort: Int) -> ResolvedHost? {
        switch endpoint {
        case .hostPort(let host, let port):
            let hostString: String
            switch host {
            case .name(let name, _): hostString = name
            case .ipv4(let addr): hostString = "\(addr)"
            case .ipv6(let addr): hostString = "\(addr)"
            @unknown default: return nil
            }
            return ResolvedHost(host: hostString, port: Int(port.rawValue))
        default:
            return nil
        }
    }
}
