//
//  SDLUDPHoleSession.swift
//  punchnet
//
//  Created by 安礼成 on 2026/5/27.
//
import Foundation
import NIOCore

actor SDLUDPHoleSession {
    private let proberActor: SDLNATProberActor
    private let onEvent: SDLUDPHoleService.EventHandler
    private let onData: SDLUDPHoleService.DataHandler

    private var udpHole: SDLUDPHole?
    private var localAddress: SocketAddress?

    init(
        proberActor: SDLNATProberActor,
        onEvent: @escaping SDLUDPHoleService.EventHandler,
        onData: @escaping SDLUDPHoleService.DataHandler
    ) {
        self.proberActor = proberActor
        self.onEvent = onEvent
        self.onData = onData
    }

    func run() async throws {
        do {
            try await self.runV4()
            await self.stop()
        } catch {
            await self.stop()
            throw error
        }
    }

    func stop() async {
        let udpHole = self.udpHole
        self.udpHole = nil
        self.localAddress = nil
        udpHole?.stop()

        await self.proberActor.cancelAll()
    }

    func send(type: SDLPacketType, data: Data, remoteAddress: SocketAddress) async {
        guard case .v4 = remoteAddress else {
            SDLLogger.log("[SDLUDPHoleSession] unsupported socket family: \(remoteAddress)", for: .debug)
            return
        }

        guard let udpHole else {
            SDLLogger.log("[SDLUDPHoleSession] udpHole is nil for remoteAddress: \(remoteAddress)", for: .debug)
            return
        }

        udpHole.send(type: type, data: data, remoteAddress: remoteAddress)
    }

    private func runV4() async throws {
        let udpHole = try SDLUDPHole()
        let localAddress = try udpHole.start()
        self.udpHole = udpHole
        self.localAddress = localAddress

        SDLLogger.log("[SDLUDPHoleSession] udpHole started, on address: \(localAddress)")
        await self.onEvent(.ready(localAddress))

        do {
            try await withThrowingTaskGroup(of: Void.self) { group in
                defer {
                    group.cancelAll()
                }

                group.addTask {
                    try await self.readV4Loop(udpHole: udpHole)
                }

                group.addTask {
                    await self.probeNatType(udpHole: udpHole)
                }

                try await group.waitForAll()
            }
        } catch {
            udpHole.stop()
            if self.udpHole === udpHole {
                self.udpHole = nil
                self.localAddress = nil
            }
            throw error
        }
    }

    private func readV4Loop(udpHole: SDLUDPHole) async throws {
        for try await datagram in udpHole.messageStream {
            try Task.checkCancellation()
            try await self.handleV4Message(remoteAddress: datagram.remoteAddress, message: datagram.message)
        }
    }

    private func probeNatType(udpHole: SDLUDPHole) async {
        if Task.isCancelled {
            return
        }

        let natType = await self.proberActor.probeNatType(using: udpHole)
        if Task.isCancelled {
            return
        }

        await self.onEvent(.natType(natType))
    }

    private func handleV4Message(remoteAddress: SocketAddress, message: SDLHoleMessage) async throws {
        switch message {
        case .control(let control):
            switch control {
            case .stunProbeReply(let probeReply):
                await self.proberActor.handleProbeReply(localAddress: self.localAddress, reply: probeReply)
            default:
                await self.onEvent(.packet(remoteAddress, control))
            }
        case .data(let data):
            await self.onData(data)
        }
    }
}