terebi/Sources/App/SponsorBlock.swift

import Foundation
import CryptoKit

struct SponsorSegment: Sendable {
    let start: Double
    let end: Double
    let videoDuration: Double

    static func visibleRanges(_ segments: [SponsorSegment], duration: Double) -> [ClosedRange<Double>] {
        guard duration.isFinite, duration > 0 else { return [] }
        let ranges = segments.compactMap { segment -> ClosedRange<Double>? in
            guard segment.start.isFinite, segment.end.isFinite,
                  segment.start >= 0, segment.end > segment.start, segment.end <= duration,
                  segment.videoDuration.isFinite, segment.videoDuration >= 0,
                  segment.videoDuration == 0 || abs(segment.videoDuration - duration) < 1 else { return nil }
            return segment.start...segment.end
        }.sorted { $0.lowerBound < $1.lowerBound }
        var merged: [ClosedRange<Double>] = []
        for range in ranges {
            if let previous = merged.last, range.lowerBound <= previous.upperBound {
                merged[merged.count - 1] = previous.lowerBound...max(previous.upperBound, range.upperBound)
            } else {
                merged.append(range)
            }
        }
        return merged
    }
}

private enum SponsorBlockError: LocalizedError {
    case invalidVideoID, invalidResponse, status(Int), tooLarge, malformedJSON

    var errorDescription: String? {
        switch self {
        case .invalidVideoID: return "SponsorBlock requires a canonical YouTube video ID."
        case .invalidResponse: return "SponsorBlock returned an invalid response."
        case .status(let status): return "SponsorBlock returned HTTP \(status)."
        case .tooLarge: return "SponsorBlock returned too much data."
        case .malformedJSON: return "SponsorBlock returned malformed segment data."
        }
    }
}

private final class SponsorBlockRedirectPolicy: NSObject, URLSessionTaskDelegate, Sendable {
    func urlSession(_ session: URLSession, task: URLSessionTask,
                    willPerformHTTPRedirection response: HTTPURLResponse,
                    newRequest request: URLRequest,
                    completionHandler: @escaping (URLRequest?) -> Void) {
        completionHandler(nil)
    }
}

enum SponsorBlockResolver {
    static let responseLimit = 2 * 1024 * 1024

    private static func validID(_ id: String) -> Bool {
        id.utf8.count == 11 && id.utf8.allSatisfy {
            (65...90).contains($0) || (97...122).contains($0) || (48...57).contains($0) || $0 == 45 || $0 == 95
        }
    }

    static func videoID(for url: URL) -> String? {
        let prefix = "https://www.youtube.com/watch?v="
        let text = url.absoluteString
        guard text.hasPrefix(prefix) else { return nil }
        let id = String(text.dropFirst(prefix.count))
        return validID(id) ? id : nil
    }

    static func request(videoID: String) throws -> URLRequest {
        guard validID(videoID) else { throw SponsorBlockError.invalidVideoID }
        let prefix = SHA256.hash(data: Data(videoID.utf8)).prefix(2).map { String(format: "%02x", $0) }.joined()
        let url = URL(string: "https://sponsor.ajay.app/api/skipSegments/\(prefix)?category=sponsor&actionType=skip&service=YouTube")!
        var request = URLRequest(url: url, cachePolicy: .reloadIgnoringLocalAndRemoteCacheData, timeoutInterval: 10)
        request.httpMethod = "GET"
        request.httpShouldHandleCookies = false
        return request
    }

    static func resolve(videoID: String, configuration: URLSessionConfiguration? = nil) async throws -> [SponsorSegment] {
        let request = try request(videoID: videoID)
        let configuration = (configuration?.copy() as? URLSessionConfiguration) ?? .ephemeral
        configuration.httpCookieStorage = nil
        configuration.httpShouldSetCookies = false
        configuration.urlCredentialStorage = nil
        configuration.urlCache = nil
        configuration.requestCachePolicy = .reloadIgnoringLocalAndRemoteCacheData
        configuration.httpAdditionalHeaders = nil
        configuration.timeoutIntervalForRequest = 10
        configuration.timeoutIntervalForResource = 15
        let session = URLSession(configuration: configuration)
        defer { session.invalidateAndCancel() }
        try Task.checkCancellation()
        let (bytes, response) = try await session.bytes(for: request, delegate: SponsorBlockRedirectPolicy())
        try Task.checkCancellation()
        guard let response = response as? HTTPURLResponse else { throw SponsorBlockError.invalidResponse }
        if response.statusCode == 404 { return [] }
        guard response.statusCode == 200 else { throw SponsorBlockError.status(response.statusCode) }
        guard response.expectedContentLength <= Int64(responseLimit) else { throw SponsorBlockError.tooLarge }
        var data = Data()
        data.reserveCapacity(min(responseLimit, max(0, Int(response.expectedContentLength))))
        for try await byte in bytes {
            try Task.checkCancellation()
            guard data.count < responseLimit else { throw SponsorBlockError.tooLarge }
            data.append(byte)
        }
        try Task.checkCancellation()
        return try decode(data, videoID: videoID)
    }

    static func decode(_ data: Data, videoID: String) throws -> [SponsorSegment] {
        guard data.count <= responseLimit else { throw SponsorBlockError.tooLarge }
        let buckets: [Lossy<Bucket>]
        do { buckets = try JSONDecoder().decode([Lossy<Bucket>].self, from: data) }
        catch { throw SponsorBlockError.malformedJSON }
        return buckets.compactMap(\.value).filter { $0.videoID == videoID }.flatMap { bucket in
            bucket.segments.compactMap { row in
                guard let row = row.value, row.category == "sponsor", row.actionType == "skip",
                      row.segment.count == 2,
                      row.segment[0].isFinite, row.segment[1].isFinite,
                      row.segment[0] >= 0, row.segment[1] > row.segment[0],
                      row.videoDuration.isFinite, row.videoDuration >= 0 else { return nil }
                return SponsorSegment(start: row.segment[0], end: row.segment[1], videoDuration: row.videoDuration)
            }
        }
    }

    private struct Lossy<Value: Decodable>: Decodable {
        let value: Value?
        init(from decoder: Decoder) throws { value = try? Value(from: decoder) }
    }

    private struct Bucket: Decodable {
        let videoID: String
        let segments: [Lossy<Row>]
    }

    private struct Row: Decodable {
        let category: String
        let actionType: String
        let segment: [Double]
        let videoDuration: Double
    }
}