stremio/apps/ios/VMC/Playback/ClearKeyDashProxy.swift

import Foundation
import Network
import os.log

private let proxyLog = Logger(subsystem: "nl.vonas.mediacluster.ios", category: "ClearKeyProxy")

/// Lokale HTTP-proxy: haalt DASH MPD/segmenten op, strip DRM-metadata, decrypt CENC met ClearKey.
/// VLC speelt daarna plain DASH vanaf 127.0.0.1.
final class ClearKeyDashProxy {
    static let shared = ClearKeyDashProxy()

    private let queue = DispatchQueue(label: "nl.vonas.vmc.clearkey-proxy")
    private var listener: NWListener?
    private var port: UInt16 = 0
    private var remoteManifest: URL?
    private var keyMap: ClearKeyKeyMap?
    private var cryptoState = CencMp4Decryptor.SessionState()
    private var generation = 0
    private let urlSession: URLSession = {
        let cfg = URLSessionConfiguration.ephemeral
        cfg.timeoutIntervalForRequest = 25
        cfg.timeoutIntervalForResource = 60
        cfg.httpMaximumConnectionsPerHost = 8
        cfg.requestCachePolicy = .reloadIgnoringLocalCacheData
        return URLSession(configuration: cfg)
    }()

    private init() {}

    var isRunning: Bool { listener != nil && port > 0 }

    /// Start (of herstart) proxy voor dit manifest + keys. Geeft lokale MPD-URL terug.
    /// Warmt de MPD-fetch op zodat VLC niet meteen een lege/fout-response krijgt.
    func start(manifestURL: URL, keys: [ClearKeyEntry]) throws -> URL {
        let map = ClearKeyKeyMap(entries: keys)
        guard !map.isEmpty else {
            throw ProxyError.noKeys
        }
        stop()
        remoteManifest = manifestURL
        keyMap = map
        cryptoState = CencMp4Decryptor.SessionState()
        generation += 1
        let gen = generation

        let params = NWParameters.tcp
        params.allowLocalEndpointReuse = true
        let listener = try NWListener(using: params, on: .any)
        self.listener = listener

        let ready = DispatchSemaphore(value: 0)
        var readyError: Error?

        listener.stateUpdateHandler = { [weak self] state in
            guard let self else { return }
            switch state {
            case .ready:
                if let p = listener.port?.rawValue {
                    self.port = p
                }
                proxyLog.info("listener ready on port \(self.port, privacy: .public)")
                ready.signal()
            case .failed(let err):
                proxyLog.error("listener failed: \(String(describing: err), privacy: .public)")
                readyError = err
                ready.signal()
            default:
                break
            }
        }

        listener.newConnectionHandler = { [weak self] conn in
            self?.handle(connection: conn)
        }

        listener.start(queue: queue)
        _ = ready.wait(timeout: .now() + 5)
        if let readyError { throw readyError }
        guard port > 0 else { throw ProxyError.bindFailed }

        // Warm-up: MPD ophalen + herschrijven vóór VLC start (live dynamic MPD's falen anders meteen).
        let warm = DispatchSemaphore(value: 0)
        var warmError: Error?
        Task {
            do {
                let (body, _) = try await self.payload(for: "/manifest.mpd")
                guard !body.isEmpty, gen == self.generation else {
                    throw ProxyError.emptyManifest
                }
                let preview = String(data: body.prefix(200), encoding: .utf8) ?? ""
                proxyLog.info("warmup MPD ok bytes=\(body.count, privacy: .public) head=\(preview, privacy: .public)")
            } catch {
                warmError = error
                proxyLog.error("warmup failed: \(error.localizedDescription, privacy: .public)")
            }
            warm.signal()
        }
        _ = warm.wait(timeout: .now() + 20)
        if let warmError { throw warmError }

        return URL(string: "http://127.0.0.1:\(port)/manifest.mpd")!
    }

