package copilot

import (
	"bytes"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

// GetEngine returns information about a specific engine.
func GetEngine(c *gin.Context) {
	modelName := os.Getenv("CODEX_API_MODEL_NAME")
	c.JSON(http.StatusOK, gin.H{
		"id":      c.Param("model-name"),
		"object":  "engine",
		"owner":   "openai",
		"ready":   true,
		"model":   modelName,
		"max_tokens": 4096,
	})
}


// CodeCompletions code代码补全接口
func CodeCompletions(c *gin.Context) {
	ctx := c.Request.Context()

	body, err := io.ReadAll(c.Request.Body)
	if nil != err {
		c.AbortWithStatus(http.StatusBadRequest)
		return
	}

	codexAPIURL := os.Getenv("CODEX_API_BASE")
	apiKey := os.Getenv("CODEX_API_KEY")
	modelName := os.Getenv("CODEX_API_MODEL_NAME")
	maxTokens, _ := strconv.Atoi(os.Getenv("CODEX_MAX_TOKENS"))
	temperature, _ := strconv.ParseFloat(os.Getenv("CODEX_TEMPERATURE"), 64)

	// 如果设置了温度为-1, 则跟随插件的设置
	if temperature != -1 {
		body, _ = sjson.SetBytes(body, "temperature", temperature)
	}

	body, _ = sjson.SetBytes(body, "model", modelName)
	if int(gjson.GetBytes(body, "max_tokens").Int()) > maxTokens {
		body, _ = sjson.SetBytes(body, "max_tokens", maxTokens)
	}

	limitPrompt, _ := strconv.Atoi(os.Getenv("CODEX_LIMIT_PROMPT"))
	if limitPrompt > 0 {
		prompt := gjson.GetBytes(body, "prompt").String()
		promptLines := strings.Split(prompt, "\n")
		if len(promptLines) > limitPrompt {
			prompt = strings.Join(promptLines[len(promptLines)-limitPrompt:], "\n")
			body, _ = sjson.SetBytes(body, "prompt", prompt)
		}
	}

	if gjson.GetBytes(body, "n").Int() > 1 {
		body, _ = sjson.SetBytes(body, "n", 1)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, codexAPIURL, io.NopCloser(bytes.NewBuffer(body)))
	if nil != err {
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}

	req.Header.Set("Content-Type", "application/json")
	if strings.Contains(apiKey, " ") {
		split := strings.Split(apiKey, " ")
		req.Header.Set(split[0], split[1])
	} else {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	httpClientTimeout, _ := time.ParseDuration(os.Getenv("HTTP_CLIENT_TIMEOUT") + "s")
	client := &http.Client{
		Timeout: httpClientTimeout,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		},
	}
	resp, err := client.Do(req)
	if nil != err {
		if errors.Is(err, context.Canceled) {
			c.AbortWithStatus(http.StatusRequestTimeout)
			return
		}

		log.Println("request completions failed:", err.Error())
		c.AbortWithStatus(http.StatusInternalServerError)
		return
	}
	defer CloseIO(resp.Body)

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		log.Println("request completions failed:", string(body))

		resp.Body = io.NopCloser(bytes.NewBuffer(body))
	}

	c.Status(resp.StatusCode)
	_, _ = io.Copy(c.Writer, resp.Body)
}