import ai.protomolt.proto.projection.MessageProjection;
import ai.protomolt.proto.projection.SourceResolver;
import ai.protomolt.proto.validate.ProtoValidator;
import ai.protomolt.proto.search.index.spi.ProtoOptionsIndexingHintSource;
import ai.protomolt.proto.search.chunk.PolicyDerivation;
import ai.protomolt.proto.search.embedding.EmbeddingProvider;
import com.google.protobuf.*;
import com.google.protobuf.util.JsonFormat;
import com.google.gson.*;
import io.grpc.*;
import io.grpc.protobuf.ProtoUtils;
import io.grpc.stub.ClientCalls;
import io.grpc.stub.ServerCalls;
import demo.CourtTypes.*;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.TimeUnit;

/** A bounded local gRPC demonstration. Requires an already-running TEI server. */
public class ContractDemo {
  static final Gson JSON = new GsonBuilder().setPrettyPrinting().create();
  static <A extends Message,B extends Message> MethodDescriptor<A,B> method(String name,A a,B b) {
    return MethodDescriptor.<A,B>newBuilder().setType(MethodDescriptor.MethodType.UNARY)
        .setFullMethodName(name).setRequestMarshaller(ProtoUtils.marshaller(a))
        .setResponseMarshaller(ProtoUtils.marshaller(b)).build();
  }
  static <A,B> B call(Channel channel,MethodDescriptor<A,B> method,A value) {
    return ClientCalls.blockingUnaryCall(channel,method,CallOptions.DEFAULT.withDeadlineAfter(30,TimeUnit.SECONDS),value);
  }
  static void require(boolean ok,String text) { if(!ok) throw new IllegalStateException(text); }
  static JsonElement json(Message message) throws Exception {
    return JsonParser.parseString(JsonFormat.printer().preservingProtoFieldNames().print(message));
  }
  public static void main(String[] args) throws Exception {
    // Read the vendored TEI wire contract compiled with protoc alongside the demo.
    var set=DescriptorProtos.FileDescriptorSet.parseFrom(Files.readAllBytes(Path.of(args[0])));
    var linked=new HashMap<String,Descriptors.FileDescriptor>();
    for(var file:set.getFileList()) linked.put(file.getName(),Descriptors.FileDescriptor.buildFrom(file,
        file.getDependencyList().stream().map(linked::get).toArray(Descriptors.FileDescriptor[]::new)));
    var tei=linked.values().stream().filter(f->f.getPackage().equals("tei.v1")).findFirst().orElseThrow();
    var infoRequest=DynamicMessage.getDefaultInstance(tei.findMessageTypeByName("InfoRequest"));
    var infoResponse=DynamicMessage.getDefaultInstance(tei.findMessageTypeByName("InfoResponse"));
    var embedRequest=DynamicMessage.getDefaultInstance(tei.findMessageTypeByName("EmbedRequest"));
    var embedResponse=DynamicMessage.getDefaultInstance(tei.findMessageTypeByName("EmbedResponse"));
    var teiChannel=ManagedChannelBuilder.forTarget(args[1]).usePlaintext().build();
    var info=call(teiChannel,method("tei.v1.Info/Info",infoRequest,infoResponse),infoRequest);
    String model=(String)info.getField(info.getDescriptorForType().findFieldByName("model_id"));
    var read=method("website.court.v1.CaseSource/Read",CaseKey.getDefaultInstance(),Case.getDefaultInstance());
    var accept=method("website.court.v1.DocumentSink/Accept",SearchDoc.getDefaultInstance(),Receipt.getDefaultInstance());
    var validator=ProtoValidator.forMessageType(SearchDoc.getDescriptor());
    var source=Case.newBuilder().setCaseId("24-cv-00117").setStyleOfCause("State v. Example").build();
    var server=io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder.forAddress(new java.net.InetSocketAddress("127.0.0.1",0))
      .addService(ServerServiceDefinition.builder("website.court.v1.CaseSource")
        .addMethod(read,ServerCalls.asyncUnaryCall((request,observer)->{
          if(!request.getCaseId().equals(source.getCaseId())) { observer.onError(Status.NOT_FOUND.asRuntimeException());return; }
          observer.onNext(source);observer.onCompleted();
        })).build())
      .addService(ServerServiceDefinition.builder("website.court.v1.DocumentSink")
        .addMethod(accept,ServerCalls.asyncUnaryCall((request,observer)->{
          var result=validator.validate(request);
          if(!result.valid()) { observer.onError(Status.INVALID_ARGUMENT.withDescription(result.violations().toString()).asRuntimeException());return; }
          observer.onNext(Receipt.newBuilder().setDocId(request.getDocId()).build());observer.onCompleted();
        })).build()).build().start();
    var channel=ManagedChannelBuilder.forAddress("127.0.0.1",server.getPort()).usePlaintext().build();
    try {
      var input=call(channel,read,CaseKey.newBuilder().setCaseId(source.getCaseId()).build());
      var projection=MessageProjection.forTarget(SearchDoc.getDescriptor(),SourceResolver.of(Case.getDescriptor())).orElseThrow();
      var projected=SearchDoc.parseFrom(projection.project(input).toByteArray());
      require(projected.getDocId().equals("24-cv-00117")&&projected.getTitle().equals("State v. Example"),"Projection differs from fixture");
      require(validator.validate(projected).valid(),"Valid fixture was rejected");
      var invalid=projected.toBuilder().setTitle("X").build();
      var failure=validator.validate(invalid);
      require(!failure.valid(),"Short title should fail");
      String refused="";
      try { call(channel,accept,invalid);throw new IllegalStateException("Invalid RPC accepted"); }
      catch(StatusRuntimeException e) { require(e.getStatus().getCode()==Status.Code.INVALID_ARGUMENT,"Unexpected gRPC status");refused=e.getStatus().getCode().name(); }
      var receipt=call(channel,accept,projected);
      var hints=new ProtoOptionsIndexingHintSource();
      var idHint=hints.resolve(SearchDoc.getDescriptor().findFieldByName("doc_id")).orElseThrow();
      var titleHint=hints.resolve(SearchDoc.getDescriptor().findFieldByName("title")).orElseThrow();
      var policy=titleHint.chunkingPolicy();
      require(policy.embedding().model().equals(model),"Configured TEI model differs from contract");
      EmbeddingProvider provider=new EmbeddingProvider() {
        public String providerId(){return model;}
        public int dimension(){return policy.embedding().dims();}
        public float[] embed(String text){
          var descriptor=embedRequest.getDescriptorForType();
          var req=embedRequest.toBuilder().setField(descriptor.findFieldByName("inputs"),text)
            .setField(descriptor.findFieldByName("truncate"),false).setField(descriptor.findFieldByName("normalize"),false).build();
          var response=call(teiChannel,method("tei.v1.Embed/Embed",embedRequest,embedResponse),req);
          var values=(List<?>)response.getField(response.getDescriptorForType().findFieldByName("embeddings"));
          float[] result=new float[values.size()];for(int i=0;i<result.length;i++)result[i]=((Number)values.get(i)).floatValue();return result;
        }
      };
      var chunks=new PolicyDerivation(provider).derive(projected.getTitle(),policy);
      require(chunks.size()==1,"Expected one title chunk");
      var vector=chunks.getFirst().vector();require(vector.length==384,"Expected 384 actual coordinates");
      double norm=0;for(float value:vector){require(Float.isFinite(value),"Nonfinite vector");norm+=(double)value*value;}norm=Math.sqrt(norm);
      require(Math.abs(norm-1)<1e-5,"Vector not L2 normalized");
      var output=new LinkedHashMap<String,Object>();
      output.put("source",json(input));output.put("projected",json(projected));output.put("valid",true);
      output.put("invalidInput",json(invalid));output.put("violations",failure.violations());output.put("invalidGrpcStatus",refused);
      output.put("receipt",json(receipt));output.put("indexFields",Map.of("doc_id",idHint.type().name(),"title",titleHint.type().name()));
      output.put("policy",policy);output.put("modelServer",json(info));output.put("embeddingInput",chunks.getFirst().chunk().text());
      output.put("vector",vector);output.put("vectorNorm",norm);output.put("chunkCount",chunks.size());
      output.put("executed",List.of("loopback CaseSource.Read RPC","MessageProjection","ProtoValidator valid and invalid cases","loopback DocumentSink.Accept success and rejection","ProtoOptionsIndexingHintSource","PolicyDerivation chunking and L2 normalization","real TEI gRPC model inference"));
      output.put("notExecuted",List.of("persistent search-index write","production service deployment"));
      Files.writeString(Path.of(args[2]),JSON.toJson(output)+"\n");
      System.out.println("PASS: projection, validation, two loopback gRPC methods, real TEI vector (384), norm="+norm);
    } finally { channel.shutdownNow();teiChannel.shutdownNow();server.shutdownNow();server.awaitTermination(5,TimeUnit.SECONDS); }
  }
}