    func stop() {
        generation += 1
        listener?.cancel()
        listener = nil
        port = 0
        remoteManifest = nil
        keyMap = nil
        cryptoState = CencMp4Decryptor.SessionState()
        proxyLog.info("proxy stopped")
    }

    // MARK: - HTTP

    private func handle(connection: NWConnection) {
        connection.start(queue: queue)
        receiveRequest(on: connection, buffer: Data())
    }

    private func receiveRequest(on connection: NWConnection, buffer: Data) {
        connection.receive(minimumIncompleteLength: 1, maximumLength: 64 * 1024) { [weak self] data, _, isComplete, error in
            guard let self else {
                connection.cancel()
                return
            }
            if error != nil {
                connection.cancel()
                return
            }
            var buf = buffer
            if let data { buf.append(data) }
            if let range = buf.range(of: Data("\r\n\r\n".utf8)) {
                let headerData = buf.subdata(in: 0..<range.lowerBound)
                self.serve(connection: connection, headerData: headerData)
                return
            }
            if isComplete {
                connection.cancel()
                return
            }
            if buf.count > 256 * 1024 {
                connection.cancel()
                return
            }
            self.receiveRequest(on: connection, buffer: buf)
        }
    }

    private func serve(connection: NWConnection, headerData: Data) {
        guard let header = String(data: headerData, encoding: .utf8),
              let firstLine = header.split(separator: "\r\n", omittingEmptySubsequences: false).first
        else {
            respond(connection, status: 400, contentType: "text/plain", body: Data("bad request".utf8))
            return
        }
        let parts = firstLine.split(separator: " ")
        guard parts.count >= 2 else {
            respond(connection, status: 400, contentType: "text/plain", body: Data("bad request".utf8))
            return
        }
        let method = String(parts[0])
        let pathAndQuery = String(parts[1])
        guard method == "GET" || method == "HEAD" else {
            respond(connection, status: 405, contentType: "text/plain", body: Data("method".utf8))
            return
        }

        let rangeHeader = Self.parseHeader(header, name: "Range")

        Task {
            do {
                let (body, contentType) = try await self.payload(for: pathAndQuery)
                if let rangeHeader, let byteRange = Self.parseByteRange(rangeHeader, total: body.count) {
                    let slice = body.subdata(in: byteRange.start..<(byteRange.end + 1))
                    if method == "HEAD" {
                        self.respond(
                            connection,
                            status: 206,
                            contentType: contentType,
                            body: Data(),
                            headOnly: true,
                            contentLength: slice.count,
                            extraHeaders: [
                                "Accept-Ranges: bytes",
                                "Content-Range: bytes \(byteRange.start)-\(byteRange.end)/\(body.count)",
                            ]
                        )
                    } else {
                        self.respond(
                            connection,
                            status: 206,
                            contentType: contentType,
                            body: slice,
                            extraHeaders: [
                                "Accept-Ranges: bytes",
                                "Content-Range: bytes \(byteRange.start)-\(byteRange.end)/\(body.count)",
                            ]
                        )
                    }
                } else if method == "HEAD" {
                    self.respond(
                        connection,
                        status: 200,
                        contentType: contentType,
                        body: Data(),
                        headOnly: true,
                        contentLength: body.count,
                        extraHeaders: ["Accept-Ranges: bytes"]
                    )
                } else {
                    self.respond(
                        connection,
                        status: 200,
                        contentType: contentType,
                        body: body,
                        extraHeaders: ["Accept-Ranges: bytes"]
                    )
                }
            } catch {
                let msg = error.localizedDescription
                proxyLog.error("serve \(pathAndQuery, privacy: .public) → \(msg, privacy: .public)")
                self.respond(connection, status: 502, contentType: "text/plain; charset=utf-8", body: Data(msg.utf8))
            }
        }
    }

