// SPDX-License-Identifier: GPL-3.0-only

import CryptoKit
import Foundation
import Network
import NetworkExtension
import Security

final class TunnelWebSocketChannel: @unchecked Sendable {
    private static let webSocketGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    private static let maximumHandshakeBytes = 16 * 1024
    private static let maximumFrameBytes = 1 * 1024 * 1024

    private let queue = DispatchQueue(label: "net.hivetechnology.o20010.06.runtime-websocket")
    private let connection: NWConnection
    private let hostHeader: String
    private let resource: String
    private let webSocketKey: String
    private let expectedAccept: String
    private let tlsDiagnostics: TLSValidationDiagnostics
    private var opened: CheckedContinuation<Void, Error>?
    private var openTimeout: DispatchWorkItem?
    private var receiveBuffer = Data()
    private var hasOpened = false
    private var handshakeStarted = false
    private var closing = false
    private var disconnectReported = false
    private var networkState = L10n.text("network.setup", fallback: "setup")
    var onDisconnect: (() -> Void)?
    var onText: ((String) -> Void)?

    @available(iOS 18.0, *)
    init(provider: NEPacketTunnelProvider, url: URL, fingerprint: String?, caCertificatePem: String?) throws {
        guard let interface = provider.virtualInterface else {
            throw TunnelChannelError.missingTunnelInterface
        }
        guard let host = url.host else {
            throw TunnelChannelError.invalidURL
        }
        let portValue = url.port ?? (url.scheme == "wss" ? 443 : 80)
        guard let port = NWEndpoint.Port(rawValue: UInt16(portValue)) else {
            throw TunnelChannelError.invalidURL
        }

        let keyData = try Self.randomBytes(count: 16)
        webSocketKey = keyData.base64EncodedString()
        let acceptSource = Data((webSocketKey + Self.webSocketGUID).utf8)
        expectedAccept = Data(Insecure.SHA1.hash(data: acceptSource)).base64EncodedString()

        let components = URLComponents(url: url, resolvingAgainstBaseURL: false)
        let path = components?.percentEncodedPath.isEmpty == false ? components!.percentEncodedPath : "/"
        if let query = components?.percentEncodedQuery, !query.isEmpty {
            resource = "\(path)?\(query)"
        } else {
            resource = path
        }
        hostHeader = host.contains(":") ? "[\(host)]:\(portValue)" : "\(host):\(portValue)"

        tlsDiagnostics = TLSValidationDiagnostics(secure: url.scheme == "wss")
        let tls = url.scheme == "wss"
            ? Self.tlsOptions(
                fingerprint: fingerprint,
                caCertificatePem: caCertificatePem,
                diagnostics: tlsDiagnostics,
                queue: queue)
            : nil
        let parameters = NWParameters(tls: tls, tcp: NWProtocolTCP.Options())
        parameters.requiredInterface = interface
        connection = NWConnection(host: NWEndpoint.Host(host), port: port, using: parameters)
    }

    func connect(timeout: TimeInterval) async throws {
        try await withCheckedThrowingContinuation { continuation in
            queue.async {
                self.opened = continuation
                self.hasOpened = false
                self.handshakeStarted = false
                self.closing = false
                self.disconnectReported = false
                self.networkState = L10n.text("network.starting", fallback: "starting")
                self.receiveBuffer.removeAll(keepingCapacity: true)
                let timeoutWork = DispatchWorkItem { [weak self] in
                    guard let self else { return }
                    self.fail(TunnelChannelError.timedOut(timeout, self.connectionDiagnostics))
                }
                self.openTimeout = timeoutWork
                self.connection.stateUpdateHandler = { [weak self] state in
                    self?.handle(state)
                }
                self.connection.start(queue: self.queue)
                self.queue.asyncAfter(deadline: .now() + timeout, execute: timeoutWork)
            }
        }
    }

