package main

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

	"github.com/edgedb/edgedb-go"
)

type HuggingfaceChatCompletionRequest struct {
	Model       string           `json:"model"`
	Messages    []RequestMessage `json:"messages"`
	Temperature float64          `json:"temperature"`
	Stream      bool             `json:"stream"`
}

type HuggingfaceChatCompletionResponse struct {
	ID      string              `json:"id"`
	Object  string              `json:"object"`
	Created int64               `json:"created"`
	Model   string              `json:"model"`
	Choices []HuggingfaceChoice `json:"choices"`
}

type HuggingfaceUsage struct {
	PromptTokens     int32 `json:"prompt_tokens"`
	CompletionTokens int32 `json:"completion_tokens"`
	TotalTokens      int32 `json:"total_tokens"`
}

type HuggingfaceChoice struct {
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
	Index        int     `json:"index"`
}

func addHuggingfaceMessage(llm LLM, selected bool) edgedb.UUID {
	Messages := getAllSelectedMessages()

	chatCompletion, err := RequestHuggingface(llm, Messages, float64(llm.Temperature))
	if err != nil {
		fmt.Println("Error fetching user profile")
		panic(err)
	} else if len(chatCompletion.Choices) == 0 {
		fmt.Println("No response from Endpoint")
		id := insertBotMessage("No response from Endpoint", selected, llm.ID)
		return id
	} else {
		Content := chatCompletion.Choices[0].Message.Content
		id := insertBotMessage(Content, selected, llm.ID)
		return id
	}
}

func RequestHuggingface(llm LLM, messages []Message, temperature float64) (HuggingfaceChatCompletionResponse, error) {
	url := llm.Endpoint.Endpoint

	requestBody := HuggingfaceChatCompletionRequest{
		Model:       "tgi",
		Messages:    Message2RequestMessage(messages, llm.Context),
		Temperature: temperature,
		Stream:      false,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return HuggingfaceChatCompletionResponse{}, fmt.Errorf("error marshaling JSON: %w", err)
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return HuggingfaceChatCompletionResponse{}, fmt.Errorf("error creating request: %w", err)
	}

	req.Header.Set("Authorization", "Bearer "+llm.Endpoint.Key)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return HuggingfaceChatCompletionResponse{}, fmt.Errorf("error sending request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return HuggingfaceChatCompletionResponse{}, fmt.Errorf("error reading response body: %w", err)
	}

	var chatCompletionResponse HuggingfaceChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return HuggingfaceChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err)
	}

	addUsage(0, 0, 0, 0, llm.Model.ModelID)

	return chatCompletionResponse, nil
}