    private func payload(for pathAndQuery: String) async throws -> (Data, String) {
        guard let keyMap, let remoteManifest else {
            throw ProxyError.notStarted
        }
        let split = pathAndQuery.split(separator: "?", maxSplits: 1).map(String.init)
        let pathOnly = split.first ?? pathAndQuery
        let query = split.count > 1 ? split[1] : nil

        if pathOnly == "/" || pathOnly == "/manifest.mpd" {
            let raw = try await fetch(remoteManifest)
            guard let xml = String(data: raw, encoding: .utf8), !xml.isEmpty else {
                // Sommige CDN's sturen UTF-16 BOM; probeer lossy.
                let lossy = String(decoding: raw, as: UTF8.self)
                guard !lossy.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
                    throw ProxyError.emptyManifest
                }
                let rewritten = rewriteMpd(lossy, mpdURL: remoteManifest)
                return (Data(rewritten.utf8), "application/dash+xml")
            }
            let rewritten = rewriteMpd(xml, mpdURL: remoteManifest)
            if !rewritten.localizedCaseInsensitiveContains("<MPD") {
                throw ProxyError.emptyManifest
            }
            return (Data(rewritten.utf8), "application/dash+xml")
        }

        if pathOnly.hasPrefix("/a/") {
            let b64 = String(pathOnly.dropFirst(3))
            guard var url = decodeURL(b64) else { throw ProxyError.badPath }
            if let query, url.query == nil {
                var comps = URLComponents(url: url, resolvingAgainstBaseURL: false)
                comps?.percentEncodedQuery = query
                if let u = comps?.url { url = u }
            }
            return try await fetchAndDecrypt(url, keys: keyMap)
        }

        if pathOnly.hasPrefix("/p/") {
            let rest = String(pathOnly.dropFirst(3))
            guard let slash = rest.firstIndex(of: "/") else { throw ProxyError.badPath }
            let b64 = String(rest[..<slash])
            var rel = String(rest[rest.index(after: slash)...])
            if let query {
                rel += (rel.contains("?") ? "&" : "?") + query
            }
            guard let base = decodeURL(b64) else { throw ProxyError.badPath }
            let joined = join(base: asDirectory(base), relative: rel.removingPercentEncoding ?? rel)
            return try await fetchAndDecrypt(joined, keys: keyMap)
        }

