package gpt

import (
	"context"
	"fmt"

	"github.com/juju/errors"
	"github.com/pingcap/log"
	gogpt "github.com/sashabaranov/go-openai"
	"go.uber.org/zap"
)

var ModelAllowedMaxTokens = map[string]int{
	"text-davinci-003":   4097,
	"code-davinci-002":   8001,
	"gpt-3.5-turbo":      4096,
	"gpt-3.5-turbo-0301": 4096,
	"gpt-4":              8192,
	"gpt-4-0314":         8192,
	"gpt-4-32k":          32768,
	"gpt-4-32k-0314":     32768,
}

var DefaultMaxOutputTokens = map[string]int{
	"text-davinci-003":   2048,
	"code-davinci-002":   4000,
	"gpt-3.5-turbo":      2048,
	"gpt-3.5-turbo-0301": 2048,
	"gpt-4":              4096,
	"gpt-4-0314":         4096,
	"gpt-4-32k":          8192,
	"gpt-4-32k-0314":     8192,
}

type ModelParameters struct {
	Model            string
	MaxTokens        int
	Temperature      float32
	TopP             float32
	PresencePenalty  float32
	FrequencyPenalty float32
	BestOf           int
}

// text-davinci-003
func DefaultTD003Model() *ModelParameters {
	return &ModelParameters{
		Model:            "text-davinci-003",
		MaxTokens:        DefaultMaxOutputTokens["text-davinci-003"],
		Temperature:      0,
		TopP:             1,
		BestOf:           1,
		PresencePenalty:  0,
		FrequencyPenalty: 0,
	}
}

type GPT struct {
	client *gogpt.Client
	*ModelParameters
}

func NewGPT(c *gogpt.Client) *GPT {
	return &GPT{
		client:          c,
		ModelParameters: DefaultTD003Model(),
	}
}

func (g *GPT) Request(ctx context.Context, prompt string) (string, error) {
	if g.MaxTokens+len(prompt) > ModelAllowedMaxTokens[g.Model] {
		return "", errors.Errorf("mode %s max tokens limit exceeded, max tokens: %d, output token allowed: %s, prompt length: %d", g.Model, ModelAllowedMaxTokens[g.Model], g.MaxTokens, len(prompt))
	}

	req := gogpt.CompletionRequest{
		Model:            g.Model,
		MaxTokens:        g.MaxTokens,
		TopP:             g.TopP,
		BestOf:           g.BestOf,
		Temperature:      g.Temperature,
		PresencePenalty:  g.PresencePenalty,
		FrequencyPenalty: g.FrequencyPenalty,
		Prompt:           prompt,
	}
	resp, err := g.client.CreateCompletion(ctx, req)
	if err != nil {
		log.Error("request gpt api", zap.Error(err))
		return "", err
	}
	log.Debug("openai api", zap.String("prompt", prompt), zap.String("predict", resp.Choices[0].Text))
	return resp.Choices[0].Text, nil
}

func (g *GPT) CreateEmbedding(ctx context.Context, inputs []string) ([][]float32, error) {
	resp, err := g.client.CreateEmbeddings(ctx, gogpt.EmbeddingRequest{
		Input: inputs,
		Model: gogpt.AdaEmbeddingV2,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}

	log.Info("usage", zap.Int("prompt tokens", resp.Usage.PromptTokens), zap.Int("completion tokens", resp.Usage.CompletionTokens), zap.Int("total tokens", resp.Usage.TotalTokens))
	embeddings := make([][]float32, 0, len(inputs))
	for _, d := range resp.Data {
		fmt.Println(d.Index)
		embeddings = append(embeddings, d.Embedding)
	}

	if len(embeddings) != len(inputs) {
		return nil, errors.Errorf("embedding response missing some input inputs(%d) vs embedding(%d)", len(inputs), len(embeddings))
	}

	return embeddings, nil
}