{.push raises: [Defect]}

import
  std/[sequtils, strformat, strutils, tables],
  stew/[base64, base32, byteutils, results],
  nimcrypto/[hash, keccak],
  ./tree

export tree

## A collection of utilities for constructing a Merkle Tree
## encoding a list of ENR and link entries.
## The tree consists of DNS TXT records.
## 
## Discovery via DNS is based on https://eips.ethereum.org/EIPS/eip-1459
## 
## This implementation is based on the Go implementation of EIP-1459
## at https://github.com/ethereum/go-ethereum/blob/master/p2p/dnsdisc


## How we determine MaxChildren: (Adapted from go-ethereum)
## https://github.com/ethereum/go-ethereum/blob/4d88974864c3ee84a24b1088064162b8dbab8ad5/p2p/dnsdisc/tree.go#L116-L146
## 
## We want to keep the UDP size below 512 bytes. The UDP size is roughly:
## UDP length = 8 + UDP payload length ( 229 )
## UPD Payload length:
##  - dns.id 2
##  - dns.flags 2
##  - dns.count.queries 2
##  - dns.count.answers 2
##  - dns.count.auth_rr 2
##  - dns.count.add_rr 2
##  - queries (query-size + 6)
##  - answers :
##  	- dns.resp.name 2
##  	- dns.resp.type 2
##  	- dns.resp.class 2
##  	- dns.resp.ttl 4
##  	- dns.resp.len 2
##  	- dns.txt.length 1
##  	- dns.txt resp_data_size
## So the total size is roughly a fixed overhead of `39`, and the size of the
## query (domain name) and response.
## The query size is, for example, FVY6INQ6LZ33WLCHO3BPR3FH6Y.snap.mainnet.ethdisco.net (52)
## We also have some static data in the response, such as `enrtree-branch:`, and potentially
## splitting the response up with `" "`, leaving us with a size of roughly `400` that we need
## to stay below.
## The number `370` is used to have some margin for extra overhead (for example, the dns query
## may be larger - more subdomains).
const
  HashAbbrevSize = 1 + (16*13)/8                 # Size of an encoded hash (plus comma)
  MaxChildren* = 370 div toInt(HashAbbrevSize)   # 13 children. See comment above for explanation.

type
  BuilderResult*[T] = Result[T, string]

  Subtree* = object # A coherent section of a larger tree. Can contain branch and leaf nodes.
    subtreeRoot*: SubtreeEntry
    subtreeEntries*: seq[SubtreeEntry]

####################
# Helper functions #
####################

proc toTXTRecord*(rootEntry: RootEntry): BuilderResult[string] =
  ## Converts a root entry into its corresponding
  ## TXT record representation
  var txtRecord: string

  let sig = Base64Url.encode(rootEntry.signature)
  try:
    txtRecord = fmt"{RootPrefix} e={rootEntry.eroot} l={rootEntry.lroot} seq={rootEntry.seqNo} sig={sig}"
  except ValueError:
    return err("Failed to format root entry")

  return ok(txtRecord)

proc toTXTRecord*(subtreeEntry: SubtreeEntry): BuilderResult[string] =
  ## Converts a subtree entry into its corresponding
  ## TXT record representation
  var txtRecord: string

  case subtreeEntry.kind:
    of Enr:
      txtRecord = subtreeEntry.enrEntry.record.toURI()
    of Link:
      txtRecord = LinkPrefix & subtreeEntry.linkEntry.str
    of Branch:
      txtRecord = BranchPrefix & subtreeEntry.branchEntry.children.join(",")
  
  return ok(txtRecord)

proc subdomain*(subtreeEntry: SubtreeEntry): BuilderResult[string] =
  ## Computes the subdomain hash for a subtree entry
  ## The subdomain name of any entry is the base32
  ## encoding of the (abbreviated) keccak256 hash of
  ## its TXT record content.
  var txtRecord: string
  
  try:
    txtRecord = subtreeEntry.toTXTRecord().tryGet()
  except ValueError:
    return err("Failed to format subtree entry")
  
  let
    keccakHash = keccak256.digest(txtRecord.toBytes()).data[0..15]
    subdomain = Base32.encode(keccakHash)
  
  return ok(subdomain)

###############
# Builder API #
###############

proc buildTXT*(tree: Tree, domain: string): BuilderResult[Table[string, string]] =
  ## Builds the TXT records for a given tree at
  ## a given domain. Returns a map of (sub)domain
  ## to TXT record for the full tree.

  var treeRecords = initTable[string, string]()

  # Add root entry
  let rootRecordRes = tree.rootEntry.toTXTRecord()

  if rootRecordRes.isErr:
    return err("Failed to build: " & rootRecordRes.error)
  else:
    treeRecords[domain] = rootRecordRes.get()

  # Add subtree entries
  for subtreeEntry in tree.entries:
    let
      subdomainRes = subtreeEntry.subdomain()
      txtRecordRes = subtreeEntry.toTXTRecord()
    
    if subdomainRes.isErr:
      return err("Failed to build: " & subdomainRes.error)
    
    if txtRecordRes.isErr:
      return err("Failed to build: " & txtRecordRes.error)

    treeRecords[subdomainRes.get() & "." & domain] = txtRecordRes.get()
  
  return ok(treeRecords)

