//
//  Cleaner.swift
//  SwiftSoup
//
//  Created by Nabil Chatbi on 15/10/16.
//

import Foundation

open class Cleaner {
  fileprivate let headWhitelist: Whitelist?
  fileprivate let bodyWhitelist: Whitelist

  /// Create a new cleaner, that sanitizes documents' `<head>` and  `<body>` using the supplied whitelist.
  /// - Parameters:
  ///   - headWhitelist: Whitelist to clean the head with
  ///   - bodyWhitelist: Whitelist to clean the body with
  public init(headWhitelist: Whitelist?, bodyWhitelist: Whitelist) {
    self.headWhitelist = headWhitelist
    self.bodyWhitelist = bodyWhitelist
  }

  /// Create a new cleaner, that sanitizes documents' `<body>` using the supplied whitelist.
  /// - Parameter whitelist: Whitelist to clean the body with
  convenience init(_ whitelist: Whitelist) {
    self.init(headWhitelist: nil, bodyWhitelist: whitelist)
  }

  /// Creates a new, clean document, from the original dirty document, containing only elements allowed by the whitelist.
  /// The original document is not modified. Only elements from the dirt document's `<body>` are used.
  /// - Parameter dirtyDocument: Untrusted base document to clean.
  /// - Returns: A cleaned document.
  public func clean(_ dirtyDocument: Document) throws -> Document {
    let clean = Document.createShell(dirtyDocument.getBaseUri())
    if let headWhitelist, let dirtHead = dirtyDocument.head(), let cleanHead = clean.head() {  // frameset documents won't have a head. the clean doc will have empty head.
      try copySafeNodes(dirtHead, cleanHead, whitelist: headWhitelist)
    }
    if let dirtBody = dirtyDocument.body(), let cleanBody = clean.body() {  // frameset documents won't have a body. the clean doc will have empty body.
      try copySafeNodes(dirtBody, cleanBody, whitelist: bodyWhitelist)
    }
    return clean
  }

  /// Determines if the input document is valid, against the whitelist. It is considered valid if all the tags and attributes
  /// in the input HTML are allowed by the whitelist.
  ///
  /// This method can be used as a validator for user input forms. An invalid document will still be cleaned successfully
  /// using the ``clean(_:)`` document. If using as a validator, it is recommended to still clean the document
  /// to ensure enforced attributes are set correctly, and that the output is tidied.
  /// - Parameter dirtyDocument: document to test
  /// - Returns: true if no tags or attributes need to be removed; false if they do
  public func isValid(_ dirtyDocument: Document) throws -> Bool {
    let clean = Document.createShell(dirtyDocument.getBaseUri())
    let numDiscarded = try copySafeNodes(
      dirtyDocument.body()!, clean.body()!, whitelist: bodyWhitelist)
    return numDiscarded == 0
  }

  @discardableResult
  fileprivate func copySafeNodes(_ source: Element, _ dest: Element, whitelist: Whitelist) throws
    -> Int
  {
    let cleaningVisitor = Cleaner.CleaningVisitor(source, dest, whitelist)
    try NodeTraversor(cleaningVisitor).traverse(source)
    return cleaningVisitor.numDiscarded
  }
}

extension Cleaner {
  fileprivate final class CleaningVisitor: NodeVisitor {
    private(set) var numDiscarded = 0

    private let root: Element
    private var destination: Element?  // current element to append nodes to

    private let whitelist: Whitelist

    public init(_ root: Element, _ destination: Element, _ whitelist: Whitelist) {
      self.root = root
      self.destination = destination
      self.whitelist = whitelist
    }

    public func head(_ source: Node, _ depth: Int) throws {
      if let sourceEl = source as? Element {
        if whitelist.isSafeTag(sourceEl.tagNameUTF8()) {  // safe, clone and copy safe attrs
          let meta = try createSafeElement(sourceEl)
          let destChild = meta.el
          try destination?.appendChild(destChild)

          numDiscarded += meta.numAttribsDiscarded
          destination = destChild
        } else if source != root {  // not a safe tag, so don't add. don't count root against discarded.
          numDiscarded += 1
        }
      } else if let sourceText = source as? TextNode {
        let destText = TextNode(sourceText.getWholeTextUTF8(), source.getBaseUriUTF8())
        try destination?.appendChild(destText)
      } else if let sourceData = source as? DataNode {
        if sourceData.parent() != nil && whitelist.isSafeTag(sourceData.parent()!.nodeNameUTF8()) {
          let destData = DataNode(sourceData.getWholeDataUTF8(), source.getBaseUriUTF8())
          try destination?.appendChild(destData)
        } else {
          numDiscarded += 1
        }
      } else {  // else, we don't care about comments, xml proc instructions, etc
        numDiscarded += 1
      }
    }

    public func tail(_ source: Node, _ depth: Int) throws {
      if let x = source as? Element {
        if whitelist.isSafeTag(x.nodeNameUTF8()) {
          // would have descended, so pop destination stack
          destination = destination?.parent()
        }
      }
    }

    private func createSafeElement(_ sourceEl: Element) throws -> ElementMeta {
      let sourceTag = sourceEl.tagName()
      let destAttrs = Attributes()
      var numDiscarded = 0

      if let sourceAttrs = sourceEl.getAttributes() {
        for sourceAttr in sourceAttrs {
          if let destAttr = try whitelist.safeAttribute(sourceTag, sourceEl, sourceAttr) {
            destAttrs.put(attribute: destAttr)
          } else {
            numDiscarded += 1
          }
        }
      }
      let enforcedAttrs = try whitelist.getEnforcedAttributes(sourceTag)
      destAttrs.addAll(incoming: enforcedAttrs)

      let dest = try Element(Tag.valueOf(sourceTag.utf8Array), sourceEl.getBaseUriUTF8(), destAttrs)
      if let treeBuilder = sourceEl.treeBuilder {
        dest.treeBuilder = treeBuilder
      }
      return ElementMeta(dest, numDiscarded)
    }
  }
}

extension Cleaner {
  fileprivate struct ElementMeta {
    let el: Element
    let numAttribsDiscarded: Int

    init(_ el: Element, _ numAttribsDiscarded: Int) {
      self.el = el
      self.numAttribsDiscarded = numAttribsDiscarded
    }
  }
}
