package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strconv"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/encoding/protowire"
)

type wireMessage []byte

type wireCodec struct{}

func (wireCodec) Name() string { return "proto" }

func (wireCodec) Marshal(value any) ([]byte, error) {
	message, ok := value.(*wireMessage)
	if !ok {
		return nil, errors.New("unexpected gRPC message")
	}
	return *message, nil
}

func (wireCodec) Unmarshal(data []byte, value any) error {
	message, ok := value.(*wireMessage)
	if !ok {
		return errors.New("unexpected gRPC message")
	}
	*message = append((*message)[:0], data...)
	return nil
}

func responseJWKS(data []byte) (string, error) {
	for len(data) > 0 {
		number, wireType, size := protowire.ConsumeTag(data)
		if size < 0 {
			return "", protowire.ParseError(size)
		}
		data = data[size:]
		if number == 1 && wireType == protowire.BytesType {
			jwks, consumed := protowire.ConsumeString(data)
			if consumed < 0 {
				return "", protowire.ParseError(consumed)
			}
			var set struct {
				Keys []json.RawMessage `json:"keys"`
			}
			if json.Unmarshal([]byte(jwks), &set) != nil || len(set.Keys) == 0 {
				return "", errors.New("account returned invalid JWKS")
			}
			return jwks, nil
		}
		consumed := protowire.ConsumeFieldValue(number, wireType, data)
		if consumed < 0 {
			return "", protowire.ParseError(consumed)
		}
		data = data[consumed:]
	}
	return "", errors.New("account response has no JWKS")
}

func main() {
	if len(os.Args) != 4 {
		fmt.Fprintln(os.Stderr, "usage: account-jwks <address> <org-id> <cluster-id>")
		os.Exit(2)
	}
	orgID, err := strconv.ParseUint(os.Args[2], 10, 64)
	if err != nil {
		fmt.Fprintln(os.Stderr, "invalid org ID:", err)
		os.Exit(2)
	}
	clusterID, err := strconv.ParseUint(os.Args[3], 10, 64)
	if err != nil {
		fmt.Fprintln(os.Stderr, "invalid cluster ID:", err)
		os.Exit(2)
	}

	var request wireMessage
	request = protowire.AppendTag(request, 1, protowire.VarintType)
	request = protowire.AppendVarint(request, orgID)
	request = protowire.AppendTag(request, 3, protowire.VarintType)
	request = protowire.AppendVarint(request, clusterID)
	request = protowire.AppendTag(request, 4, protowire.VarintType)
	request = protowire.AppendVarint(request, protowire.EncodeBool(true))

	connection, err := grpc.NewClient(
		os.Args[1],
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithDefaultCallOptions(grpc.ForceCodec(wireCodec{})),
	)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	var response wireMessage
	invokeErr := connection.Invoke(
		ctx,
		"/account.AccountAPIService/GetClusterAuthPubKey",
		&request,
		&response,
	)
	cancel()
	_ = connection.Close()
	if invokeErr != nil {
		fmt.Fprintln(os.Stderr, invokeErr)
		os.Exit(1)
	}
	jwks, err := responseJWKS(response)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	fmt.Print(jwks)
}