package org.studiorailgun.kobold;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.MalformedURLException;
import java.net.ProtocolException;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLConnection;
import java.nio.charset.StandardCharsets;
import java.util.UUID;

/**
 * A request to koboldcpp
 */
public class KoboldRequest {

    /**
     * Base Kobold URL
     */
    static final String BASE_URL = "http://localhost:5001";

    /**
     * Endpoint to start generating
     */
    static final String STREAM_ENDPOINT = "/api/v1/generate";

    /**
     * Endpoint to query generation status
     */
    static final String CHECK_ENDPOINT = "/api/extra/generate/check";

    /**
     * The prompt
     */
    String prompt;

    /**
     * The gen key for streaming
     */
    String genKey;

    /**
     * The response string
     */
    String response = "";

    /**
     * True if has received the complete generation response from kobold, false otherwise 
     */
    boolean completed = false;

    /**
     * Constructor
     * @param prompt The prompt
     */
    public KoboldRequest(String prompt){
        this.prompt = prompt;
        this.genKey = UUID.randomUUID().toString();
    }
    

    /**
     * Fires the request to generate
     */
    public void fire(){
        try {
            String urlPath = BASE_URL + STREAM_ENDPOINT;
            URL url = new URI(urlPath).toURL();
            URLConnection con = url.openConnection();
            HttpURLConnection http = (HttpURLConnection)con;
            http.setRequestMethod("POST"); // PUT is another valid option
            http.setDoOutput(true);

            String body = "{\"prompt\": \"" + this.prompt + "\", \"genkey\": \"" + genKey + "\"}";
            byte[] out = body.getBytes(StandardCharsets.UTF_8);
            int length = out.length;

            http.setFixedLengthStreamingMode(length);
            http.setRequestProperty("Content-Type", "application/json; charset=UTF-8");
            http.connect();
            try(OutputStream os = http.getOutputStream()) {
                os.write(out);
            }
            http.disconnect();
            //if uncommented, will block until response is fully generated
            // try(InputStream is = http.getInputStream()){
            //     byte[] response = is.readAllBytes();
            //     String responseParsed = new String(response, StandardCharsets.UTF_8);
            //     System.out.println("parsed: " + responseParsed);
            // }
        } catch (MalformedURLException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (ProtocolException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (URISyntaxException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    /**
     * Checks for new text generated from the prompt
     */
    public void check(){
        try {
            String urlPath = BASE_URL + CHECK_ENDPOINT;
            URL url = new URI(urlPath).toURL();
            URLConnection con = url.openConnection();
            HttpURLConnection http = (HttpURLConnection)con;
            http.setRequestMethod("POST"); // PUT is another valid option
            http.setDoOutput(true);
            http.setDoInput(true);
            String body = "{\"genkey\": \"" + genKey + "\"}";
            byte[] out = body.getBytes(StandardCharsets.UTF_8);
            int length = out.length;

            http.setFixedLengthStreamingMode(length);
            http.setRequestProperty("Content-Type", "application/json; charset=UTF-8");
            http.connect();
            try(OutputStream os = http.getOutputStream()) {
                os.write(out);
            }
            try(InputStream is = http.getInputStream()){
                byte[] responseBytes = is.readAllBytes();
                String responseParsed = new String(responseBytes, StandardCharsets.UTF_8);
                String responseBody = responseParsed.substring(23, responseParsed.length() - 4);
                if(this.response.equals(responseBody) && this.response.length() > 0){
                    this.completed = true;
                } else {
                    this.response = responseBody;
                }
            }
        } catch (MalformedURLException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (ProtocolException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (URISyntaxException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    /**
     * Gets the response
     * @return The response
     */
    public String getResponse(){
        return response;
    }

    /**
     * Gets whether this request has completed generation or not
     * @return true if has completed generation, false otherwise
     */
    public boolean isCompleted(){
        return completed;
    }

    /**
     * Flags the request as completed
     */
    public void markCompleted(){
        this.completed = true;
    }

    /**
     * Gets the prompt of the request
     * @return The prompt
     */
    public String getPrompt(){
        return prompt;
    }

}
