import Foundation
import NetworkExtension

@MainActor
final class PacketTunnelManager: ObservableObject {
    @Published private(set) var status: NEVPNStatus = .invalid
    @Published private(set) var lastDisconnectError: String?
    @Published private(set) var isLoaded = false
    private var manager: NETunnelProviderManager?
    private var observer: NSObjectProtocol?
    private var disconnectErrorTask: Task<Void, Never>?
    private var disconnectGeneration: UInt64 = 0
    private var userRequestedStop = false
    var isUserRequestedStop: Bool { userRequestedStop }

    init() {
        observer = NotificationCenter.default.addObserver(
            forName: .NEVPNStatusDidChange,
            object: nil,
            queue: .main
        ) { [weak self] notification in
            guard let connection = notification.object as? NEVPNConnection else { return }
            let newStatus = connection.status
            Task { @MainActor [weak self] in
                guard let self, connection === self.manager?.connection else { return }
                self.handleStatusChange(newStatus)
            }
        }
    }

    deinit {
        if let observer { NotificationCenter.default.removeObserver(observer) }
    }

    func load() async throws {
        let managers = try await NETunnelProviderManager.loadAllFromPreferences()
        manager = managers.first { ($0.protocolConfiguration as? NETunnelProviderProtocol)?.providerBundleIdentifier == providerId }
        status = manager?.connection.status ?? .invalid
        isLoaded = true
    }

    func installAndStart(credentials: SessionCredentials) async throws {
        userRequestedStop = false
        invalidateDisconnectRecovery()
        lastDisconnectError = nil
        if manager == nil { try await load() }
        let current = manager ?? NETunnelProviderManager()
        let proto = NETunnelProviderProtocol()
        proto.providerBundleIdentifier = providerId
        proto.serverAddress = "HT Portal"
        proto.providerConfiguration = ["sessionId": credentials.sessionId]
        current.protocolConfiguration = proto
        current.localizedDescription = "HT Portal VPN"
        current.isEnabled = true
        try await current.saveToPreferences()
        try await current.loadFromPreferences()
        manager = current
        let encodedCredentials = try JSONEncoder().encode(credentials)
        try current.connection.startVPNTunnel(options: [
            "sessionId": credentials.sessionId as NSString,
            "sessionCredentials": encodedCredentials as NSData,
        ])
        // startVPNTunnel returning without error means that launch was accepted.
        // connection.status can still report the previous .disconnected value
        // until the first NetworkExtension notification arrives.
        status = .connecting
    }

    func stop() {
        userRequestedStop = true
        invalidateDisconnectRecovery()
        lastDisconnectError = nil
        manager?.connection.stopVPNTunnel()
    }

    func diagnostics() async -> TunnelDiagnostics? {
        guard let session = manager?.connection as? NETunnelProviderSession,
              status == .connected || status == .reasserting else { return nil }
        return await withCheckedContinuation { continuation in
            do {
                try session.sendProviderMessage(Data("diagnostics".utf8)) { data in
                    let value = data.flatMap { try? JSONDecoder().decode(TunnelDiagnostics.self, from: $0) }
                    continuation.resume(returning: value)
                }
            } catch {
                continuation.resume(returning: nil)
            }
        }
    }

    private func handleStatusChange(_ newStatus: NEVPNStatus) {
        let previousStatus = status
        status = newStatus
        if newStatus == .connecting || newStatus == .connected || newStatus == .reasserting {
            invalidateDisconnectRecovery()
            lastDisconnectError = nil
        }
        if newStatus == .disconnected,
           previousStatus != .invalid,
           previousStatus != .disconnected,
           !userRequestedStop {
            captureDisconnectError()
        }
    }

    private func captureDisconnectError() {
        guard let connection = manager?.connection else { return }
        invalidateDisconnectRecovery()
        let generation = disconnectGeneration
        disconnectErrorTask = Task { [weak self] in
            for attempt in 0..<6 {
                guard !Task.isCancelled else { return }
                if let error = await Self.disconnectError(connection) {
                    guard let self,
                          !Task.isCancelled,
                          self.disconnectGeneration == generation,
                          self.status == .disconnected else { return }
                    self.lastDisconnectError = Self.describe(error)
                    return
                }
                if attempt < 5 { try? await Task.sleep(for: .milliseconds(400)) }
            }
        }
    }

    private func invalidateDisconnectRecovery() {
        disconnectGeneration &+= 1
        disconnectErrorTask?.cancel()
        disconnectErrorTask = nil
    }

    nonisolated private static func disconnectError(_ connection: NEVPNConnection) async -> Error? {
        await withCheckedContinuation { continuation in
            connection.fetchLastDisconnectError { continuation.resume(returning: $0) }
        }
    }

    nonisolated private static func describe(_ error: Error) -> String {
        let value = error as NSError
        let message = value.localizedDescription.trimmingCharacters(in: .whitespacesAndNewlines)
        let technical = "\(value.domain) (\(value.code))"
        return message.isEmpty ? technical : "\(message)\n\(technical)"
    }

    private let providerId = "net.hivetechnology.o20010.06.packet-tunnel"
}
