package runpod

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

var Endpoint string
var Authorization string

type Request struct {
	Id     string
	Status string
}

type StatusResp struct {
	DelayTime     int                 `json:"delayTime"`
	ExecutionTime int                 `json:"executionTime"`
	Id            string              `json:"id"`
	Status        string              `json:"status"`
	Output        []map[string]string `json:"output"`
}

func postReq(data []byte, path string) (resp *http.Response, err error) {
	req, err := http.NewRequest("POST", Endpoint+path, bytes.NewBuffer(data))
	if err != nil {
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Authorization", Authorization)

	client := &http.Client{}
	resp, err = client.Do(req)

	statusOK := resp.StatusCode >= 200 && resp.StatusCode < 300
	if !statusOK {
		return resp, fmt.Errorf("runpod '%s' resp HTTP code '%d'", path, resp.StatusCode)
	}
	return
}

func getReq(path string) (resp *http.Response, err error) {
	req, err := http.NewRequest("GET", Endpoint+path, nil)
	if err != nil {
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Authorization", Authorization)

	client := &http.Client{}
	resp, err = client.Do(req)

	statusOK := resp.StatusCode >= 200 && resp.StatusCode < 300
	if !statusOK {
		return resp, fmt.Errorf("runpod '%s' resp HTTP code '%d'", path, resp.StatusCode)
	}

	return
}

// Returns the status of the serverless endpoint, can be used to check if it's operational
// returns a map because I'm too lazy to make the proper struct
func Health() (res map[string]interface{}, err error) {
	resp, err := getReq("health")
	if err != nil {
		return
	}
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	return
}

// Return the Status of a job
func Status(id string) (status StatusResp, err error) {
	resp, err := getReq("status/" + id)
	if err != nil {
		return
	}

	err = json.NewDecoder(resp.Body).Decode(&status)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	return
}

// Run a job
func Run(data []byte) (req Request, err error) {
	resp, err := postReq(data, "run")
	if err != nil {
		return
	}

	err = json.NewDecoder(resp.Body).Decode(&req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	return
}

// Poll job until it succeeded or failed and return it
func Poll(id string, delay time.Duration) (status StatusResp, err error) {
	for {
		status, err = Status(id)
		if err != nil {
			return
		}

		switch status.Status {
		case "COMPLETED":
			return
		case "IN_QUEUE", "IN_PROGRESS":
			time.Sleep(delay)
			continue
		default:
			return status, fmt.Errorf("job failed with status '%s'", status.Status)
		}
	}
}