    func ping() async throws {
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            queue.async {
                guard self.hasOpened, !self.closing else {
                    continuation.resume(throwing: TunnelChannelError.closed)
                    return
                }
                self.sendFrame(opcode: 0x09, payload: Data()) { error in
                    if let error { continuation.resume(throwing: error) }
                    else { continuation.resume() }
                }
            }
        }
    }

    func close() {
        queue.async {
            guard !self.closing else { return }
            self.closing = true
            self.openTimeout?.cancel()
            self.openTimeout = nil
            let continuation = self.opened
            self.opened = nil
            continuation?.resume(throwing: TunnelChannelError.closed)
            self.connection.cancel()
        }
    }

    private var connectionDiagnostics: String {
        let path = connection.currentPath.map { String(describing: $0.status) }
            ?? L10n.text("network.path_unavailable", fallback: "unavailable")
        return L10n.text(
            "network.phase",
            fallback: "network phase: %@, path: %@, TLS: %@",
            networkState,
            path,
            tlsDiagnostics.value
        )
    }

    private func handle(_ state: NWConnection.State) {
        switch state {
        case .setup:
            networkState = L10n.text("network.setup", fallback: "setup")
        case .preparing:
            networkState = L10n.text("network.preparing", fallback: "preparing TCP/TLS")
        case .ready:
            networkState = L10n.text(
                "network.websocket_upgrade",
                fallback: "TCP/TLS ready, WebSocket upgrade"
            )
            startHandshake()
        case .failed(let error):
            networkState = L10n.text(
                "network.error",
                fallback: "error: %@",
                error.localizedDescription
            )
            fail(TunnelChannelError.network(error, connectionDiagnostics))
        case .waiting(let error):
            networkState = L10n.text(
                "network.waiting",
                fallback: "waiting: %@",
                error.localizedDescription
            )
        case .cancelled:
            if !closing { fail(TunnelChannelError.closed) }
        @unknown default:
            networkState = L10n.text("network.unknown", fallback: "unknown network state")
        }
    }

    private func startHandshake() {
        guard !handshakeStarted, !closing else { return }
        handshakeStarted = true
        let request = "GET \(resource) HTTP/1.1\r\n" +
            "Host: \(hostHeader)\r\n" +
            "Upgrade: websocket\r\n" +
            "Connection: Upgrade\r\n" +
            "Sec-WebSocket-Key: \(webSocketKey)\r\n" +
            "Sec-WebSocket-Version: 13\r\n\r\n"
        connection.send(content: Data(request.utf8), completion: .contentProcessed { [weak self] error in
            guard let self else { return }
            self.queue.async {
                if let error {
                    self.fail(TunnelChannelError.network(error, self.connectionDiagnostics))
                } else {
                    self.receive()
                }
            }
        })
    }

    private func receive() {
        connection.receive(minimumIncompleteLength: 1, maximumLength: 65_535) {
            [weak self] data, _, isComplete, error in
            guard let self else { return }
            self.queue.async {
                guard !self.closing else { return }
                if let data, !data.isEmpty {
                    self.receiveBuffer.append(data)
                    if self.hasOpened { self.processFrames() }
                    else { self.processHandshake() }
                }
                guard !self.closing else { return }
                if let error {
                    self.fail(TunnelChannelError.network(error, self.connectionDiagnostics))
                } else if isComplete {
                    self.fail(TunnelChannelError.closedByPeer(nil))
                } else {
                    self.receive()
                }
            }
        }
    }

    private func processHandshake() {
        let terminator = Data("\r\n\r\n".utf8)
        guard let range = receiveBuffer.range(of: terminator) else {
            if receiveBuffer.count > Self.maximumHandshakeBytes {
                fail(TunnelChannelError.invalidHandshake(
                    L10n.text("network.http_too_large", fallback: "HTTP response too large")
                ))
            }
            return
        }
        let headerData = receiveBuffer[..<range.lowerBound]
        receiveBuffer.removeSubrange(..<range.upperBound)
        guard let header = String(data: headerData, encoding: .utf8) else {
            fail(TunnelChannelError.invalidHandshake(
                L10n.text("network.http_invalid", fallback: "invalid HTTP response")
            ))
            return
        }
        let lines = header.components(separatedBy: "\r\n")
        guard let statusLine = lines.first,
              statusLine.split(separator: " ").dropFirst().first == "101" else {
            fail(TunnelChannelError.invalidHandshake(
                lines.first ?? L10n.text("network.http_empty", fallback: "empty HTTP response")
            ))
            return
        }
        var headers: [String: String] = [:]
        for line in lines.dropFirst() {
            guard let separator = line.firstIndex(of: ":") else { continue }
            let name = line[..<separator].trimmingCharacters(in: .whitespaces).lowercased()
            let value = line[line.index(after: separator)...].trimmingCharacters(in: .whitespaces)
            headers[name] = value
        }
        guard headers["upgrade"]?.lowercased() == "websocket",
              headers["sec-websocket-accept"] == expectedAccept else {
            fail(TunnelChannelError.invalidHandshake(
                L10n.text(
                    "network.websocket_upgrade_invalid",
                    fallback: "invalid WebSocket upgrade"
                )
            ))
            return
        }
        networkState = L10n.text("network.websocket_ready", fallback: "WebSocket ready")
        completeOpen()
        if !receiveBuffer.isEmpty { processFrames() }
    }

    private func completeOpen() {
        openTimeout?.cancel()
        openTimeout = nil
        let continuation = opened
        opened = nil
        hasOpened = true
        continuation?.resume()
    }

    private func processFrames() {
        while !closing {
            let bytes = [UInt8](receiveBuffer)
            guard bytes.count >= 2 else { return }
            let first = bytes[0]
            let second = bytes[1]
            let opcode = first & 0x0f
            let isControl = opcode >= 0x08
            if isControl && first & 0x80 == 0 {
                fail(TunnelChannelError.invalidFrame(
                    L10n.text(
                        "network.fragmented_control_frame",
                        fallback: "fragmented control frame"
                    )
                ))
                return
            }

            var offset = 2
            var payloadLength = Int(second & 0x7f)
            if payloadLength == 126 {
                guard bytes.count >= offset + 2 else { return }
                payloadLength = Int(bytes[offset]) << 8 | Int(bytes[offset + 1])
                offset += 2
            } else if payloadLength == 127 {
                guard bytes.count >= offset + 8 else { return }
                var length: UInt64 = 0
                for byte in bytes[offset..<(offset + 8)] { length = (length << 8) | UInt64(byte) }
                guard length <= UInt64(Self.maximumFrameBytes) else {
                    fail(TunnelChannelError.invalidFrame(
                        L10n.text("network.frame_too_large", fallback: "frame too large")
                    ))
                    return
                }
                payloadLength = Int(length)
                offset += 8
            }
            if payloadLength > Self.maximumFrameBytes || (isControl && payloadLength > 125) {
                fail(TunnelChannelError.invalidFrame(
                    L10n.text("network.invalid_frame_size", fallback: "invalid frame size")
                ))
                return
            }

            let masked = second & 0x80 != 0
            var mask: ArraySlice<UInt8>?
            if masked {
                guard bytes.count >= offset + 4 else { return }
                mask = bytes[offset..<(offset + 4)]
                offset += 4
            }
            guard bytes.count >= offset + payloadLength else { return }
            var payload = Array(bytes[offset..<(offset + payloadLength)])
            if let mask {
                let key = Array(mask)
                for index in payload.indices { payload[index] ^= key[index % 4] }
            }
            receiveBuffer.removeFirst(offset + payloadLength)
            handleFrame(opcode: opcode, payload: Data(payload))
        }
    }

    private func handleFrame(opcode: UInt8, payload: Data) {
        switch opcode {
        case 0x08:
            let code: Int?
            if payload.count >= 2 {
                code = Int(payload[payload.startIndex]) << 8 | Int(payload[payload.startIndex + 1])
            } else {
                code = nil
            }
            fail(TunnelChannelError.closedByPeer(code))
        case 0x09:
            sendFrame(opcode: 0x0a, payload: payload, completion: nil)
        case 0x01:
            if let text = String(data: payload, encoding: .utf8) {
                onText?(text)
            } else {
                fail(TunnelChannelError.invalidFrame(
                    L10n.text("network.invalid_text_frame", fallback: "invalid text frame")
                ))
            }
        case 0x00, 0x02, 0x0a:
            break
        default:
            fail(TunnelChannelError.invalidFrame(
                L10n.text(
                    "network.unsupported_opcode",
                    fallback: "unsupported opcode %d",
                    Int(opcode)
                )
            ))
        }
    }

    private func sendFrame(opcode: UInt8, payload: Data, completion: ((Error?) -> Void)?) {
        guard !closing else {
            completion?(TunnelChannelError.closed)
            return
        }
        do {
            let frame = try Self.maskedFrame(opcode: opcode, payload: payload)
            connection.send(content: frame, completion: .contentProcessed { [weak self] error in
                guard let self else { return }
                self.queue.async {
                    if let error { self.fail(TunnelChannelError.network(error, self.connectionDiagnostics)) }
                    completion?(error)
                }
            })
        } catch {
            completion?(error)
            fail(error)
        }
    }

    private func fail(_ error: Error) {
        guard !closing else { return }
        openTimeout?.cancel()
        openTimeout = nil
        let continuation = opened
        opened = nil
        let shouldReport = hasOpened && !disconnectReported
        if shouldReport { disconnectReported = true }
        closing = true
        connection.cancel()
        continuation?.resume(throwing: error)
        if shouldReport { onDisconnect?() }
    }

    private static func maskedFrame(opcode: UInt8, payload: Data) throws -> Data {
        var frame = Data([0x80 | opcode])
        let count = payload.count
        if count <= 125 {
            frame.append(UInt8(0x80 | count))
        } else if count <= Int(UInt16.max) {
            frame.append(0x80 | 126)
            frame.append(UInt8((count >> 8) & 0xff))
            frame.append(UInt8(count & 0xff))
        } else {
            frame.append(0x80 | 127)
            let length = UInt64(count)
            for shift in stride(from: 56, through: 0, by: -8) {
                frame.append(UInt8((length >> UInt64(shift)) & 0xff))
            }
        }
        let mask = try randomBytes(count: 4)
        frame.append(mask)
        let maskBytes = [UInt8](mask)
        for (index, byte) in payload.enumerated() {
            frame.append(byte ^ maskBytes[index % 4])
        }
        return frame
    }

    private static func randomBytes(count: Int) throws -> Data {
        var bytes = [UInt8](repeating: 0, count: count)
        guard SecRandomCopyBytes(kSecRandomDefault, bytes.count, &bytes) == errSecSuccess else {
            throw TunnelChannelError.randomGeneration
        }
        return Data(bytes)
    }

    private static func tlsOptions(
        fingerprint: String?,
        caCertificatePem: String?,
        diagnostics: TLSValidationDiagnostics,
        queue: DispatchQueue
    ) -> NWProtocolTLS.Options {
        let options = NWProtocolTLS.Options()
        let expected = normalizedFingerprint(fingerprint)
        let suppliedFingerprint = fingerprint?.trimmingCharacters(in: .whitespacesAndNewlines)
        let caCertificate = certificate(from: caCertificatePem)
        guard suppliedFingerprint?.isEmpty == false || caCertificate != nil else {
            diagnostics.set(L10n.text("tls.system_trust", fallback: "system trust validation"))
            return options
        }

        sec_protocol_options_set_verify_block(options.securityProtocolOptions, { _, secTrust, complete in
            let trust = sec_trust_copy_ref(secTrust).takeRetainedValue()
            if suppliedFingerprint?.isEmpty == false {
                guard let expected, expected.count == 64 else {
                    diagnostics.set(L10n.text(
                        "tls.invalid_fingerprint",
                        fallback: "invalid received fingerprint"
                    ))
                    complete(false)
                    return
                }
                guard let chain = SecTrustCopyCertificateChain(trust) as? [SecCertificate],
                      let leaf = chain.first else {
                    diagnostics.set(L10n.text(
                        "tls.missing_certificate",
                        fallback: "TLS certificate not provided"
                    ))
                    complete(false)
                    return
                }
                let raw = SecCertificateCopyData(leaf) as Data
                let actual = SHA256.hash(data: raw).map { String(format: "%02x", $0) }.joined()
                if actual == expected {
                    diagnostics.set(L10n.text(
                        "tls.fingerprint_verified",
                        fallback: "fingerprint verified"
                    ))
                    complete(true)
                } else {
                    diagnostics.set(L10n.text(
                        "tls.fingerprint_mismatch",
                        fallback: "fingerprint mismatch"
                    ))
                    complete(false)
                }
                return
            }

            guard let caCertificate else {
                diagnostics.set(L10n.text("tls.invalid_ca", fallback: "configured CA is invalid"))
                complete(false)
                return
            }
            SecTrustSetPolicies(trust, SecPolicyCreateBasicX509())
            SecTrustSetAnchorCertificates(trust, [caCertificate] as CFArray)
            SecTrustSetAnchorCertificatesOnly(trust, true)
            var trustError: CFError?
            let accepted = SecTrustEvaluateWithError(trust, &trustError)
            diagnostics.set(accepted
                ? L10n.text("tls.ca_verified", fallback: "CA verified")
                : L10n.text("tls.ca_untrusted", fallback: "CA not trusted"))
            complete(accepted)
        }, queue)
        return options
    }

    private static func normalizedFingerprint(_ value: String?) -> String? {
        guard let value else { return nil }
        let suffix = value.split(separator: "/").last.map(String.init) ?? value
        return suffix.lowercased().filter(\.isHexDigit)
    }

    private static func certificate(from pem: String?) -> SecCertificate? {
        guard let pem,
              let data = Data(base64Encoded: pem.components(separatedBy: .newlines)
                .filter { !$0.hasPrefix("---") }.joined()) else { return nil }
        return SecCertificateCreateWithData(nil, data as CFData)
    }
}

