package de.extio.lmlib.token;

import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpMethod;
import org.springframework.web.client.RestClient;

import de.extio.lmlib.profile.ModelProfile;

final class LlamaServerTokenizer implements Tokenizer {
	
	static final String NAME = "llamaServer";
	
	private static final Logger LOGGER = LoggerFactory.getLogger(LlamaServerTokenizer.class);
	
	private final RestClient.Builder restClientBuilder;
	
	public LlamaServerTokenizer(final RestClient.Builder restClientBuilder) {
		this.restClientBuilder = restClientBuilder;
	}
	
	@Override
	public String getName() {
		return NAME;
	}
	
	@Override
	public List<Long> tokenize(final String txt, final ModelProfile modelProfile) {
		LOGGER.debug("Requesting tokenizing");
		final var request = new TokenizeRequest(txt);
		final var restClient = this.restClientBuilder.baseUrl(modelProfile.url()).build();
		final var response = restClient
				.method(HttpMethod.POST)
				.uri("/tokenize")
				.body(request)
				.retrieve()
				.onStatus(status -> {
					if (status.getStatusCode().isError()) {
						LOGGER.error("Tokenization request failed with status code {} {}: {}", status.getStatusCode(), status.getStatusText(), new String(status.getBody().readAllBytes(), StandardCharsets.UTF_8));
						throw new RuntimeException("Tokenization request failed with status code " + status.getStatusCode() + " " + status.getStatusText());
					}
					return false;
				})
				.body(TokenizeResponse.class);
		return response.tokens();
	}

	@Override
	public int count(final String txt, final ModelProfile modelProfile) {
		return this.tokenize(txt, modelProfile).size();
	}
	
	@Override
	public String detokenize(final List<Long> tokens, final ModelProfile modelProfile) {
		LOGGER.debug("Requesting detokenizing");
		final var request = new DetokenizeRequest(tokens);
		final var restClient = this.restClientBuilder.baseUrl(modelProfile.url()).build();
		final var response = restClient
				.method(HttpMethod.POST)
				.uri("/detokenize")
				.body(request)
				.retrieve()
				.onStatus(status -> {
					if (status.getStatusCode().isError()) {
						LOGGER.error("Detokenization request failed with status code {} {}: {}", status.getStatusCode(), status.getStatusText(), new String(status.getBody().readAllBytes(), StandardCharsets.UTF_8));
						throw new RuntimeException("Detokenization request failed with status code " + status.getStatusCode() + " " + status.getStatusText());
					}
					return false;
				})				
				.body(DetokenizeResponse.class);
		return response.content();
	}
	
	private record TokenizeRequest(String content) {
	}
	
	private record TokenizeResponse(List<Long> tokens) {
	}
	
	private record DetokenizeRequest(List<Long> tokens) {
	}
	
	private record DetokenizeResponse(String content) {
	}
	
}