package doubao

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
)

type RequestBody struct {
	Model               string    `json:"model"`
	Messages            []Message `json:"messages"`
	MaxCompletionTokens int32     `json:"max_completion_tokens"`
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type SearchOptions struct {
	SearchStrategy string `json:"search_strategy"`
}

type ResponseBody struct {
	Choices []struct {
		FinishReason string  `json:"finish_reason"`
		Message      Message `json:"message"`
	} `json:"choices"`
	Usage struct {
		InputTokens  int32 `json:"prompt_tokens"`
		OutputTokens int32 `json:"completion_tokens"`
		TotalTokens  int32 `json:"total_tokens"`
	} `json:"usage"`
}

func request(model, key string, messages []Message) (*ResponseBody, error) {
	requestBody := RequestBody{
		Model:               model,
		Messages:            messages,
		MaxCompletionTokens: 65535,
	}
	jsonData, _ := json.Marshal(requestBody)

	req, err := http.NewRequest("POST", "https://ark.cn-beijing.volces.com/api/v3/chat/completions", bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Content-Type", "application/json")

	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var res ResponseBody
	if err = json.Unmarshal(body, &res); err != nil {
		return nil, err
	}
	return &res, nil
}