package main

import (
	"bytes"
	"context"
	"encoding/hex"
	"fmt"
	"gitlab.com/Blockdaemon/go-tsm-sdkv2/v68/tsm"
	"golang.org/x/sync/errgroup"
	"sync"
)

func main() {

	// Create clients for two of the nodes

	configs := []*tsm.Configuration{
		tsm.Configuration{URL: "http://localhost:8501"}.WithAPIKeyAuthentication("apikey1"),
		tsm.Configuration{URL: "http://localhost:8502"}.WithAPIKeyAuthentication("apikey2"),
	}

	clients := make([]*tsm.Client, len(configs))
	for i, config := range configs {
		var err error
		if clients[i], err = tsm.NewClient(config); err != nil {
			panic(err)
		}
	}

	// Generate an EdDSA master key

	threshold := 1         // The security threshold to use for this key
	players := []int{1, 2} // We want to use MPC node 1 and 2
	keyGenSessionID := tsm.GenerateSessionID()
	keyGenSessionConfig := tsm.NewSessionConfig(keyGenSessionID, players, nil)
	masterkeyIDs := make([]string, len(clients))
	var eg errgroup.Group
	for i, client := range clients {
		client, i := client, i
		eg.Go(func() error {
			var err error
			masterkeyIDs[i], err = client.Schnorr().GenerateKey(context.TODO(), keyGenSessionConfig, threshold, "ED-25519", "")
			return err
		})
	}

	if err := eg.Wait(); err != nil {
		panic(err)
	}

	// Validate key IDs

	for i := 1; i < len(masterkeyIDs); i++ {
		if masterkeyIDs[0] != masterkeyIDs[i] {
			panic("key IDs do not match")
		}
	}
	masterKeyID := masterkeyIDs[0]
	fmt.Println("Generated master key with ID", masterKeyID, "on MPC nodes", players)

	// Get derived public key

	// In this example we will obtain the non-hardened derived key m/42/2/3
	derivationPath := []uint32{42, 2, 3}

	publicKeys := make([][]byte, len(clients))
	for i, client := range clients {
		var err error
		if publicKeys[i], err = client.Schnorr().PublicKey(context.TODO(), masterKeyID, derivationPath); err != nil {
			panic(err)
		}
	}

	// Validate derived public keys

	for i := 1; i < len(publicKeys); i++ {
		if bytes.Compare(publicKeys[0], publicKeys[i]) != 0 {
			panic("public keys do not match")
		}
	}
	publicKey := publicKeys[0]
	fmt.Println("Public key for derivation path", derivationPath, ":", hex.EncodeToString(publicKey))

	// We can now sign with the created key

	message := []byte("The message to be signed")

	fmt.Println("Creating signature using key derived according to path", derivationPath, " and players", players)
	partialSignaturesLock := sync.Mutex{}
	partialSignatures := make([][]byte, 0)
	signSessionID := tsm.GenerateSessionID()
	signSessionConfig := tsm.NewSessionConfig(signSessionID, players, nil)
	for _, client := range clients {
		client := client
		eg.Go(func() error {
			partialSignResult, err := client.Schnorr().Sign(context.TODO(), signSessionConfig, masterKeyID, derivationPath, message)
			if err != nil {
				return err
			}
			partialSignaturesLock.Lock()
			partialSignatures = append(partialSignatures, partialSignResult.PartialSignature)
			partialSignaturesLock.Unlock()
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		panic(err)
	}

	signature, err := tsm.SchnorrFinalizeSignature(message, partialSignatures)
	if err != nil {
		panic(err)
	}

	// Verify the signature relative to the signed message and the public key

	err = tsm.SchnorrVerifySignature(publicKey, message, signature)
	if err != nil {
		panic(err)
	}

	fmt.Println("Signature:", hex.EncodeToString(signature))
}

package com.example;

import com.sepior.tsm.sdkv2.*;

import java.util.*;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.function.Supplier;

public class SchnorrSignDerivedKeyExample {

    public static void main(String[] args) throws Exception {

        // Create a client for each MPC node

        Configuration[] configs = {
            new Configuration("http://localhost:8500"),
            new Configuration("http://localhost:8501"),
        };
        configs[0].withApiKeyAuthentication("Node0LoginPassword");
        configs[1].withApiKeyAuthentication("Node1LoginPassword");

        Client[] clients = {
            new Client(configs[0]),
            new Client(configs[1]),
        };


        // Generate an ECDSA key

        final int[] keyGenPlayers = {0, 1 }; // The key should be secret shared among the two MPC nodes
        final int threshold = 1; // The security threshold for this key
        final String curveName = "ED-25519"; // We want the key to be a Ed25519 key (e.g., for Polkadot)
        final int[] derivationPath = {42, 2, 3}; // In this example we will obtain the non-hardened derived key m/42/2/3

        String keyGenSessionId = SessionConfig.generateSessionId();
        final SessionConfig keyGenSessionConfig = SessionConfig.newSessionConfig(keyGenSessionId, keyGenPlayers, null);
        System.out.println("Generating key using players " + Arrays.toString(keyGenPlayers));
        List<String> results = runConcurrent(
                () -> clients[0].getSchnorr().generateKey(keyGenSessionConfig, threshold, curveName, null),
                () -> clients[1].getSchnorr().generateKey(keyGenSessionConfig, threshold, curveName, null));
        String keyId = results.get(0);
        System.out.println("Generated key with ID: " + keyId);


        // Get the public key

        byte[] publicKey = clients[0].getSchnorr().publicKey(keyId, derivationPath);
        System.out.println("Public key for path m/42/2/3: 0x" + bytesToHex(publicKey));

        // Remember: Depending on your use case, you may need to check that all or at least threshold + 1 clients agree on the
        // public key, before using it, e.g. for creating a cryptocurrency account.


        // Sign a message using the private key

        final int[] signPlayers = {0, 1}; // We sign with threshold + 1 players, in this case MPC node 0, 1
        final byte[] message = new byte[117];  // This is the message to be signed

        String signSessionId = SessionConfig.generateSessionId();
        final SessionConfig signSessionConfig = SessionConfig.newSessionConfig(signSessionId, signPlayers, null);

        System.out.println("Creating signature using derived key m/42/2/3 and players " + Arrays.toString(signPlayers));
        List<SchnorrPartialSignResult> signResults = runConcurrent(
                () -> clients[0].getSchnorr().sign(signSessionConfig, keyId, derivationPath, message),
                () -> clients[1].getSchnorr().sign(signSessionConfig, keyId, derivationPath, message));

        byte[][] partialSignatures = {signResults.get(0).getPartialSignature(), signResults.get(1).getPartialSignature()};
        byte[] signature = Schnorr.finalizeSignature(message, partialSignatures);
        System.out.println("Signature: 0x" + bytesToHex(signature));

        // Validate the signature

        boolean valid = Schnorr.verifySignature(publicKey, message, signature);
        System.out.println("Signature validity: " + valid);

    }


    @SafeVarargs
    static <T> List<T> runConcurrent(Supplier<T>... players) throws Exception {
        List<T> result = new ArrayList<T>(players.length);
        Queue<Exception> errors = new ConcurrentLinkedQueue<Exception>();
        for (int i = 0; i < players.length; i++) {
            result.add(null);
        }
        Thread[] threads = new Thread[players.length];
        for (int i = 0; i < players.length; i++) {
            final int index = i;
            Thread thread = new Thread() {
                public void run() {
                    try {
                        T runResult = players[index].get();
                        result.set(index, runResult);
                    } catch (Exception e) {
                        errors.add(e);
                    }
                }
            };
            threads[i] = thread;
            thread.start();
        }
        for (int i = 0; i < players.length; i++) {
            threads[i].join();
        }
        if (!errors.isEmpty()) {
            throw new RuntimeException("One of the threads failed executing command", errors.remove());
        }
        return result;
    }

    static char[] HEX_ARRAY = "0123456789ABCDEF".toCharArray();
    public static String bytesToHex(byte[] bytes) {
        char[] hexChars = new char[bytes.length * 2];
        for (int j = 0; j < bytes.length; j++) {
            int v = bytes[j] & 0xFF;
            hexChars[j * 2] = HEX_ARRAY[v >>> 4];
            hexChars[j * 2 + 1] = HEX_ARRAY[v & 0x0F];
        }
        return new String(hexChars);
    }

}

const { TSMClient, Configuration, SessionConfig, curves } = require("@sepior/tsmsdkv2");

async function main() {
  // Create clients for two of the nodes

  const configs = [
    {
      url: "http://localhost:8501",
      apiKey: "apikey1",
    },
    {
      url: "http://localhost:8502",
      apiKey: "apikey2",
    },
  ];

  const clients = [];

  for (const rawConfig of configs) {
    const config = await new Configuration(rawConfig.url);
    await config.withAPIKeyAuthentication(rawConfig.apiKey);
    const client = await TSMClient.withConfiguration(config);

    clients.push(client);
  }

  // Generate an EdDSA master key

  const threshold = 1; // The security threshold for this key
  const players = [1, 2]; // We want to use MPC node 1 and 2
  const keygenSessionConfig = await SessionConfig.newSessionConfig(
    await SessionConfig.GenerateSessionID(),
    new Uint32Array(players),
    {}
  );

  const masterKeyIds = ["", ""];

  const generateKeyPromises = [];

  for (const [i, client] of clients.entries()) {
    const func = async () => {
      const schnorrApi = client.Schnorr();
      masterKeyIds[i] = await schnorrApi.generateKey(
        keygenSessionConfig,
        threshold,
        curves.ED25519,
        ""
      );
    };
    generateKeyPromises.push(func());
  }

  await Promise.all(generateKeyPromises);

  // Validate key IDs

  for (let i = 1; i < masterKeyIds.length; i++) {
    if (masterKeyIds[0] !== masterKeyIds[i]) {
      console.log("Key ids do not match");
      return;
    }
  }

  const masterKeyId = masterKeyIds[0];

  console.log(
    `Generated master key with id: ${masterKeyId} on MPC nodes ${players}`
  );

  // Get derived public key

  // In this example we will obtain the non-hardened derived key m/42/2/3
  const derivationPath = new Uint32Array([42, 2, 3]);

  const publickeys = [];

  for (const client of clients) {
    const schnorrApi = client.Schnorr();
    publickeys.push(await schnorrApi.publicKey(masterKeyId, derivationPath));
  }

  // Validate public keys

  for (let i = 1; i < publickeys.length; i++) {
    if (!Buffer.from(publickeys[0]).equals(publickeys[i])) {
      console.log("Public keys does not match");
      return;
    }
  }

  const publicKey = publickeys[0];

  console.log(
    `Public key for derivation path m/42/2/3: ${Buffer.from(publicKey).toString(
      "hex"
    )}`
  );

  // We can now sign with the created key

  const message = Buffer.from("This is a message to be signed");

  const sessionId = await SessionConfig.GenerateSessionID();
  const signSessionConfig = await SessionConfig.newSessionConfig(
    sessionId,
    new Uint32Array(players),
    {}
  );

  console.log(
    `Creating signature using key derived according to path ${derivationPath} and players ${players}`
  );

  const partialSignatures = [];

  const partialSignaturePromises = [];

  for (let playerIdx = 0; playerIdx < players.length; playerIdx++) {
    const func = async () => {
      const schnorrApi = clients[playerIdx].Schnorr();

      const partialSignResult = await schnorrApi.sign(
        signSessionConfig,
        masterKeyId,
        derivationPath,
        message
      );

      partialSignatures.push(partialSignResult);
    };

    partialSignaturePromises.push(func());
  }

  await Promise.all(partialSignaturePromises);

  const schnorrApi = clients[0].Schnorr();

  const signature = await schnorrApi.finalizeSignature(
    message,
    partialSignatures
  );

  // Verify the signature relative to the signed message and the public key

  try {
    const result = await schnorrApi.verifySignature(
      publicKey,
      message,
      signature.signature
    );
    console.log(result);
    console.log(
      `Signature: ${Buffer.from(signature.signature).toString("hex")}`
    );
  } catch (e) {
    console.log(e);
  }
}

main().catch((e) => console.log(e));