package ai

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

type OpenAIResponse struct {
	ID          string   `json:"id"`
	Object      string   `json:"object"`
	Created     int64    `json:"created"`
	Model       string   `json:"model"`
	Choices     []Choice `json:"choices"`
	Usage       Usage    `json:"usage"`
	ServiceTier string   `json:"service_tier"`
}

type Choice struct {
	Index        int64       `json:"index"`
	Message      Message     `json:"message"`
	Logprobs     interface{} `json:"logprobs"`
	FinishReason string      `json:"finish_reason"`
}

type Message struct {
	Role        string        `json:"role"`
	Content     string        `json:"content"`
	Refusal     interface{}   `json:"refusal"`
	Annotations []interface{} `json:"annotations"`
}

type Usage struct {
	PromptTokens            int64                   `json:"prompt_tokens"`
	CompletionTokens        int64                   `json:"completion_tokens"`
	TotalTokens             int64                   `json:"total_tokens"`
	PromptTokensDetails     PromptTokensDetails     `json:"prompt_tokens_details"`
	CompletionTokensDetails CompletionTokensDetails `json:"completion_tokens_details"`
}

type CompletionTokensDetails struct {
	ReasoningTokens          int64 `json:"reasoning_tokens"`
	AudioTokens              int64 `json:"audio_tokens"`
	AcceptedPredictionTokens int64 `json:"accepted_prediction_tokens"`
	RejectedPredictionTokens int64 `json:"rejected_prediction_tokens"`
}

type PromptTokensDetails struct {
	CachedTokens int64 `json:"cached_tokens"`
	AudioTokens  int64 `json:"audio_tokens"`
}

type OpenAIAPI struct {
	Token   string
	Model   string
	BaseURL string
}

func NewOpenAIAPI(baseURL, token, model string) *OpenAIAPI {
	return &OpenAIAPI{
		Token:   token,
		Model:   model,
		BaseURL: baseURL,
	}
}

type CreateCompletionReq struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
}

func (o *OpenAIAPI) CreateCompletion(request CreateCompletionReq) (*OpenAIResponse, error) {
	url := fmt.Sprintf("%s/v1/chat/completions", o.BaseURL)
	data, err := json.Marshal(request)
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(data)
	req, err := http.NewRequest("POST", url, buf)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", o.Token))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	response := new(OpenAIResponse)
	err = json.Unmarshal(body, response)
	return response, err
}