import Foundation

private final class SponsorBlockProtocol: URLProtocol, @unchecked Sendable {
    struct Reply {
        var status = 200
        var body = Data("[]".utf8)
        var headers: [String: String] = [:]
        var chunkSize = 8192
    }
    final class State: @unchecked Sendable {
        let lock = NSLock()
        var reply = Reply()
        var requests: [URLRequest] = []
        func reset(_ reply: Reply) {
            lock.lock(); defer { lock.unlock() }
            self.reply = reply
            requests = []
        }
        func receive(_ request: URLRequest) -> Reply {
            lock.lock(); defer { lock.unlock() }
            requests.append(request)
            return reply
        }
        var captured: [URLRequest] {
            lock.lock(); defer { lock.unlock() }
            return requests
        }
    }
    static let state = State()
    override class func canInit(with request: URLRequest) -> Bool { true }
    override class func canonicalRequest(for request: URLRequest) -> URLRequest { request }
    override func startLoading() {
        let reply = Self.state.receive(request)
        let response = HTTPURLResponse(url: request.url!, statusCode: reply.status,
                                       httpVersion: "HTTP/1.1", headerFields: reply.headers)!
        client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
        for offset in stride(from: 0, to: reply.body.count, by: reply.chunkSize) {
            client?.urlProtocol(self, didLoad: reply.body.subdata(in: offset..<min(offset + reply.chunkSize, reply.body.count)))
        }
        client?.urlProtocolDidFinishLoading(self)
    }
    override func stopLoading() {}
}

