/*
 * SPDX-License-Identifier: (LGPL-3.0-only WITH LGPL-3.0-linking-exception) OR MPL-2.0
 */

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package mozilla.components.lib.publicsuffixlist

import java.net.IDN
import mozilla.components.lib.publicsuffixlist.ext.binarySearch

/** Class wrapping the public suffix list data and offering methods for accessing rules in it. */
internal class PublicSuffixListData(
  private val rules: ByteArray,
  private val exceptions: ByteArray,
) {

  private fun binarySearchRules(labels: List<ByteArray>, labelIndex: Int): String? {
    return rules.binarySearch(labels, labelIndex)
  }

  private fun binarySearchExceptions(labels: List<ByteArray>, labelIndex: Int): String? {
    return exceptions.binarySearch(labels, labelIndex)
  }

  @Suppress("ReturnCount")
  fun getPublicSuffixOffset(domain: String): PublicSuffixOffset? {
    if (domain.isEmpty()) {
      return null
    }

    val domainLabels = IDN.toUnicode(domain).split('.')
    if (domainLabels.find { it.isEmpty() } != null) {
      // At least one of the labels is empty: Bail out.
      return null
    }

    val rule = findMatchingRule(domainLabels)

    if (domainLabels.size == rule.size && rule[0][0] != EXCEPTION_MARKER) {
      // The domain is a public suffix.
      return if (rule == PREVAILING_RULE) {
        PublicSuffixOffset.PrevailingRule
      } else {
        PublicSuffixOffset.PublicSuffix
      }
    }

    return if (rule[0][0] == EXCEPTION_MARKER) {
      // Exception rules hold the effective TLD plus one.
      PublicSuffixOffset.Offset(domainLabels.size - rule.size)
    } else {
      // Otherwise the rule is for a public suffix, so we must take one more label.
      PublicSuffixOffset.Offset(domainLabels.size - (rule.size + 1))
    }
  }

  /**
   * Find a matching rule for the given domain labels.
   *
   * This algorithm is based on OkHttp's PublicSuffixDatabase class:
   * https://github.com/square/okhttp/blob/master/okhttp/src/main/java/okhttp3/internal/publicsuffix/PublicSuffixDatabase.java
   */
  private fun findMatchingRule(domainLabels: List<String>): List<String> {
    // Break apart the domain into UTF-8 labels, i.e. foo.bar.com turns into [foo, bar, com].
    val domainLabelsBytes = domainLabels.map { it.toByteArray(Charsets.UTF_8) }

    val exactMatch = findExactMatch(domainLabelsBytes)
    val wildcardMatch = findWildcardMatch(domainLabelsBytes)
    val exceptionMatch = findExceptionMatch(domainLabelsBytes, wildcardMatch)

    if (exceptionMatch != null) {
      return ("$EXCEPTION_MARKER$exceptionMatch").split('.')
    }

    if (exactMatch == null && wildcardMatch == null) {
      return PREVAILING_RULE
    }

    val exactRuleLabels = exactMatch?.split('.') ?: EMPTY_RULE
    val wildcardRuleLabels = wildcardMatch?.split('.') ?: EMPTY_RULE

    return if (exactRuleLabels.size > wildcardRuleLabels.size) {
      exactRuleLabels
    } else {
      wildcardRuleLabels
    }
  }

  /** Returns an exact match or null. */
  private fun findExactMatch(labels: List<ByteArray>): String? {
    // Start by looking for exact matches. We start at the leftmost label. For example,
    // foo.bar.com
    // will look like: [foo, bar, com], [bar, com], [com]. The longest matching rule wins.

    for (i in labels.indices) {
      val rule = binarySearchRules(labels, i)

      if (rule != null) {
        return rule
      }
    }

    return null
  }

  /** Returns a wildcard match or null. */
  private fun findWildcardMatch(labels: List<ByteArray>): String? {
    // In theory, wildcard rules are not restricted to having the wildcard in the leftmost
    // position.
    // In practice, wildcards are always in the leftmost position. For now, this implementation
    // cheats and does not attempt every possible permutation. Instead, it only considers
    // wildcards
    // in the leftmost position. We assert this fact when we generate the public suffix file. If
    // this assertion ever fails we'll need to refactor this implementation.
    if (labels.size > 1) {
      val labelsWithWildcard = labels.toMutableList()
      for (labelIndex in 0 until labelsWithWildcard.size) {
        labelsWithWildcard[labelIndex] = WILDCARD_LABEL
        val rule = binarySearchRules(labelsWithWildcard, labelIndex)
        if (rule != null) {
          return rule
        }
      }
    }

    return null
  }

  private fun findExceptionMatch(labels: List<ByteArray>, wildcardMatch: String?): String? {
    // Exception rules only apply to wildcard rules, so only try it if we matched a wildcard.
    if (wildcardMatch == null) {
      return null
    }

    for (labelIndex in labels.indices) {
      val rule = binarySearchExceptions(labels, labelIndex)
      if (rule != null) {
        return rule
      }
    }

    return null
  }

  companion object {

    val WILDCARD_LABEL = byteArrayOf('*'.code.toByte())
    val PREVAILING_RULE = listOf("*")
    val EMPTY_RULE = listOf<String>()
    const val EXCEPTION_MARKER = '!'
  }
}

internal sealed class PublicSuffixOffset {
  data class Offset(val value: Int) : PublicSuffixOffset()

  object PublicSuffix : PublicSuffixOffset()

  object PrevailingRule : PublicSuffixOffset()
}