private final class TLSValidationDiagnostics: @unchecked Sendable {
    private let lock = NSLock()
    private var current: String

    init(secure: Bool) {
        current = secure
            ? L10n.text("tls.not_run", fallback: "not run yet")
            : L10n.text("tls.not_applicable", fallback: "not applicable")
    }

    var value: String {
        lock.lock()
        defer { lock.unlock() }
        return current
    }

    func set(_ value: String) {
        lock.lock()
        current = value
        lock.unlock()
    }
}

private enum TunnelChannelError: LocalizedError {
    case missingTunnelInterface
    case invalidURL
    case randomGeneration
    case timedOut(TimeInterval, String)
    case network(Error, String)
    case invalidHandshake(String)
    case invalidFrame(String)
    case closed
    case closedByPeer(Int?)

    var errorDescription: String? {
        switch self {
        case .missingTunnelInterface:
            L10n.text("error.missing_tunnel_interface", fallback: "VPN virtual interface unavailable")
        case .invalidURL:
            L10n.text("error.invalid_runtime_address", fallback: "Invalid runtime address")
        case .randomGeneration:
            L10n.text(
                "error.websocket_initialization",
                fallback: "Could not initialize the WebSocket channel"
            )
        case .timedOut(let value, let detail):
            L10n.text(
                "error.runtime_timeout_diagnostics",
                fallback: "The runtime did not respond within %d seconds (%@)",
                Int(value),
                detail
            )
        case .network(let error, let detail):
            L10n.text(
                "error.runtime_connection_diagnostics",
                fallback: "Could not connect to the runtime: %@ (%@)",
                error.localizedDescription,
                detail
            )
        case .invalidHandshake(let detail):
            L10n.text(
                "error.runtime_rejected_websocket",
                fallback: "The runtime rejected the WebSocket channel: %@",
                detail
            )
        case .invalidFrame(let detail):
            L10n.text(
                "error.runtime_invalid_frame",
                fallback: "The runtime sent an invalid WebSocket frame: %@",
                detail
            )
        case .closed:
            L10n.text("error.runtime_channel_closed", fallback: "Runtime channel closed")
        case .closedByPeer(let code):
            code.map {
                L10n.text(
                    "error.runtime_closed_code",
                    fallback: "The runtime closed the VPN channel (%d)",
                    $0
                )
            } ?? L10n.text(
                "error.runtime_closed",
                fallback: "The runtime closed the VPN channel"
            )
        }
    }
}
