//
// Copyright (c) Double Symmetry GmbH
// Commercial use requires a license. See https://rntp.dev/pricing
//

import Foundation

/// Rewrites the child URIs in an HLS playlist so AVPlayer routes segment and
/// variant requests back through the cache proxy.
///
/// Every child URI (segments, variant/media playlists, encryption keys, init
/// maps, etc.) is resolved against the playlist's own absolute URL per
/// RFC 3986, then mapped through `transform` (in production: a proxy-URL
/// builder). Tag lines that carry a `URI="..."` attribute (per RFC 8216 §4.2)
/// have only that quoted value rewritten; everything else on the line is
/// preserved byte-for-byte, including blank lines, comments, ordering, and
/// unrelated attributes. FairPlay (`skd://`) key URIs and `data:` URIs are left
/// untouched.
///
/// The rewriter is pure by contract and performs no idempotence guard: callers
/// must never feed already-rewritten output back in, as proxy URLs are absolute
/// `http(s)` URIs and would be proxied a second time.
enum HLSManifestRewriter {

  // MARK: - Rewrite

  /// Resolve and proxy every child URI in the playlist, preserving structure.
  ///
  /// - Parameters:
  ///   - manifest: The raw playlist text (may contain CRLF or LF endings).
  ///   - baseURL: The playlist's own absolute URL; relative URIs resolve here.
  ///   - transform: Maps a resolved absolute upstream URL to its proxy URL.
  /// - Returns: The rewritten playlist, with lines joined by `\n`.
  static func rewrite(manifest: String, baseURL: URL, transform: (URL) -> URL) -> String {
    // Split on "\n" and strip a trailing "\r" from each line so CRLF input is
    // tolerated. Output is always "\n"-joined.
    let lines = manifest.components(separatedBy: "\n").map { line -> String in
      line.hasSuffix("\r") ? String(line.dropLast()) : line
    }

    let rewritten = lines.map { line -> String in
      rewriteLine(line, baseURL: baseURL, transform: transform)
    }

    return rewritten.joined(separator: "\n")
  }

  /// Rewrite a single playlist line, leaving it untouched when it carries no
  /// rewritable URI or when resolution would fail.
  private static func rewriteLine(
    _ line: String,
    baseURL: URL,
    transform: (URL) -> URL
  ) -> String {
    let trimmed = line.trimmingCharacters(in: .whitespaces)

    // Blank lines pass through verbatim.
    if trimmed.isEmpty { return line }

    if trimmed.hasPrefix("#") {
      // Only `#EXT...` tags may carry a rewritable `URI="..."` attribute.
      // Plain comment lines (`#` not followed by `EXT`) are ignored per
      // RFC 8216 §4.1 and pass through verbatim, even if they happen to
      // contain `URI="..."` in prose.
      guard trimmed.hasPrefix("#EXT") else { return line }
      return rewriteURIAttribute(in: line, baseURL: baseURL, transform: transform)
    }

    // Otherwise this is a bare URI (segment or variant/media playlist).
    guard !isPreservedScheme(trimmed) else { return line }
    guard let resolved = resolve(trimmed, against: baseURL) else { return line }
    return transform(resolved).absoluteString
  }

  /// Rewrite the value of the `URI` attribute on a `#EXT...` tag in place,
  /// preserving the rest of the line exactly. Returns the line unchanged when
  /// no such attribute is present, when the value is empty or uses a preserved
  /// scheme, or when resolution fails.
  ///
  /// This is intentionally generic: any tag bearing a `URI="..."` attribute is
  /// covered (EXT-X-KEY, EXT-X-MAP, EXT-X-MEDIA, EXT-X-I-FRAME-STREAM-INF,
  /// EXT-X-SESSION-KEY, EXT-X-SESSION-DATA, EXT-X-PART, EXT-X-PRELOAD-HINT,
  /// EXT-X-RENDITION-REPORT, …), with no per-tag hardcoding. The match is
  /// anchored to the attribute *name* `URI` (preceded by the tag's `:` or an
  /// attribute-list `,`), so it never fires on an attribute whose name merely
  /// ends in `URI` (e.g. a vendor `X-FOO-URI`) nor on `URI="` appearing inside
  /// another attribute's quoted value.
  private static func rewriteURIAttribute(
    in line: String,
    baseURL: URL,
    transform: (URL) -> URL
  ) -> String {
    guard let valueRange = uriAttributeValueRange(in: line) else { return line }

    let value = String(line[valueRange])

    guard !value.isEmpty else { return line }
    guard !isPreservedScheme(value) else { return line }
    guard let resolved = resolve(value, against: baseURL) else { return line }

    let replacement = transform(resolved).absoluteString
    return line.replacingCharacters(in: valueRange, with: replacement)
  }

