AeroToss/app/src/desktopMain/kotlin/com/aerotoss/discovery/JmdnsDiscovery.kt

package com.aerotoss.discovery

import com.aerotoss.model.Device
import com.aerotoss.model.DeviceType
import com.aerotoss.util.generateDeviceId
import com.aerotoss.util.getDeviceName
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.asStateFlow
import java.net.InetAddress
import java.net.NetworkInterface
import java.util.concurrent.ConcurrentHashMap
import javax.jmdns.JmDNS
import javax.jmdns.ServiceEvent
import javax.jmdns.ServiceListener
import javax.jmdns.ServiceInfo

class JmdnsDiscovery : DiscoveryManager {
    private val _devices = MutableStateFlow<List<Device>>(emptyList())
    override val devices: Flow<List<Device>> = _devices.asStateFlow()

    private var jmdns: JmDNS? = null
    private var serviceListener: ServiceListener? = null
    private var discoveryThread: Thread? = null
    private val deviceId = generateDeviceId()
    private val deviceName = getDeviceName()
    private val discoveredDevices = ConcurrentHashMap<String, Device>()

    override fun startDiscovery(servicePort: Int) {
        if (discoveryThread != null) return

        discoveryThread = Thread {
            try {
                val addr = findLocalAddress() ?: return@Thread
                println("AeroToss: Starting mDNS on ${addr.hostAddress}")
                jmdns = JmDNS.create(addr, "aerotoss-$deviceId")

                val serviceInfo = ServiceInfo.create(
                    SERVICE_TYPE,
                    SERVICE_NAME,
                    servicePort,
                    "path=/ aerotoss=1 id=$deviceId name=$deviceName"
                )
                jmdns?.registerService(serviceInfo)
                println("AeroToss: Registered service on port $servicePort")

                serviceListener = object : ServiceListener {
                    override fun serviceAdded(event: ServiceEvent) {
                        println("AeroToss: Service added: ${event.name} type=${event.type}")
                        jmdns?.requestServiceInfo(event.type, event.name, true)
                    }

                    override fun serviceRemoved(event: ServiceEvent) {
                        println("AeroToss: Service removed: ${event.name}")
                        discoveredDevices.remove(event.name)
                        _devices.value = discoveredDevices.values.toList()
                    }

                    override fun serviceResolved(event: ServiceEvent) {
                        val info = event.info
                        val hostAddresses = info.hostAddresses
                        println("AeroToss: Service resolved: ${event.name} addresses=${hostAddresses.toList()}")
                        if (hostAddresses.isNotEmpty()) {
                            val attributes = info.textBytes?.let { parseAttributes(it) } ?: emptyMap()
                            val id = attributes["id"] ?: event.name
                            val name = attributes["name"] ?: event.name
                            val device = Device(
                                id = id,
                                name = name,
                                type = DeviceType.DESKTOP,
                                hostAddress = hostAddresses.first(),
                                port = info.port
                            )
                            if (id != deviceId) {
                                discoveredDevices[event.name] = device
                                _devices.value = discoveredDevices.values.toList()
                            }
                        }
                    }
                }

                jmdns?.addServiceListener(SERVICE_TYPE, serviceListener)

                // Also listen without the trailing dot for compatibility
                jmdns?.addServiceListener("_aerotoss._tcp", serviceListener)

                println("AeroToss: Listening for services on $SERVICE_TYPE")
                Thread.currentThread().join()
            } catch (e: Exception) {
                e.printStackTrace()
                println("AeroToss: mDNS error: ${e.message}")
            }
        }.apply {
            isDaemon = true
            name = "aerotoss-mdns"
            start()
        }
    }

    override fun stopDiscovery() {
        discoveryThread?.interrupt()
        discoveryThread = null
        try {
            serviceListener?.let { jmdns?.removeServiceListener(SERVICE_TYPE, it) }
            serviceListener?.let { jmdns?.removeServiceListener("_aerotoss._tcp", it) }
            jmdns?.unregisterAllServices()
            jmdns?.close()
        } catch (_: Exception) {
        }
        jmdns = null
        serviceListener = null
        discoveredDevices.clear()
        _devices.value = emptyList()
    }

    override fun dispose() {
        stopDiscovery()
    }

    private fun findLocalAddress(): InetAddress? {
        return try {
            NetworkInterface.getNetworkInterfaces()?.toList()
                ?.filter { it.isUp && !it.isLoopback && !isVirtual(it) }
                ?.flatMap { it.inetAddresses.toList() }
                ?.firstOrNull { it is java.net.Inet4Address }
                ?: InetAddress.getLocalHost()
        } catch (_: Exception) {
            try { InetAddress.getLocalHost() } catch (_: Exception) { null }
        }
    }

    private fun isVirtual(iface: NetworkInterface): Boolean {
        return iface.isVirtual || iface.name.startsWith("vmnet") || iface.name.startsWith("veth")
    }

    private fun parseAttributes(raw: ByteArray): Map<String, String> {
        val result = mutableMapOf<String, String>()
        var i = 0
        while (i < raw.size) {
            val keyLen = raw[i].toInt() and 0xFF
            i++
            if (i + keyLen > raw.size) break
            val key = String(raw, i, keyLen, Charsets.UTF_8)
            i += keyLen
            if (i >= raw.size) break
            val valueLen = raw[i].toInt() and 0xFF
            i++
            if (i + valueLen > raw.size) break
            val value = String(raw, i, valueLen, Charsets.UTF_8)
            i += valueLen
            result[key] = value
        }
        return result
    }

    companion object {
        const val SERVICE_TYPE = "_aerotoss._tcp."
        const val SERVICE_NAME = "AeroToss"
    }
}