@main
struct SponsorBlockChecks {
    struct Failure: Error { let message: String }
    static let id = "YOelRv7fMxY"
    static func require(_ condition: @autoclosure () -> Bool, _ message: String) throws {
        if !condition() { throw Failure(message: message) }
    }
    static func reject(_ message: String, _ operation: () async throws -> Void) async throws {
        do { try await operation() }
        catch { return }
        throw Failure(message: "Accepted \(message)")
    }
    static func row(_ start: Double = 10.25, _ end: Double = 20.75, duration: Double = 100) -> [String: Any] {
        ["category": "sponsor", "actionType": "skip", "segment": [start, end], "videoDuration": duration]
    }
    static func wire(_ value: Any) throws -> Data {
        try JSONSerialization.data(withJSONObject: value, options: [.sortedKeys, .fragmentsAllowed])
    }
    static func bucket(_ rows: [Any], videoID: String = id) -> [String: Any] {
        ["videoID": videoID, "segments": rows]
    }
    static func configuration() -> URLSessionConfiguration {
        let configuration = URLSessionConfiguration.ephemeral
        configuration.protocolClasses = [SponsorBlockProtocol.self]
        // Even an injected configuration must not send ambient credentials.
        configuration.httpAdditionalHeaders = ["Authorization": "secret", "Cookie": "secret"]
        return configuration
    }
    private static func resolve(_ reply: SponsorBlockProtocol.Reply) async throws -> [SponsorSegment] {
        SponsorBlockProtocol.state.reset(reply)
        return try await SponsorBlockResolver.resolve(videoID: id, configuration: configuration())
    }
    static func checks() async throws {
        let canonical = "https://www.youtube.com/watch?v=\(id)"
        try require(SponsorBlockResolver.videoID(for: URL(string: canonical)!) == id, "canonical identity")
        for invalid in ["http://www.youtube.com/watch?v=\(id)", "https://youtube.com/watch?v=\(id)",
                        "https://youtu.be/\(id)", "https://www.youtube.com:443/watch?v=\(id)",
                        "https://user:secret@www.youtube.com/watch?v=\(id)", canonical + "&t=10",
                        canonical + "#chapter", canonical + "&v=\(id)", canonical + "x",
                        "https://www.youtube.com/watch?v=YOelRv7fMxé",
                        "https://www.youtube.com/watch?v=%59OelRv7fMxY",
                        "https://www.youtube.com/watch?v=YOelRv7fMx.",
                        "https://www.youtube.com/embed/\(id)", "https://example.org/watch?v=\(id)"] {
            try require(SponsorBlockResolver.videoID(for: URL(string: invalid)!) == nil, "reject noncanonical identity: \(invalid)")
        }
        let request = try SponsorBlockResolver.request(videoID: id)
        try require(request.url?.absoluteString == "https://sponsor.ajay.app/api/skipSegments/5db6?category=sponsor&actionType=skip&service=YouTube", "known SHA256 prefix and exact endpoint")
        try require(request.httpMethod == "GET" && !request.httpShouldHandleCookies && request.timeoutInterval == 10, "bounded credential-free GET")
        try await reject("invalid direct identity") { _ = try SponsorBlockResolver.request(videoID: "../secret") }

        var invalidRows: [Any] = [NSNull(), "bad row", 4, ["segment": [1, 2]], row(-1, 2), row(3, 2), row(2, 2)]
        for (key, value) in [("category", "intro" as Any), ("actionType", "mute" as Any),
                             ("segment", [1] as Any), ("segment", [1, 2, 3] as Any),
                             ("segment", [true, false] as Any), ("segment", ["1", "2"] as Any),
                             ("videoDuration", -1 as Any), ("videoDuration", "100" as Any),
                             ("videoDuration", true as Any), ("videoDuration", NSNull() as Any)] {
            var malformed = row(); malformed[key] = value; invalidRows.append(malformed)
        }
        var missing = row(); missing.removeValue(forKey: "videoDuration"); invalidRows.append(missing)
        let body = try wire([NSNull(), 42, ["videoID": id],
                             bucket([row(1, 2)], videoID: "aqz-KE-bpKQ"),
                             bucket(invalidRows + [row(), row(80.125, 90.5, duration: 0)])])
        let segments = try SponsorBlockResolver.decode(body, videoID: id)
        try require(segments.count == 2 && segments[0].start == 10.25 && segments[0].end == 20.75 && segments[0].videoDuration == 100 && segments[1].videoDuration == 0, "lossy individual buckets/rows, exact match, fractional times and unknown duration")
        let absent = try SponsorBlockResolver.decode(body, videoID: "abcdefghijk")
        try require(absent.isEmpty, "no exact bucket match")
        for malformed in [Data("[".utf8), try wire(["segments": []]), try wire(NSNull()), Data("[NaN]".utf8)] {
            try await reject("malformed top-level JSON") { _ = try SponsorBlockResolver.decode(malformed, videoID: id) }
        }
        let overflow = Data("[{\"videoID\":\"\(id)\",\"segments\":[{\"category\":\"sponsor\",\"actionType\":\"skip\",\"segment\":[0,1e999],\"videoDuration\":100}]}]".utf8)
        let overflowRows = try SponsorBlockResolver.decode(overflow, videoID: id)
        try require(overflowRows.isEmpty, "nonfinite row numbers are rejected without poisoning the bucket")

        func segment(_ start: Double, _ end: Double, _ reference: Double = 100) -> SponsorSegment {
            SponsorSegment(start: start, end: end, videoDuration: reference)
        }
        let ranges = SponsorSegment.visibleRanges([segment(30, 40), segment(10.25, 20.75), segment(20.75, 31), segment(12, 13)], duration: 100)
        try require(ranges == [10.25...40], "sort, union overlapping and touching ranges")
        for duration in [0, -1, Double.nan, .infinity] {
            try require(SponsorSegment.visibleRanges(segments, duration: duration).isEmpty, "unknown or invalid actual duration hides paint")
        }
        for reference in [99.0, 101.0, -1, Double.nan, .infinity] {
            try require(SponsorSegment.visibleRanges([segment(10, 20, reference)], duration: 100).isEmpty, "strict stale boundary and invalid reference duration")
        }
        for reference in [99.001, 100.999, 0] {
            try require(SponsorSegment.visibleRanges([segment(10, 20, reference)], duration: 100) == [10...20], "under-one-second mismatch or unknown reference accepted")
        }
        let invalid = [segment(-1, 2), segment(2, 2), segment(3, 2), segment(.nan, 3), segment(0, .infinity), segment(90, 101)]
        try require(SponsorSegment.visibleRanges(invalid, duration: 100).isEmpty, "never clip, rescale, or retain invalid timestamps")
        let changing = [segment(90, 100, 0)]
        try require(SponsorSegment.visibleRanges(changing, duration: 100) == [90...100] && SponsorSegment.visibleRanges(changing, duration: 99).isEmpty && SponsorSegment.visibleRanges(changing, duration: 110) == [90...100], "changing duration recomputes visibility from untouched ranges")

        let fetched = try await resolve(.init(body: body))
        try require(fetched.count == 2, "native async streamed transport decodes fixture")
        let captured = SponsorBlockProtocol.state.captured
        try require(captured.count == 1 && captured[0].url == request.url && captured[0].httpMethod == "GET" && captured[0].value(forHTTPHeaderField: "Authorization") == nil && captured[0].value(forHTTPHeaderField: "Cookie") == nil, "one fixed-host request without injected credentials")
        let notFound = try await resolve(.init(status: 404, body: Data("not JSON".utf8)))
        try require(notFound.isEmpty, "404 is normal absence")
        for status in [401, 429, 500, 302] {
            try await reject("HTTP \(status)") {
                _ = try await resolve(.init(status: status, headers: ["Location": "https://example.org/secret"]))
            }
            try require(SponsorBlockProtocol.state.captured.count == 1, "no retries or follow-up requests")
        }
        try await reject("malformed network payload") { _ = try await resolve(.init(body: Data("oops".utf8))) }
        try await reject("oversize advertised payload") { _ = try await resolve(.init(headers: ["Content-Length": "2097153"])) }
        try await reject("oversize streamed payload without length") {
            _ = try await resolve(.init(body: Data(repeating: 32, count: SponsorBlockResolver.responseLimit + 1)))
        }
        let cancelled = Task {
            withUnsafeCurrentTask { $0?.cancel() }
            return try await SponsorBlockResolver.resolve(videoID: id, configuration: configuration())
        }
        do {
            _ = try await cancelled.value
            throw Failure(message: "Cancelled lookup succeeded")
        } catch is CancellationError {} catch let failure as Failure { throw failure }
        print("SponsorBlock checks passed")
    }

    static func main() async {
        do {
            if Array(CommandLine.arguments.dropFirst()) == ["--live"] {
                for videoID in [id, "aqz-KE-bpKQ"] {
                    let segments = try await SponsorBlockResolver.resolve(videoID: videoID)
                    let intervals = segments.map { ["start": $0.start, "end": $0.end, "videoDuration": $0.videoDuration] }
                    let data = try wire(["videoID": videoID, "intervals": intervals])
                    print(String(decoding: data, as: UTF8.self))
                }
            } else if CommandLine.arguments.count == 1 {
                try await checks()
            } else {
                throw Failure(message: "Usage: sponsorblock-check [--live]")
            }
        } catch {
            let message = (error as? Failure)?.message ?? error.localizedDescription
            FileHandle.standardError.write(Data("FAIL: \(message)\n".utf8))
            exit(1)
        }
    }
}