import std/sugar

import pkg/chronos
import pkg/libp2p/cid

import pkg/codex/codextypes
import pkg/codex/stores
import pkg/codex/merkletree
import pkg/codex/manifest
import pkg/codex/blocktype as bt
import pkg/codex/chunker
import pkg/codex/indexingstrategy
import pkg/codex/slots
import pkg/codex/rng
import pkg/codex/utils/poseidon2digest

import ../helpers

proc storeManifest*(store: BlockStore, manifest: Manifest): Future[?!bt.Block] {.async.} =
  without encodedVerifiable =? manifest.encode(), err:
    trace "Unable to encode manifest"
    return failure(err)

  without blk =? bt.Block.new(data = encodedVerifiable, codec = ManifestCodec), error:
    trace "Unable to create block from manifest"
    return failure(error)

  if err =? (await store.putBlock(blk)).errorOption:
    trace "Unable to store manifest block", cid = blk.cid, err = err.msg
    return failure(err)

  success blk

proc makeManifest*(
  cids: seq[Cid],
  datasetSize: NBytes,
  blockSize: NBytes,
  store: BlockStore,
  hcodec = Sha256HashCodec,
  dataCodec = BlockCodec): Future[?!Manifest] {.async.} =
  without tree =? CodexTree.init(cids), err:
    return failure(err)

  without treeCid =? tree.rootCid(CIDv1, dataCodec), err:
    return failure(err)

  for index, cid in cids:
    without proof =? tree.getProof(index), err:
      return failure(err)

    if err =? (await store.putCidAndProof(treeCid, index, cid, proof)).errorOption:
      # TODO add log here
      return failure(err)

  let
    manifest = Manifest.new(
      treeCid = treeCid,
      blockSize = blockSize,
      datasetSize = datasetSize,
      version = CIDv1,
      hcodec = hcodec,
      codec = dataCodec)

  without manifestBlk =? await store.storeManifest(manifest), err:
    trace "Unable to store manifest"
    return failure(err)

  success manifest

proc createBlocks*(
  chunker: Chunker,
  store: BlockStore): Future[seq[bt.Block]] {.async.} =

  collect(newSeq):
    while (let chunk = await chunker.getBytes(); chunk.len > 0):
      let blk = bt.Block.new(chunk).tryGet()
      discard await store.putBlock(blk)
      blk

proc createProtectedManifest*(
  datasetBlocks: seq[bt.Block],
  store: BlockStore,
  numDatasetBlocks: int,
  ecK: int, ecM: int,
  blockSize: NBytes,
  originalDatasetSize: int,
  totalDatasetSize: int):
  Future[tuple[manifest: Manifest, protected: Manifest]] {.async.} =

  let
    cids = datasetBlocks.mapIt(it.cid)
    datasetTree = CodexTree.init(cids[0..<numDatasetBlocks]).tryGet()
    datasetTreeCid = datasetTree.rootCid().tryGet()

    protectedTree = CodexTree.init(cids).tryGet()
    protectedTreeCid = protectedTree.rootCid().tryGet()

  for index, cid in cids[0..<numDatasetBlocks]:
    let proof = datasetTree.getProof(index).tryget()
    (await store.putCidAndProof(datasetTreeCid, index, cid, proof)).tryGet

  for index, cid in cids:
    let proof = protectedTree.getProof(index).tryget()
    (await store.putCidAndProof(protectedTreeCid, index, cid, proof)).tryGet

  let
    manifest = Manifest.new(
      treeCid = datasetTreeCid,
      blockSize = blockSize,
      datasetSize = originalDatasetSize.NBytes)

    protectedManifest = Manifest.new(
      manifest = manifest,
      treeCid = protectedTreeCid,
      datasetSize = totalDatasetSize.NBytes,
      ecK = ecK,
      ecM = ecM,
      strategy = SteppedStrategy)

    manifestBlock = bt.Block.new(
      manifest.encode().tryGet(),
      codec = ManifestCodec).tryGet()

    protectedManifestBlock = bt.Block.new(
      protectedManifest.encode().tryGet(),
      codec = ManifestCodec).tryGet()

  (await store.putBlock(manifestBlock)).tryGet()
  (await store.putBlock(protectedManifestBlock)).tryGet()

  (manifest, protectedManifest)

proc createVerifiableManifest*(
  store: BlockStore,
  numDatasetBlocks: int,
  ecK: int, ecM: int,
  blockSize: NBytes,
  cellSize: NBytes):
  Future[tuple[manifest: Manifest, protected: Manifest, verifiable: Manifest]] {.async.} =

  let
    numSlots = ecK + ecM
    numTotalBlocks = calcEcBlocksCount(numDatasetBlocks, ecK, ecM)  # total number of blocks in the dataset after
                                                                    # EC (should will match number of slots)
    originalDatasetSize = numDatasetBlocks * blockSize.int
    totalDatasetSize    = numTotalBlocks * blockSize.int

    chunker = RandomChunker.new(Rng.instance(), size = totalDatasetSize, chunkSize = blockSize)
    datasetBlocks = await chunker.createBlocks(store)

    (manifest, protectedManifest) =
        await createProtectedManifest(
          datasetBlocks,
          store,
          numDatasetBlocks,
          ecK, ecM,
          blockSize,
          originalDatasetSize,
          totalDatasetSize)

    builder = Poseidon2Builder.new(store, protectedManifest, cellSize = cellSize).tryGet
    verifiableManifest = (await builder.buildManifest()).tryGet

  # build the slots and manifest
  (manifest, protectedManifest, verifiableManifest)