import std/sequtils
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 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(it.cellData.mapIt(it.toBig.toJsonDecimal)),
    "merklePaths": input.samples.mapIt(it.merklePaths.mapIt(it.toBig.toJsonDecimal)),
  }

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

func jsonToProofInput*(
    _: type Poseidon2Hash, inputJson: JsonNode
): ProofInputs[Poseidon2Hash] =
  let
    cellData = inputJson["cellData"].mapIt(
      it.mapIt(
        block:
          var
            big: BigInt[256]
            hash: Poseidon2Hash
            data: array[32, byte]
          assert bool(big.fromDecimal(it.str))
          assert data.marshal(big, littleEndian)

          Poseidon2Hash.fromBytes(data).get
      ).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])
      ),
  )