proc buildSubtree*(entries: seq[SubtreeEntry]): BuilderResult[Subtree] =
  ## Builds a subtree from a given list of entries.
  ## Returns the built subtree entries and the root of the subtree.
  var subtree: Subtree

  if entries.len() == 1:
    # Single entry is its own root
    subtree.subtreeRoot = entries[0]
    return ok(subtree)
  
  if entries.len() <= MaxChildren:
    # Entries will fit in single branch
    # Determine subdomain hashes
    var children: seq[string]

    for entry in entries:
      let subdomainRes = entry.subdomain()
      
      if subdomainRes.isErr:
        return err("Failed to build subtree: " & subdomainRes.error)
      
      children.add(subdomainRes.get())
    
    # Return branch as subtree
    subtree.subtreeRoot = SubtreeEntry(kind: Branch,
                                       branchEntry: BranchEntry(children: children))
    subtree.subtreeEntries = entries
    return ok(subtree)
  
  ## Several branches required. The algorithm is now:
  ## 1. Create a branch subtree for each slice of entries that fits within MaxChildren
  ## 2. Create a subtree consisting of the subtree root entries of all branches in (1)
  ## 3. Combine entries from (1) and (2) into a single subtree
  var
    subtrees: seq[Subtree]
    entriesToAdd = entries
  while entriesToAdd.len > 0:
    # Iterate until all entries are part of a subtree
    let sliceSize = if entriesToAdd.len < MaxChildren: entriesToAdd.len
                    else: MaxChildren

    let subtreeRes = buildSubtree(entriesToAdd[0..<sliceSize])

    if subtreeRes.isErr:
      return err(subtreeRes.error)

    subtrees.add(subtreeRes.get())

    entriesToAdd = entriesToAdd[sliceSize..<entriesToAdd.len] # Continue with remainder of entries

  # Collect entries and roots from all subtrees
  let
    combinedEntries = subtrees.mapIt(it.subtreeEntries).concat()
    combinedRoots = subtrees.mapIt(it.subtreeRoot)

  # Build a subtree from the roots of the existing subtrees. Recursion is fun!
  let rootsSubtreeRes = buildSubtree(combinedRoots)
  if rootsSubtreeRes.isErr:
    return err(rootsSubtreeRes.error)
  
  let rootsSubtree = rootsSubtreeRes.get()
  
  # Return combined subtree
  subtree.subtreeRoot = rootsSubtree.subtreeRoot
  subtree.subtreeEntries = concat(rootsSubtree.subtreeEntries, combinedEntries)
  return ok(subtree)

proc signTree*(tree: var Tree, privateKey: PrivateKey): BuilderResult[void] =
  ## Signs the tree with the provided private key
  var
    sig: Signature
    rootEntry = tree.rootEntry
  
  try:
    sig = sign(privateKey, hashableContent(rootEntry))
  except ValueError:
    return err("Failed to sign root entry")

  tree.rootEntry.signature = @(sig.toRaw())
  return ok()

proc buildTree*(seqNo: uint32,
                enrRecords: seq[Record],
                links: seq[LinkEntry]): BuilderResult[Tree] =
  ## Builds a tree from given lists of ENR and links.
  
  var tree: Tree
  
  # @TODO verify ENR here - should be signed

  # Convert ENR and links to subtree sequences
  var
    enrEntries: seq[SubtreeEntry]
    linkEntries: seq[SubtreeEntry]
  
  enrEntries = enrRecords.mapIt(SubtreeEntry(kind: Enr, enrEntry: EnrEntry(record: it)))
  linkEntries = links.mapIt(SubtreeEntry(kind: Link, linkEntry: it))
  
  # Build ENR and link subtrees
  
  let enrSubtreeRes = buildSubtree(enrEntries)

  if enrSubtreeRes.isErr:
    return err("Failed to build ENR subtree: " & enrSubtreeRes.error)

  let linkSubtreeRes = buildSubtree(linkEntries)

  if linkSubtreeRes.isErr:
    return err("Failed to build link subtree: " & linkSubtreeRes.error)

  # Create a root entry
  let
    erootRes = enrSubtreeRes.get().subtreeRoot.subdomain()
    lrootRes = linkSubtreeRes.get().subtreeRoot.subdomain()
  
  if erootRes.isErr or lrootRes.isErr:
    return err("Failed to determine subtree root subdomain")

  let rootEntry = RootEntry(eroot: erootRes.get(),
                            lroot: lrootRes.get(),
                            seqNo: seqNo)
  
  # Combine subtrees
  let entries = concat(@[enrSubtreeRes.get().subtreeRoot,
                         linkSubtreeRes.get().subtreeRoot],
                       enrSubtreeRes.get().subtreeEntries,
                       linkSubtreeRes.get().subtreeEntries)
  
  tree = Tree(rootEntry: rootEntry,
              entries: entries)
  
  return ok(tree)