package main

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"flag"
	"log"
	"os"
	"path/filepath"
	"time"

	"go.vocdoni.io/dvote/api/censusdb"
	"go.vocdoni.io/dvote/censustree"
	"go.vocdoni.io/dvote/data/compressor"
	"go.vocdoni.io/dvote/data/downloader"
	"go.vocdoni.io/dvote/data/ipfs"
	"go.vocdoni.io/dvote/data/ipfs/ipfsconnect"
	"go.vocdoni.io/dvote/db"
	"go.vocdoni.io/dvote/db/metadb"
	"go.vocdoni.io/dvote/types"
	"go.vocdoni.io/proto/build/go/models"
)

func main() {
	log.SetFlags(log.Lshortfile | log.LstdFlags)
	// parse flags
	ipfsKey := flag.String("ipfs-key", "", "IPFS key to use")
	uri := flag.String("uri", "", "IPFS URI of the census")
	dumpPath := flag.String("dump", "", "the path of the census dump")
	save := flag.Bool("save", false, "save the census to a file")
	flag.Parse()
	if (uri == nil || *uri == "") && (dumpPath == nil || *dumpPath == "") {
		log.Fatal("IPFS URI or dump path required. Use -uri or -dump flags")
	}
	// init temp folder and defer its removal
	testTempPath := "./temp"
	defer os.RemoveAll(testTempPath)

	ipfsData := []byte{}
	if uri != nil && *uri != "" {
		log.Println("Downloading census from IPFS...")
		// init ipfs service
		storage := new(ipfs.Handler)
		if err := storage.Init(&types.DataStore{Datadir: testTempPath}); err != nil {
			log.Fatal(err)
		}
		var ipfsConn *ipfsconnect.IPFSConnect
		if ipfsKey != nil && *ipfsKey != "" {
			ipfsConn = ipfsconnect.New(*ipfsKey, storage)
			ipfsConn.Start()
		}

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
		defer cancel()

		var err error
		ipfsData, err = storage.Retrieve(ctx, *uri, downloader.MaxFileSize)
		if err != nil {
			log.Fatal(err)
		}
	} else {
		log.Println("Reading census dump...")
		var err error
		if ipfsData, err = os.ReadFile(*dumpPath); err != nil {
			log.Fatal(err)
		}
	}
	log.Println("Importing census from dump...")

	// unmarshal the census dump
	census := censusdb.CensusDump{}
	if err := json.Unmarshal(ipfsData, &census); err != nil {
		log.Fatal(err)
	}
	// init test key-value database
	database, err := metadb.New(db.TypePebble, filepath.Join(testTempPath, "censusdb"))
	if err != nil {
		log.Fatal(err)
	}
	// init the target tree
	dbTree, err := censustree.New(censustree.Options{
		ParentDB:   database,
		CensusType: models.Census_Type(census.Type),
		MaxLevels:  census.MaxLevels,
		Name:       "test",
	})
	if err != nil {
		log.Fatal(err)
	}
	// decompress the retrieved dump and import it into the target tree
	decompressDump := compressor.NewCompressor().DecompressBytes(census.Data)
	if err := dbTree.ImportDump(decompressDump); err != nil {
		log.Fatal(err)
	}
	// get root
	newRoot, err := dbTree.Root()
	if err != nil {
		log.Fatal(err)
	}

	// print results
	log.Printf("Received root: %s, Computed root: %s\n", hex.EncodeToString(census.RootHash), hex.EncodeToString(newRoot))

	holders := map[string]string{}
	dbTree.IterateLeaves(func(key, value []byte) bool {
		holders["0x"+hex.EncodeToString(key)] = dbTree.BytesToBigInt(value).String()
		return true
	})
	// save the census to a file
	if save != nil && *save {
		log.Println("Saving census to file...")
		fd, err := os.Create("census.json")
		if err != nil {
			log.Fatal(err)
		}
		defer fd.Close()
		enc := json.NewEncoder(fd)
		enc.SetIndent("", "  ")
		if err := enc.Encode(holders); err != nil {
			log.Fatal(err)
		}
	}
}