  /// Locate the quoted value of the `URI` attribute on a tag line, returning the
  /// range of the value *between* the quotes (the quotes themselves are not
  /// included). Returns `nil` when there is no well-formed `URI="..."` attribute.
  ///
  /// Exposed `internal` so `HLSPlaylistParser` can reuse the same RFC 8216
  /// `URI="..."` scanner rather than re-implement subtly-different logic.
  static func uriAttributeValueRange(in line: String) -> Range<String.Index>? {
    attributeValueRange(named: "URI", in: line)
  }

  /// Locate the value of the attribute named `name` (case-insensitive) on a
  /// `#EXT...` tag line, returning the range of the value; for a quoted value
  /// the surrounding quotes are excluded. Returns `nil` when the attribute is
  /// absent or, for a quoted value, the closing quote is missing.
  ///
  /// Scans the comma-separated attribute list after the tag's first `:` per
  /// RFC 8216 §4.2. The scan is *quote-aware*: a `,` inside a quoted value (e.g.
  /// `CODECS="avc1.4d401f,mp4a.40.2"`) does not split attributes, so an
  /// attribute appearing after such a value is still matched correctly. Optional
  /// whitespace around `=` is tolerated for lenient producers.
  ///
  /// Exposed `internal` so `HLSPlaylistParser` reads BANDWIDTH/METHOD through the
  /// same quote-aware scanner instead of a naive comma split.
  static func attributeValueRange(named name: String, in line: String) -> Range<String.Index>? {
    // Attributes live after the tag name's first ":" (e.g. "#EXT-X-KEY:...").
    guard let colon = line.firstIndex(of: ":") else { return nil }

    var index = line.index(after: colon)
    let end = line.endIndex

    while index < end {
      while index < end, line[index] == " " || line[index] == "\t" {
        index = line.index(after: index)
      }

      // Read the attribute name up to "=" — names are unquoted tokens.
      let nameStart = index
      while index < end, line[index] != "=", line[index] != "," {
        index = line.index(after: index)
      }
      let attrName = line[nameStart..<index].trimmingCharacters(in: .whitespaces)
      let matches = attrName.caseInsensitiveCompare(name) == .orderedSame

      // No "=" for this token (a value-less token or trailing junk): skip it.
      guard index < end, line[index] == "=" else {
        if index < end { index = line.index(after: index) } // step over the ","
        continue
      }

      // Move past "=" and any whitespace to the value.
      index = line.index(after: index)
      while index < end, line[index] == " " || line[index] == "\t" {
        index = line.index(after: index)
      }

      if index < end, line[index] == "\"" {
        // Quoted value: find the closing quote.
        let valueStart = line.index(after: index)
        guard let closingQuote = line[valueStart..<end].firstIndex(of: "\"") else {
          return nil // Malformed (unterminated quote): leave the line as-is.
        }
        if matches {
          return valueStart..<closingQuote
        }
        // Advance past the closing quote, then the following "," if present.
        index = line.index(after: closingQuote)
      } else {
        // Unquoted value: read up to the next "," (the attribute separator).
        let valueStart = index
        while index < end, line[index] != "," {
          index = line.index(after: index)
        }
        if matches {
          return valueStart..<index
        }
      }

      // Step over the separating "," before the next attribute.
      if index < end, line[index] == "," {
        index = line.index(after: index)
      }
    }

    return nil
  }

