import AVFoundation
import Foundation
import os

private let logger = Logger(subsystem: "dev.shane.voice", category: "media")

/// Proxies {base}/media/* bytes to AVPlayer over the pinned-TLS session, since
/// AVPlayer cannot trust the self-signed LAN cert. AVPlayer loads https natively
/// (ignoring registered URLProtocol subclasses), so the asset uses a custom
/// `pinned-media` scheme and this delegate answers every data request by
/// fetching the real https URL through TLSPinning.session(), forwarding Range.
///
/// **Lifecycle invariants (Finding 11 / Iteration 06 Item 7):**
/// - Tasks are keyed by `ObjectIdentifier(loadingRequest)` directly, not a
///   derived hash, so two distinct requests can never collide.
/// - All mutable state (`loaderURLs`, `activeTasks`) is accessed only on the
///   main queue: the AVAssetResourceLoader delegate is set with `queue: .main`
///   and the pinned URLSession uses `delegateQueue: .main`, so both the
///   resource-loading callbacks and the data-task completion handlers run on
///   the same serialization domain. This is enforced in code, not by comment.
/// - For HTTP 206, `contentLength` is set only when `Content-Range` provides a
///   known total; it never falls back to the range's `expectedContentLength`.
/// - For HTTP 200, `expectedContentLength` is already the full resource size.
/// - Cleanup (task removal) occurs on success, error, cancellation, and item
///   replacement via `removeLoader`.
final class PinnedMediaResourceLoader: NSObject, AVAssetResourceLoaderDelegate {

    static let scheme = "pinned-media"

    private let base: URL
    /// Maps each asset's resource loader to the real https URL it proxies.
    /// Accessed only on the main queue (the delegate queue).
    private var loaderURLs: [ObjectIdentifier: URL] = [:]
    /// Maps each loading request's identity to its upstream data task so
    /// cancellation can cancel the in-flight download. Keyed directly by
    /// `ObjectIdentifier(loadingRequest)` — no hash derivation, no collisions.
    private var activeTasks: [ObjectIdentifier: URLSessionDataTask] = [:]
    /// Shared pinned session, created once and invalidated on deinit.
    /// Uses `delegateQueue: .main` so all completion handlers (and therefore
    /// all mutations of `activeTasks`) run on the main queue — the same
    /// serialization domain as the AVAssetResourceLoader delegate callbacks
    /// (which are also dispatched to `.main`). This makes map access truly
    /// serialized in code, not just by convention.
    private lazy var pinnedSession: URLSession = {
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 30
        return URLSession(configuration: config, delegate: PinnedTrustDelegate(), delegateQueue: .main)
    }()

    init(base: URL) {
        self.base = base
    }

    deinit {
        pinnedSession.invalidateAndCancel()
    }

    /// Maps a media URL (absolute https or /media/... path) to the pinned-media
    /// scheme so AVPlayer routes it through this delegate instead of loading
    /// https natively. Built as a string because URLComponents.url returns nil
    /// for schemes not in Foundation's known-scheme list.
    func rewrite(_ url: URL) -> URL {
        let absolute: String
        if let scheme = url.scheme?.lowercased(), scheme.hasPrefix("http") {
            var c = URLComponents(url: url, resolvingAgainstBaseURL: false)!
            c.scheme = Self.scheme
            return c.url!
        } else {
            let path = url.absoluteString.hasPrefix("/") ? String(url.absoluteString.dropFirst()) : url.absoluteString
            guard let host = base.host, !host.isEmpty else { return url }
            let port = base.port.map { ":\($0)" } ?? ""
            absolute = "\(Self.scheme)://\(host)\(port)/\(path)"
        }
        return URL(string: absolute) ?? url
    }

    /// Maps a pinned-media URL back to the real https URL it proxies.
    func resolve(_ url: URL) -> URL {
        guard let scheme = url.scheme?.lowercased(), scheme == Self.scheme else { return url }
        var c = URLComponents(url: url, resolvingAgainstBaseURL: false)!
        c.scheme = "https"
        return c.url ?? url
    }

    /// Creates an AVURLAsset that routes through this loader's delegate. The
    /// mapping from the asset's resource loader to the real URL is recorded so
    /// `resourceLoader(_:shouldWaitForLoadingOfRequestedResource:)` can resolve
    /// it. All mutations happen on the main queue (the delegate queue).
    func assetWithPinnedURL(_ pinned: URL) -> AVURLAsset {
        let asset = AVURLAsset(url: pinned)
        let realURL = resolve(pinned)
        // The delegate queue is .main, so this runs on the main thread.
        loaderURLs[ObjectIdentifier(asset.resourceLoader)] = realURL
        asset.resourceLoader.setDelegate(self, queue: .main)
        return asset
    }

    /// Removes the loader mapping and any in-flight tasks for an asset that is
    /// no longer needed. Called by MediaPlayerController when replacing or
    /// clearing media items to bound the dictionaries (Finding 11).
    func removeLoader(_ asset: AVURLAsset) {
        let key = ObjectIdentifier(asset.resourceLoader)
        loaderURLs[key] = nil
        // Cancel any in-flight tasks. URLSessionDataTask has no isCancelled
        // property; calling cancel() on an already-cancelled task is a no-op.
        for (_, task) in activeTasks {
            task.cancel()
        }
        activeTasks.removeAll()
    }

