import std/sequtils
import std/sugar
import std/strutils
import std/options

import pkg/poseidon2
import pkg/poseidon2/io
import pkg/constantine/math/arithmetic
import pkg/constantine/math/io/io_bigints
import pkg/constantine/math/io/io_fields

import pkg/codex/merkletree
import pkg/codex/slots
import pkg/codex/slots/types
import pkg/codex/utils/json

export types

func fromCircomData*(_: type Poseidon2Hash, cellData: seq[byte]): seq[Poseidon2Hash] =
  var
    pos = 0
    cellElms: seq[Bn254Fr]
  while pos < cellData.len:
    var
      step = 32
      offset = min(pos + step, cellData.len)
      data = cellData[pos..<offset]
    let ff = Bn254Fr.fromBytes(data.toArray32).get
    cellElms.add(ff)
    pos += data.len

  cellElms

func toJsonDecimal*(big: BigInt[254]): string =
  let s = big.toDecimal.strip( leading = true, trailing = false, chars = {'0'} )
  if s.len == 0: "0" else: s

func toJson*(g1: CircomG1): JsonNode =
  %* {
    "x": Bn254Fr.fromBytes(g1.x).get.toBig.toJsonDecimal,
    "y": Bn254Fr.fromBytes(g1.y).get.toBig.toJsonDecimal
  }

func toJson*(g2: CircomG2): JsonNode =
  %* {
    "x": [
      Bn254Fr.fromBytes(g2.x[0]).get.toBig.toJsonDecimal,
      Bn254Fr.fromBytes(g2.x[1]).get.toBig.toJsonDecimal],
    "y": [
      Bn254Fr.fromBytes(g2.y[0]).get.toBig.toJsonDecimal,
      Bn254Fr.fromBytes(g2.y[1]).get.toBig.toJsonDecimal]
  }

proc toJson*(vpk: VerifyingKey): JsonNode =
  let
    ic = toSeq(cast[ptr UncheckedArray[CircomG1]](vpk.ic).toOpenArray(0, vpk.icLen.int - 1))

  echo ic.len
  %* {
    "alpha1": vpk.alpha1.toJson,
    "beta2": vpk.beta2.toJson,
    "gamma2": vpk.gamma2.toJson,
    "delta2": vpk.delta2.toJson,
    "ic": ic.mapIt( it.toJson )
  }

func toJson*(input: ProofInputs[Poseidon2Hash]): JsonNode =
  var
    input = input

  %* {
    "dataSetRoot": input.datasetRoot.toBig.toJsonDecimal,
    "entropy": input.entropy.toBig.toJsonDecimal,
    "nCellsPerSlot": input.nCellsPerSlot,
    "nSlotsPerDataSet": input.nSlotsPerDataSet,
    "slotIndex": input.slotIndex,
    "slotRoot": input.slotRoot.toDecimal,
    "slotProof": input.slotProof.mapIt( it.toBig.toJsonDecimal ),
    "cellData": input.samples.mapIt(
      toSeq( it.cellData.elements(Poseidon2Hash) ).mapIt( it.toBig.toJsonDecimal )
    ),
    "merklePaths": input.samples.mapIt(
      it.merklePaths.mapIt( it.toBig.toJsonDecimal )
    )
  }

func jsonToProofInput*(_: type Poseidon2Hash, inputJson: JsonNode): ProofInputs[Poseidon2Hash] =
  let
    cellData =
      inputJson["cellData"].mapIt(
          it.mapIt(
            block:
              var
                big: BigInt[256]
                data = newSeq[byte](big.bits div 8)
              assert bool(big.fromDecimal( it.str ))
              data.marshal(big, littleEndian)
              data
          ).concat # flatten out elements
        )

    merklePaths =
        inputJson["merklePaths"].mapIt(
          it.mapIt(
            block:
              var
                big: BigInt[254]
                hash: Poseidon2Hash
              assert bool(big.fromDecimal( it.getStr ))
              hash.fromBig( big )
              hash
          )
        )

    slotProof = inputJson["slotProof"].mapIt(
      block:
        var
          big: BigInt[254]
          hash: Poseidon2Hash
        assert bool(big.fromDecimal( it.str ))
        hash.fromBig( big )
        hash
      )

    datasetRoot = block:
      var
        big: BigInt[254]
        hash: Poseidon2Hash
      assert bool(big.fromDecimal( inputJson["dataSetRoot"].str ))
      hash.fromBig( big )
      hash

    slotRoot = block:
      var
        big: BigInt[254]
        hash: Poseidon2Hash
      assert bool(big.fromDecimal( inputJson["slotRoot"].str ))
      hash.fromBig( big )
      hash

    entropy = block:
      var
        big: BigInt[254]
        hash: Poseidon2Hash
      assert bool(big.fromDecimal( inputJson["entropy"].str ))
      hash.fromBig( big )
      hash

    nCellsPerSlot = inputJson["nCellsPerSlot"].getInt
    nSlotsPerDataSet = inputJson["nSlotsPerDataSet"].getInt
    slotIndex = inputJson["slotIndex"].getInt

  ProofInputs[Poseidon2Hash](
    entropy: entropy,
    slotIndex: slotIndex,
    datasetRoot: datasetRoot,
    slotProof: slotProof,
    slotRoot: slotRoot,
    nCellsPerSlot: nCellsPerSlot,
    nSlotsPerDataSet: nSlotsPerDataSet,
    samples: zip(cellData, merklePaths)
      .mapIt(Sample[Poseidon2Hash](
        cellData: it[0],
        merklePaths: it[1]
      ))
  )