package ai

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

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

type Input struct {
	Messages []Message `json:"messages"`
}

type Parameters struct {
	EnableSearch  bool          `json:"enable_search,omitempty"`
	SearchOptions SearchOptions `json:"search_options,omitempty"`
	ResultFormat  string        `json:"result_format,omitempty"`
	Temperature   float32       `json:"temperature,omitempty"`
}

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

type RequestBody struct {
	Model      string     `json:"model"`
	Input      Input      `json:"input"`
	Parameters Parameters `json:"parameters"`
}

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

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

	req, err := http.NewRequest("POST", "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation", bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+"sk-4daf94ad2fa94288b198d2a1d8924cef")
	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
}