/*
 *  LanguageTool, a natural language style checker
 *  * Copyright (C) 2018 Fabian Richter
 *  *
 *  * This library is free software; you can redistribute it and/or
 *  * modify it under the terms of the GNU Lesser General Public
 *  * License as published by the Free Software Foundation; either
 *  * version 2.1 of the License, or (at your option) any later version.
 *  *
 *  * This library is distributed in the hope that it will be useful,
 *  * but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *  * Lesser General Public License for more details.
 *  *
 *  * You should have received a copy of the GNU Lesser General Public
 *  * License along with this library; if not, write to the Free Software
 *  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301
 *  * USA
 *
 */
package org.languagetool.languagemodel.bert;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import io.grpc.ManagedChannel;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.shaded.io.grpc.netty.NegotiationType;
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
import io.grpc.netty.shaded.io.netty.handler.ssl.SslContextBuilder;
import org.jetbrains.annotations.Nullable;
import org.languagetool.languagemodel.bert.grpc.BertLmGrpc;
import org.languagetool.languagemodel.bert.grpc.BertLmGrpc.BertLmBlockingStub;

import javax.net.ssl.SSLException;
import java.io.File;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

import static org.languagetool.languagemodel.bert.grpc.BertLmProto.*;

public class RemoteLanguageModel {

  private final BertLmGrpc.BertLmBlockingStub model;
  private final ManagedChannel channel;
  private final Cache<Request, List<Double>> cache = CacheBuilder.newBuilder()
    .maximumSize(1000)
    .build();

  public static class Request {
    public String text;
    public int start;
    public int end;
    public List<String> candidates;

    public Request(String text, int start, int end, List<String> candidates) {
      this.text = text;
      this.start = start;
      this.end = end;
      this.candidates = candidates;
    }

    public ScoreRequest convert() {
      List<Mask> masks = Arrays.asList(Mask.newBuilder()
        .setStart(start)
        .setEnd(end)
        .addAllCandidates(candidates)
        .build());
      return ScoreRequest.newBuilder().setText(text).addAllMask(masks).build();
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;
      Request request = (Request) o;
      return start == request.start &&
        end == request.end &&
        text.equals(request.text) &&
        candidates.equals(request.candidates);
    }

    @Override
    public int hashCode() {
      return Objects.hash(text, start, end, candidates);
    }

  }

  public RemoteLanguageModel(String host, int port, boolean useSSL,
                             @Nullable String clientPrivateKey, @Nullable  String clientCertificate,
                             @Nullable String rootCertificate) throws SSLException {
    // TODO configure deadline/retries/... here?
    channel = getChannel(host, port, useSSL, clientPrivateKey, clientCertificate, rootCertificate);
    model = BertLmGrpc.newBlockingStub(channel);
  }

  private ManagedChannel getChannel(String host, int port, boolean useSSL,
                    @Nullable String clientPrivateKey, @Nullable  String clientCertificate,
                    @Nullable String rootCertificate) throws SSLException {
    NettyChannelBuilder channelBuilder = NettyChannelBuilder.forAddress(host, port);
    if (useSSL) {
      SslContextBuilder sslContextBuilder = GrpcSslContexts.forClient();
      if (rootCertificate != null) {
        sslContextBuilder.trustManager(new File(rootCertificate));
      }
      if (clientCertificate != null && clientPrivateKey != null) {
        sslContextBuilder.keyManager(new File(clientCertificate), new File(clientPrivateKey));
      }
      channelBuilder = channelBuilder.negotiationType(NegotiationType.TLS).sslContext(sslContextBuilder.build());
    } else {
      channelBuilder = channelBuilder.usePlaintext();
    }
    return channelBuilder.build();
  }

  public void shutdown() {
    if (channel != null) {
      channel.shutdownNow();
    }
  }

  public List<List<Double>> batchScore(List<Request> requests, long timeoutMilliseconds) throws TimeoutException {
    Map<Request, List<Double>> cachedRequests = new HashMap<>();
    List<Request> uncachedRequests = new ArrayList<>();
    for (Request request : requests) {
      List<Double> result = cache.getIfPresent(request);
      if (result == null) {
        uncachedRequests.add(request);
      } else {
        cachedRequests.put(request, result);
      }
    }
    BatchScoreRequest batch = BatchScoreRequest.newBuilder().addAllRequests(
      uncachedRequests.stream().map(Request::convert).collect(Collectors.toList())
    ).build();
    // TODO multiple masks
    List<List<Double>> nonCacheResult;
    try {
      BertLmBlockingStub stub;
      if (timeoutMilliseconds > 0) {
        stub = model.withDeadlineAfter(timeoutMilliseconds, TimeUnit.MILLISECONDS);
      } else {
        stub = model;
      }
      nonCacheResult = stub.batchScore(batch)
        .getResponsesList().stream().map(r ->
          r.getScoresList().get(0).getScoreList()).collect(Collectors.toList());
    } catch (StatusRuntimeException e) {
      if (e.getStatus().getCode() == Status.DEADLINE_EXCEEDED.getCode()) {
        throw new TimeoutException(e.getMessage());
      } else {
        throw e;
      }
    }
    //
    List<List<Double>> allResults = new ArrayList<>();
    int i = 0;
    for (Request request : requests) {
      List<Double> result = cachedRequests.get(request);
      if (result != null) {
        //System.out.println("Adding result from cache");
        allResults.add(result);
      } else {
        //System.out.println("Adding result from remote");
        allResults.add(nonCacheResult.get(i++));
      }
    }

    int j = 0;
    for (List<Double> re : nonCacheResult) {
      // a CacheLoader doesn't work with batching, so add manually:
      //System.out.println("Adding request to cache");
      cache.put(uncachedRequests.get(j), re);
      j++;
    }
    return allResults;
  }

  public List<Double> score(Request req) {
    // TODO deal with max seq length, extract windows
    // TODO mask multiple tokens in a sentence
    // TODO multiple masks
    return model.score(req.convert()).getScoresList().get(0).getScoreList();
  }

}