  // MARK: - Detection

  /// Whether a response looks like an HLS playlist (master or media), based on
  /// its content type, URL extension, or a body sniff — in that order.
  ///
  /// - Parameters:
  ///   - contentType: The `Content-Type` header value, if any.
  ///   - url: The request/response URL, if any.
  ///   - body: The (leading) response bytes, if any.
  static func isPlaylist(contentType: String?, url: URL?, body: Data?) -> Bool {
    if let contentType = contentType {
      // Strip any ";charset=..." parameter and normalise.
      let main = contentType.prefix { $0 != ";" }
        .trimmingCharacters(in: .whitespaces)
        .lowercased()
      if playlistContentTypes.contains(main) { return true }
    }

    if let url = url {
      let ext = url.pathExtension.lowercased()
      if ext == "m3u8" || ext == "m3u" { return true }
    }

    if let body = body, bodyStartsWithEXTM3U(body) {
      return true
    }

    return false
  }

  /// Whether the playlist is a complete VOD playlist. True when an
  /// `#EXT-X-ENDLIST` tag is present, or when `#EXT-X-PLAYLIST-TYPE:VOD` is
  /// declared (RFC 8216 §4.3.3.4 / §4.3.3.5). Matched only as line-leading tags
  /// so the tokens are never picked up from comments or attribute values.
  static func isVOD(_ manifest: String) -> Bool {
    for rawLine in manifest.components(separatedBy: "\n") {
      let trimmed = rawLine.trimmingCharacters(in: .whitespacesAndNewlines)
      if trimmed == "#EXT-X-ENDLIST" { return true }
      if trimmed == "#EXT-X-PLAYLIST-TYPE:VOD" { return true }
    }
    return false
  }

  // MARK: - Helpers

  /// Content types recognised as HLS playlists (lowercased, no parameters).
  private static let playlistContentTypes: Set<String> = [
    "application/vnd.apple.mpegurl",
    "application/x-mpegurl",
    "application/mpegurl",
    "audio/mpegurl",
    "audio/x-mpegurl",
    "vnd.apple.mpegurl",
  ]

  /// Resolve a (possibly relative) URI against the playlist URL per RFC 3986.
  /// Returns `nil` for an empty URI or when the URI cannot be parsed/resolved.
  ///
  /// Exposed `internal` so `HLSPlaylistParser` resolves URIs identically.
  static func resolve(_ uri: String, against baseURL: URL) -> URL? {
    guard !uri.isEmpty else { return nil }
    return URL(string: uri, relativeTo: baseURL)?.absoluteURL
  }

  /// Schemes whose URIs must never be proxied: FairPlay key delivery (`skd://`)
  /// and inline data (`data:`).
  ///
  /// Exposed `internal` so `HLSPlaylistParser` skips the same URIs.
  static func isPreservedScheme(_ uri: String) -> Bool {
    let lower = uri.lowercased()
    return lower.hasPrefix("skd://") || lower.hasPrefix("data:")
  }

  /// Whether the leading bytes of `body`, decoded as UTF-8, begin with the HLS
  /// magic `#EXTM3U` once a leading UTF-8 BOM and any whitespace are skipped.
  private static func bodyStartsWithEXTM3U(_ body: Data) -> Bool {
    // A generous prefix so the magic is reachable even after a BOM or a run of
    // leading whitespace, while still avoiding a full-body decode.
    var prefix = body.prefix(256)
    // Drop a leading UTF-8 BOM (EF BB BF) if present.
    if prefix.starts(with: [0xEF, 0xBB, 0xBF]) {
      prefix = prefix.dropFirst(3)
    }
    guard let text = String(data: prefix, encoding: .utf8)
      ?? String(bytes: prefix, encoding: .ascii) else {
      return false
    }
    return text.drop(while: { $0.isWhitespace }).hasPrefix("#EXTM3U")
  }
}