package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	reflectionpb "google.golang.org/grpc/reflection/grpc_reflection_v1"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/dynamicpb"
)

type result struct {
	Code     string          `json:"code"`
	Message  string          `json:"message,omitempty"`
	Response json.RawMessage `json:"response,omitempty"`
}

func main() {
	address := flag.String("address", "", "loopback gRPC address")
	method := flag.String("method", "", "full unary gRPC method")
	requestPath := flag.String("request", "", "protobuf JSON request file")
	fencePath := flag.String("before-invoke-fence", "", "atomic marker to write immediately before the RPC")
	fencePayloadPath := flag.String("before-invoke-fence-payload", "", "marker payload file")
	timeout := flag.Duration("timeout", 5*time.Minute, "RPC timeout")
	flag.Parse()

	if err := run(*address, *method, *requestPath, *fencePath, *fencePayloadPath, *timeout); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run(address, fullMethod, requestPath, fencePath, fencePayloadPath string, timeout time.Duration) error {
	if err := validateAddress(address); err != nil {
		return err
	}
	serviceName, methodName, err := parseMethod(fullMethod)
	if err != nil {
		return err
	}
	if (fencePath == "") != (fencePayloadPath == "") {
		return errors.New("--before-invoke-fence and --before-invoke-fence-payload must be provided together")
	}
	requestJSON, err := os.ReadFile(requestPath)
	if err != nil {
		return fmt.Errorf("read request: %w", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	conn, err := grpc.NewClient(address, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return fmt.Errorf("create gRPC client: %w", err)
	}
	defer conn.Close()

	method, err := reflectMethod(ctx, conn, serviceName, methodName)
	if err != nil {
		err = fmt.Errorf("gRPC reflection at %s for %s is unavailable: %w", address, serviceName, err)
		if writeStatusError(err) {
			return nil
		}
		return err
	}
	request := dynamicpb.NewMessage(method.Input())
	if err := protojson.Unmarshal(requestJSON, request); err != nil {
		return fmt.Errorf("decode request for %s: %w", fullMethod, err)
	}
	response := dynamicpb.NewMessage(method.Output())
	if fencePath != "" {
		if err := writeInvokeFence(fencePath, fencePayloadPath); err != nil {
			return err
		}
	}
	if err := conn.Invoke(ctx, fullMethod, request, response); err != nil {
		if writeStatusError(err) {
			return nil
		}
		return err
	}
	responseJSON, err := protojson.Marshal(response)
	if err != nil {
		return fmt.Errorf("encode response for %s: %w", fullMethod, err)
	}
	return json.NewEncoder(os.Stdout).Encode(result{Code: "OK", Response: responseJSON})
}

func writeInvokeFence(path, payloadPath string) error {
	payload, err := os.ReadFile(payloadPath)
	if err != nil {
		return fmt.Errorf("read invoke fence payload: %w", err)
	}
	if !json.Valid(payload) {
		return errors.New("invoke fence payload must be valid JSON")
	}
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return fmt.Errorf("create invoke fence directory: %w", err)
	}
	temporary, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*")
	if err != nil {
		return fmt.Errorf("create invoke fence: %w", err)
	}
	temporaryPath := temporary.Name()
	defer os.Remove(temporaryPath)
	if err := temporary.Chmod(0o600); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("protect invoke fence: %w", err)
	}
	if _, err := temporary.Write(payload); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("write invoke fence: %w", err)
	}
	if err := temporary.Sync(); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("sync invoke fence: %w", err)
	}
	if err := temporary.Close(); err != nil {
		return fmt.Errorf("close invoke fence: %w", err)
	}
	if err := os.Rename(temporaryPath, path); err != nil {
		return fmt.Errorf("commit invoke fence: %w", err)
	}
	return nil
}

func validateAddress(address string) error {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return fmt.Errorf("invalid --address: %w", err)
	}
	ip := net.ParseIP(host)
	if ip == nil || !ip.IsLoopback() || port == "" {
		return errors.New("--address must be a loopback host and port")
	}
	return nil
}

func parseMethod(fullMethod string) (serviceName, methodName string, err error) {
	service, method, ok := strings.Cut(strings.TrimPrefix(fullMethod, "/"), "/")
	if !ok || service == "" || method == "" || strings.Contains(method, "/") {
		return "", "", errors.New("--method must be /package.Service/Method")
	}
	return service, method, nil
}

func reflectMethod(
	ctx context.Context,
	conn *grpc.ClientConn,
	serviceName string,
	methodName string,
) (protoreflect.MethodDescriptor, error) {
	stream, err := reflectionpb.NewServerReflectionClient(conn).ServerReflectionInfo(ctx)
	if err != nil {
		return nil, fmt.Errorf("open reflection stream: %w", err)
	}
	if err := stream.Send(&reflectionpb.ServerReflectionRequest{
		MessageRequest: &reflectionpb.ServerReflectionRequest_FileContainingSymbol{
			FileContainingSymbol: serviceName,
		},
	}); err != nil {
		return nil, fmt.Errorf("request reflection descriptor: %w", err)
	}
	response, err := stream.Recv()
	if err != nil {
		return nil, fmt.Errorf("receive reflection descriptor: %w", err)
	}
	if reflectionErr := response.GetErrorResponse(); reflectionErr != nil {
		return nil, fmt.Errorf("reflection error %d: %s", reflectionErr.ErrorCode, reflectionErr.ErrorMessage)
	}
	descriptorResponse := response.GetFileDescriptorResponse()
	if descriptorResponse == nil {
		return nil, errors.New("reflection response did not contain file descriptors")
	}

	set := &descriptorpb.FileDescriptorSet{}
	for _, encoded := range descriptorResponse.FileDescriptorProto {
		file := &descriptorpb.FileDescriptorProto{}
		if err := proto.Unmarshal(encoded, file); err != nil {
			return nil, fmt.Errorf("decode reflected file descriptor: %w", err)
		}
		set.File = append(set.File, file)
	}
	files, err := protodesc.NewFiles(set)
	if err != nil {
		return nil, fmt.Errorf("link reflected file descriptors: %w", err)
	}
	descriptor, err := files.FindDescriptorByName(protoreflect.FullName(serviceName))
	if err != nil {
		return nil, fmt.Errorf("find reflected service %s: %w", serviceName, err)
	}
	service, ok := descriptor.(protoreflect.ServiceDescriptor)
	if !ok {
		return nil, fmt.Errorf("reflected symbol %s is not a service", serviceName)
	}
	method := service.Methods().ByName(protoreflect.Name(methodName))
	if method == nil {
		return nil, fmt.Errorf("reflected service %s has no method %s", serviceName, methodName)
	}
	if method.IsStreamingClient() || method.IsStreamingServer() {
		return nil, fmt.Errorf("reflected method %s/%s is not unary", serviceName, methodName)
	}
	return method, nil
}

func writeStatusError(err error) bool {
	grpcStatus, ok := status.FromError(err)
	if !ok {
		return false
	}
	_ = json.NewEncoder(os.Stdout).Encode(result{
		Code:    grpcStatus.Code().String(),
		Message: err.Error(),
	})
	return true
}