    func resourceLoader(
        _ resourceLoader: AVAssetResourceLoader,
        shouldWaitForLoadingOfRequestedResource loadingRequest: AVAssetResourceLoadingRequest
    ) -> Bool {
        // All state access is on the main queue (the delegate queue).
        guard let dataRequest = loadingRequest.dataRequest else { return false }
        guard let url = loaderURLs[ObjectIdentifier(resourceLoader)] else {
            logger.log("no resolved URL for resource loader")
            loadingRequest.finishLoading(with: URLError(.badURL))
            return false
        }

        var upstream = URLRequest(url: url)
        let offset = dataRequest.requestedOffset
        let length = Int64(dataRequest.requestedLength)
        if length > 0 {
            upstream.setValue("bytes=\(offset)-\(offset + length - 1)", forHTTPHeaderField: "Range")
        } else {
            upstream.setValue("bytes=\(offset)-", forHTTPHeaderField: "Range")
        }
        logger.log("fetching \(url.absoluteString, privacy: .public) range=\(upstream.value(forHTTPHeaderField: "Range") ?? "-", privacy: .public)")

        // Key the task by the loading request's identity directly — no hash
        // derivation, so two distinct requests can never collide.
        let requestKey = ObjectIdentifier(loadingRequest)
        let task = pinnedSession.dataTask(with: upstream) { [weak self] data, response, error in
            guard let self else { return }
            // Cleanup: remove the task from the active map regardless of outcome.
            // Runs on .main (the session's delegateQueue), same domain as all
            // other map mutations.
            self.activeTasks[requestKey] = nil

            if let error {
                logger.log("fetch failed \(url.absoluteString, privacy: .public): \(error.localizedDescription, privacy: .public)")
                loadingRequest.finishLoading(with: error)
                return
            }
            guard let http = response as? HTTPURLResponse else {
                loadingRequest.finishLoading(with: URLError(.badServerResponse))
                return
            }
            // Reject non-200/206 responses (404, 500, etc.) instead of feeding
            // error JSON into AVPlayer as media bytes.
            guard http.statusCode == 200 || http.statusCode == 206 else {
                logger.log("unexpected status \(http.statusCode) for \(url.absoluteString, privacy: .public)")
                loadingRequest.finishLoading(with: URLError(.badServerResponse))
                return
            }
            if let info = loadingRequest.contentInformationRequest {
                info.contentType = http.mimeType ?? "application/octet-stream"
                // Content-length handling (Iteration 05 Item 7):
                // - For 206: set full content length ONLY when Content-Range
                //   provides a known total. Never fall back to the range's
                //   expectedContentLength (that would tell AVPlayer the total
                //   resource is only as large as this one range).
                // - For 200: expectedContentLength is already the full size.
                if http.statusCode == 206 {
                    if let total = Self.parseTotalLength(http), total >= 0 {
                        info.contentLength = total
                    }
                    // If Content-Range has no known total (e.g. "bytes 0-1023/*"),
                    // do NOT set contentLength — AVPlayer will handle it as
                    // unknown-length, which is correct for a range response.
                } else {
                    // 200: expectedContentLength is the full resource size.
                    if http.expectedContentLength >= 0 {
                        info.contentLength = Int64(http.expectedContentLength)
                    }
                }
                info.isByteRangeAccessSupported = true
            }
            if let data {
                dataRequest.respond(with: data)
            }
            loadingRequest.finishLoading()
        }
        activeTasks[requestKey] = task
        task.resume()
        return true
    }

    func resourceLoader(
        _ resourceLoader: AVAssetResourceLoader,
        didCancel loadingRequest: AVAssetResourceLoadingRequest
    ) {
        // Cancel the upstream data task so a seek or item swap doesn't keep
        // downloading cancelled ranges. Keyed by ObjectIdentifier directly.
        let key = ObjectIdentifier(loadingRequest)
        activeTasks[key]?.cancel()
        activeTasks[key] = nil
        logger.log("resource loader cancelled request, upstream task cancelled")
    }

    /// Parses the total resource length from an HTTP 206 response's
    /// `Content-Range` header (e.g. "bytes 0-1023/4096" → 4096). Returns nil
    /// for non-206 responses, missing headers, unknown totals ("*"), or
    /// malformed values.
    static func parseTotalLength(_ response: HTTPURLResponse) -> Int64? {
        if response.statusCode == 206,
           let contentRange = response.value(forHTTPHeaderField: "Content-Range") {
            // Format: "bytes <start>-<end>/<total>" or "bytes <start>-<end>/*"
            let afterSlash = contentRange.split(separator: "/").last?.trimmingCharacters(in: .whitespaces)
            if let afterSlash, afterSlash != "*", let total = Int64(afterSlash) {
                return total
            }
        }
        return nil
    }
}