        throw ProxyError.badPath
    }

    private func fetchAndDecrypt(_ url: URL, keys: ClearKeyKeyMap) async throws -> (Data, String) {
        var data = try await fetch(url)
        let lower = url.absoluteString.lowercased()
        let isManifest = lower.contains(".mpd")
        if isManifest, let xml = String(data: data, encoding: .utf8) ?? Optional(String(decoding: data, as: UTF8.self)) {
            let rewritten = rewriteMpd(xml, mpdURL: url)
            return (Data(rewritten.utf8), "application/dash+xml")
        }
        queue.sync {
            _ = CencMp4Decryptor.process(data: &data, keys: keys, state: &cryptoState)
        }
        let ct: String
        if lower.contains(".m4s") || lower.contains(".mp4") || lower.contains("init") || lower.contains(".cmfv") || lower.contains(".cmfa") {
            ct = "video/mp4"
        } else {
            ct = "application/octet-stream"
        }
        return (data, ct)
    }

    private func fetch(_ url: URL) async throws -> Data {
        var req = URLRequest(url: url)
        req.timeoutInterval = 25
        req.cachePolicy = .reloadIgnoringLocalCacheData
        req.setValue("VMC-iOS-ClearKey/0.5.2", forHTTPHeaderField: "User-Agent")
        req.setValue("bytes", forHTTPHeaderField: "Accept-Ranges")
        let (data, resp) = try await urlSession.data(for: req)
        if let http = resp as? HTTPURLResponse, !(200...299).contains(http.statusCode) {
            throw ProxyError.upstream(http.statusCode)
        }
        return data
    }

    private func rewriteMpd(_ xml: String, mpdURL: URL) -> String {
        var s = xml
        // Strip ContentProtection (ClearKey/Widevine/PlayReady/mp4protection).
        s = s.replacingOccurrences(
            of: #"<ContentProtection\b[^>]*/>"#,
            with: "",
            options: .regularExpression
        )
        s = s.replacingOccurrences(
            of: #"<ContentProtection\b[^>]*>[\s\S]*?</ContentProtection>"#,
            with: "",
            options: .regularExpression
        )

        // Live dynamic: Location mag niet naar remote DRM-MPD wijzen (VLC omzeilt proxy).
        s = s.replacingOccurrences(
            of: #"<Location\b[^>]*>[\s\S]*?</Location>"#,
            with: "<Location>http://127.0.0.1:\(port)/manifest.mpd</Location>",
            options: .regularExpression
        )

        // Herschrijf elke BaseURL naar de juiste resolved CDN-base via /p/<b64>/.
        if let baseRegex = try? NSRegularExpression(pattern: #"<BaseURL>([\s\S]*?)</BaseURL>"#, options: []) {
            let ns = s as NSString
            let matches = baseRegex.matches(in: s, options: [], range: NSRange(location: 0, length: ns.length))
            for match in matches.reversed() {
                guard match.numberOfRanges >= 2 else { continue }
                let content = ns.substring(with: match.range(at: 1))
                    .trimmingCharacters(in: .whitespacesAndNewlines)
                if content.contains("127.0.0.1:\(port)") {
                    continue
                }
                let resolved = resolveBaseURL(content, mpdURL: mpdURL)
                let proxyBase = "http://127.0.0.1:\(port)/p/\(encodeURL(asDirectory(resolved)))/"
                s = (s as NSString).replacingCharacters(in: match.range, with: "<BaseURL>\(proxyBase)</BaseURL>")
            }
        }

        if !s.localizedCaseInsensitiveContains("<BaseURL") {
            let proxyBase = "http://127.0.0.1:\(port)/p/\(encodeURL(asDirectory(mpdURL.deletingLastPathComponent())))/"
            if let regex = try? NSRegularExpression(pattern: #"<MPD\b[^>]*>"#, options: []) {
                let range = NSRange(s.startIndex..<s.endIndex, in: s)
                if let match = regex.firstMatch(in: s, options: [], range: range),
                   let r = Range(match.range, in: s) {
                    s.insert(contentsOf: "\n  <BaseURL>\(proxyBase)</BaseURL>", at: r.upperBound)
                }
            }
        }

        // Absolute http(s) URL's in attributen → /a/<b64>
        if let urlRegex = try? NSRegularExpression(pattern: #"https?://[^\s"'<>]+"#, options: []) {
            let ns = s as NSString
            let matches = urlRegex.matches(in: s, options: [], range: NSRange(location: 0, length: ns.length))
            for match in matches.reversed() {
                let abs = ns.substring(with: match.range)
                if abs.contains("127.0.0.1:\(port)") { continue }
                let replacement = "http://127.0.0.1:\(port)/a/\(encodeURLString(abs))"
                s = (s as NSString).replacingCharacters(in: match.range, with: replacement)
            }
        }

        return s
    }

    /// Resolve BaseURL tekst t.o.v. MPD-URL; altijd als directory-URL.
    private func resolveBaseURL(_ raw: String, mpdURL: URL) -> URL {
        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        if trimmed.isEmpty {
            return asDirectory(mpdURL.deletingLastPathComponent())
        }
        if trimmed.hasPrefix("http://") || trimmed.hasPrefix("https://"),
           let abs = URL(string: trimmed) {
            return asDirectory(abs)
        }
        if let joined = URL(string: trimmed, relativeTo: mpdURL)?.absoluteURL {
            return asDirectory(joined)
        }
        return asDirectory(mpdURL.deletingLastPathComponent())
    }

    private func asDirectory(_ url: URL) -> URL {
        var s = url.absoluteString
        if s.isEmpty { return url }
        if !s.hasSuffix("/") {
            // Als laatste path-component op een bestand lijkt (.mpd/.mp4), strip; anders slash toevoegen.
            let last = url.lastPathComponent.lowercased()
            if last.contains("."), !last.hasSuffix("/") {
                return asDirectory(url.deletingLastPathComponent())
            }
            s += "/"
        }
        return URL(string: s) ?? url
    }

    private func join(base: URL, relative: String) -> URL {
        if relative.hasPrefix("http://") || relative.hasPrefix("https://") {
            return URL(string: relative) ?? base
        }
        return URL(string: relative, relativeTo: base)?.absoluteURL ?? base.appendingPathComponent(relative)
    }

    private func encodeURL(_ url: URL) -> String {
        encodeURLString(url.absoluteString)
    }

    private func encodeURLString(_ s: String) -> String {
        Data(s.utf8).base64EncodedString()
            .replacingOccurrences(of: "+", with: "-")
            .replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
    }

    private func decodeURL(_ b64url: String) -> URL? {
        var b64 = b64url
            .replacingOccurrences(of: "-", with: "+")
            .replacingOccurrences(of: "_", with: "/")
        let pad = (4 - b64.count % 4) % 4
        if pad > 0 { b64 += String(repeating: "=", count: pad) }
        guard let data = Data(base64Encoded: b64),
              let s = String(data: data, encoding: .utf8)
        else { return nil }
        return URL(string: s)
    }

    private static func parseHeader(_ header: String, name: String) -> String? {
        let lines = header.split(separator: "\r\n")
        let prefix = name.lowercased() + ":"
        for line in lines.dropFirst() {
            let l = line.trimmingCharacters(in: .whitespaces)
            if l.lowercased().hasPrefix(prefix) {
                return String(l.dropFirst(prefix.count)).trimmingCharacters(in: .whitespaces)
            }
        }
        return nil
    }

    private static func parseByteRange(_ value: String, total: Int) -> (start: Int, end: Int)? {
        // bytes=START-END | bytes=START-
        guard value.lowercased().hasPrefix("bytes=") else { return nil }
        let spec = value.dropFirst(6)
        let parts = spec.split(separator: "-", maxSplits: 1).map(String.init)
        guard let start = Int(parts.first ?? "") else { return nil }
        let end: Int
        if parts.count > 1, let e = Int(parts[1]) {
            end = min(e, total - 1)
        } else {
            end = total - 1
        }
        guard total > 0, start >= 0, start <= end, end < total else { return nil }
        return (start, end)
    }

    private func respond(
        _ connection: NWConnection,
        status: Int,
        contentType: String,
        body: Data,
        headOnly: Bool = false,
        contentLength: Int? = nil,
        extraHeaders: [String] = []
    ) {
        let reason: String
        switch status {
        case 200: reason = "OK"
        case 206: reason = "Partial Content"
        case 400: reason = "Bad Request"
        case 405: reason = "Method Not Allowed"
        case 502: reason = "Bad Gateway"
        default: reason = "Error"
        }
        let len = contentLength ?? body.count
        var header = "HTTP/1.1 \(status) \(reason)\r\n"
        header += "Content-Type: \(contentType)\r\n"
        header += "Content-Length: \(len)\r\n"
        header += "Connection: close\r\n"
        header += "Access-Control-Allow-Origin: *\r\n"
        header += "Cache-Control: no-store\r\n"
        for h in extraHeaders {
            header += h + "\r\n"
        }
        header += "\r\n"
        var payload = Data(header.utf8)
        if !headOnly { payload.append(body) }
        connection.send(content: payload, completion: .contentProcessed { _ in
            connection.cancel()
        })
    }

    enum ProxyError: LocalizedError {
        case noKeys
        case bindFailed
        case notStarted
        case badPath
        case emptyManifest
        case upstream(Int)

        var errorDescription: String? {
            switch self {
            case .noKeys: return "Geen geldige ClearKey-sleutels"
            case .bindFailed: return "Kon lokale ClearKey-proxy niet starten"
            case .notStarted: return "ClearKey-proxy is niet actief"
            case .badPath: return "Ongeldig proxy-pad"
            case .emptyManifest: return "Leeg of ongeldig DASH-manifest ontvangen"
            case .upstream(let code): return "Upstream HTTP \(code)"
            }
        }
    }
}