/*
 * Copyright (c) 2026 Meshtastic LLC
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
@file:Suppress("SwallowedException")

package org.meshtastic.core.network.repository

import android.annotation.SuppressLint
import android.net.nsd.NsdManager
import android.net.nsd.NsdServiceInfo
import android.os.Build
import co.touchlab.kermit.Logger
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.asExecutor
import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withTimeoutOrNull
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.coroutines.resume

private const val RESOLVE_TIMEOUT_MS = 10000L
private const val RESOLVE_BACKOFF_MS = 1000L

@Suppress("TooGenericExceptionCaught")
@OptIn(ExperimentalCoroutinesApi::class)
internal fun NsdManager.serviceList(
    serviceType: String,
    protocolType: Int = NsdManager.PROTOCOL_DNS_SD,
): Flow<List<NsdServiceInfo>> = callbackFlow {
    val resolvedServices = CopyOnWriteArrayList<NsdServiceInfo>()
    val resolveChannel = Channel<NsdServiceInfo>(Channel.UNLIMITED)
    val mutex = Mutex()

    launch {
        for (service in resolveChannel) {
            mutex.withLock {
                try {
                    val resolved = withTimeoutOrNull(RESOLVE_TIMEOUT_MS) { resolveService(service) }
                    if (resolved != null) {
                        resolvedServices.removeAll { it.serviceName == resolved.serviceName }
                        resolvedServices.add(resolved)
                        trySend(resolvedServices.toList())
                    }
                } catch (e: IllegalArgumentException) {
                    Logger.e(e) { "NSD resolution failed for ${service.serviceName}" }
                    delay(RESOLVE_BACKOFF_MS)
                } catch (e: kotlinx.coroutines.CancellationException) {
                    throw e
                } catch (e: Exception) {
                    Logger.e(e) { "NSD resolution failed for ${service.serviceName}" }
                    delay(RESOLVE_BACKOFF_MS)
                }
            }
        }
    }

    val discoveryListener =
        object : NsdManager.DiscoveryListener {
            override fun onStartDiscoveryFailed(serviceType: String, errorCode: Int) {
                cancel("Start Discovery failed: Error code: $errorCode")
            }

            override fun onStopDiscoveryFailed(serviceType: String, errorCode: Int) {
                cancel("Stop Discovery failed: Error code: $errorCode")
            }

            override fun onDiscoveryStarted(serviceType: String) {
                Logger.d { "NSD Service discovery started" }
            }

            override fun onDiscoveryStopped(serviceType: String) {
                Logger.d { "NSD Service discovery stopped" }
                close()
            }

            override fun onServiceFound(serviceInfo: NsdServiceInfo) {
                Logger.d { "NSD Service found: $serviceInfo" }
                resolveChannel.trySend(serviceInfo)
            }

            override fun onServiceLost(serviceInfo: NsdServiceInfo) {
                Logger.d { "NSD Service lost: $serviceInfo" }
                resolvedServices.removeAll { it.serviceName == serviceInfo.serviceName }
                trySend(resolvedServices.toList())
            }
        }
    trySend(emptyList()) // Emit an initial empty list
    discoverServices(serviceType, protocolType, discoveryListener)

    awaitClose {
        try {
            stopServiceDiscovery(discoveryListener)
        } catch (ex: IllegalArgumentException) {
            // ignore if discovery is already stopped
        }
    }
}

@SuppressLint("NewApi")
private suspend fun NsdManager.resolveService(serviceInfo: NsdServiceInfo): NsdServiceInfo? =
    suspendCancellableCoroutine { continuation ->
        val isResumed = AtomicBoolean(false)
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.UPSIDE_DOWN_CAKE) {
            val callback =
                object : NsdManager.ServiceInfoCallback {
                    override fun onServiceInfoCallbackRegistrationFailed(errorCode: Int) {
                        if (isResumed.compareAndSet(false, true)) {
                            continuation.resume(null)
                        }
                    }

                    override fun onServiceUpdated(updatedServiceInfo: NsdServiceInfo) {
                        if (updatedServiceInfo.hostAddresses.isNotEmpty()) {
                            if (isResumed.compareAndSet(false, true)) {
                                continuation.resume(updatedServiceInfo)
                                try {
                                    unregisterServiceInfoCallback(this)
                                } catch (e: IllegalArgumentException) {
                                    Logger.w(e) { "Already unregistered" }
                                }
                            }
                        }
                    }

                    override fun onServiceLost() {
                        if (isResumed.compareAndSet(false, true)) {
                            continuation.resume(null)
                            try {
                                unregisterServiceInfoCallback(this)
                            } catch (e: IllegalArgumentException) {
                                Logger.w(e) { "Already unregistered" }
                            }
                        }
                    }

                    override fun onServiceInfoCallbackUnregistered() {
                        // No op
                    }
                }
            registerServiceInfoCallback(serviceInfo, Dispatchers.Main.asExecutor(), callback)
            continuation.invokeOnCancellation {
                try {
                    unregisterServiceInfoCallback(callback)
                } catch (e: IllegalArgumentException) {
                    Logger.w(e) { "Already unregistered" }
                }
            }
        } else {
            val listener =
                object : NsdManager.ResolveListener {
                    override fun onResolveFailed(serviceInfo: NsdServiceInfo, errorCode: Int) {
                        if (isResumed.compareAndSet(false, true)) {
                            continuation.resume(null)
                        }
                    }

                    override fun onServiceResolved(serviceInfo: NsdServiceInfo) {
                        if (isResumed.compareAndSet(false, true)) {
                            continuation.resume(serviceInfo)
                        }
                    }
                }
            @Suppress("DEPRECATION")
            resolveService(serviceInfo, listener)
        }
    }
