package com.example.support;

import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.FileDescriptor;
import com.google.protobuf.Descriptors.MethodDescriptor;
import com.google.protobuf.Descriptors.ServiceDescriptor;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.util.JsonFormat;
import io.grpc.CallOptions;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.protobuf.ProtoUtils;
import io.grpc.reflection.v1alpha.ServerReflectionGrpc;
import io.grpc.reflection.v1alpha.ServerReflectionRequest;
import io.grpc.reflection.v1alpha.ServerReflectionResponse;
import io.grpc.stub.ClientCalls;
import io.grpc.stub.StreamObserver;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Owned gRPC test client (L3git). Makes a dynamic UNARY call by name - no generated stubs and no
 * .proto file needed - using the server's own reflection service (the same mechanism grpcurl uses):
 * fetch the descriptors, parse a JSON request into a {@link DynamicMessage}, invoke the method, and
 * return the response as JSON. The grpc-java + protobuf dependencies are added to the project's pom
 * only when a test uses a gRPC step. The target server must have the reflection service enabled.
 * Generated / shipped by L3git.
 */
public class GrpcClient {

    /** target = host:port; fullMethod = "package.Service/Method"; jsonRequest = the request as JSON. */
    public String callJson(String target, String fullMethod, String jsonRequest) throws Exception {
        int slash = fullMethod.lastIndexOf('/');
        if (slash < 0) { throw new IllegalArgumentException("gRPC method must be 'package.Service/Method': " + fullMethod); }
        String serviceName = fullMethod.substring(0, slash);
        String methodName = fullMethod.substring(slash + 1);

        ManagedChannel channel = ManagedChannelBuilder.forTarget(target).usePlaintext().build();
        try {
            Map<String, FileDescriptorProto> protos = reflectFileDescriptors(channel, serviceName);
            FileDescriptor[] built = buildAll(protos);

            ServiceDescriptor service = null;
            for (FileDescriptor fd : built) {
                ServiceDescriptor s = fd.findServiceByName(simpleName(serviceName));
                if (s != null && s.getFullName().equals(serviceName)) { service = s; break; }
            }
            if (service == null) { throw new IllegalStateException("gRPC service not found via reflection: " + serviceName); }
            MethodDescriptor method = service.findMethodByName(methodName);
            if (method == null) { throw new IllegalStateException("gRPC method not found: " + methodName); }

            Descriptor inputType = method.getInputType();
            Descriptor outputType = method.getOutputType();

            DynamicMessage.Builder reqBuilder = DynamicMessage.newBuilder(inputType);
            JsonFormat.parser().ignoringUnknownFields().merge(jsonRequest == null ? "{}" : jsonRequest, reqBuilder);
            DynamicMessage request = reqBuilder.build();

            io.grpc.MethodDescriptor<DynamicMessage, DynamicMessage> grpcMethod =
                    io.grpc.MethodDescriptor.<DynamicMessage, DynamicMessage>newBuilder()
                            .setType(io.grpc.MethodDescriptor.MethodType.UNARY)
                            .setFullMethodName(serviceName + "/" + methodName)
                            .setRequestMarshaller(ProtoUtils.marshaller(DynamicMessage.getDefaultInstance(inputType)))
                            .setResponseMarshaller(ProtoUtils.marshaller(DynamicMessage.getDefaultInstance(outputType)))
                            .build();

            DynamicMessage response = ClientCalls.blockingUnaryCall(channel, grpcMethod, CallOptions.DEFAULT, request);
            return JsonFormat.printer().includingDefaultValueFields().print(response);
        } finally {
            channel.shutdownNow();
        }
    }

    // ---- server reflection: gather the FileDescriptorProtos for a symbol + its dependencies ------

    private Map<String, FileDescriptorProto> reflectFileDescriptors(ManagedChannel channel, String symbol) throws Exception {
        ServerReflectionGrpc.ServerReflectionStub stub = ServerReflectionGrpc.newStub(channel);
        Map<String, FileDescriptorProto> byName = new HashMap<>();
        List<String> pending = new ArrayList<>();
        pending.add(symbol);
        // First round asks by symbol; subsequent rounds resolve dependency file names.
        boolean bySymbol = true;
        while (!pending.isEmpty()) {
            String query = pending.remove(0);
            List<FileDescriptorProto> got = oneReflection(stub, query, bySymbol);
            for (FileDescriptorProto fdp : got) {
                if (byName.containsKey(fdp.getName())) { continue; }
                byName.put(fdp.getName(), fdp);
                for (String dep : fdp.getDependencyList()) {
                    if (!byName.containsKey(dep)) { pending.add(dep); }
                }
            }
            bySymbol = false; // after the first, resolve by file name
        }
        return byName;
    }

    private List<FileDescriptorProto> oneReflection(ServerReflectionGrpc.ServerReflectionStub stub,
                                                    String query, boolean bySymbol) throws Exception {
        final List<FileDescriptorProto> result = new ArrayList<>();
        final CountDownLatch done = new CountDownLatch(1);
        final AtomicReference<Throwable> err = new AtomicReference<>();

        StreamObserver<ServerReflectionResponse> respObs = new StreamObserver<ServerReflectionResponse>() {
            @Override public void onNext(ServerReflectionResponse resp) {
                try {
                    for (com.google.protobuf.ByteString bs : resp.getFileDescriptorResponse().getFileDescriptorProtoList()) {
                        result.add(FileDescriptorProto.parseFrom(bs));
                    }
                } catch (Exception e) { err.set(e); }
            }
            @Override public void onError(Throwable t) { err.set(t); done.countDown(); }
            @Override public void onCompleted() { done.countDown(); }
        };
        StreamObserver<ServerReflectionRequest> reqObs = stub.serverReflectionInfo(respObs);
        ServerReflectionRequest.Builder rb = ServerReflectionRequest.newBuilder();
        if (bySymbol) { rb.setFileContainingSymbol(query); } else { rb.setFileByFilename(query); }
        reqObs.onNext(rb.build());
        reqObs.onCompleted();
        if (!done.await(20, TimeUnit.SECONDS)) { throw new IllegalStateException("gRPC reflection timed out for " + query); }
        if (err.get() != null) { throw new Exception("gRPC reflection failed for " + query + ": " + err.get().getMessage()); }
        return result;
    }

    // ---- build FileDescriptors from the proto graph (dependencies first) ------------------------

    private FileDescriptor[] buildAll(Map<String, FileDescriptorProto> protos) throws Exception {
        Map<String, FileDescriptor> built = new HashMap<>();
        List<FileDescriptor> order = new ArrayList<>();
        for (String name : protos.keySet()) { buildOne(name, protos, built); }
        order.addAll(built.values());
        return order.toArray(new FileDescriptor[0]);
    }

    private FileDescriptor buildOne(String name, Map<String, FileDescriptorProto> protos,
                                    Map<String, FileDescriptor> built) throws Exception {
        if (built.containsKey(name)) { return built.get(name); }
        FileDescriptorProto fdp = protos.get(name);
        if (fdp == null) { throw new IllegalStateException("missing gRPC descriptor dependency: " + name); }
        List<FileDescriptor> deps = new ArrayList<>();
        for (String dep : fdp.getDependencyList()) { deps.add(buildOne(dep, protos, built)); }
        FileDescriptor fd = FileDescriptor.buildFrom(fdp, deps.toArray(new FileDescriptor[0]));
        built.put(name, fd);
        return fd;
    }

    private static String simpleName(String fullName) {
        int dot = fullName.lastIndexOf('.');
        return dot < 0 ? fullName : fullName.substring(dot + 1);
    }
}
