package main

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

	"github.com/edgedb/edgedb-go"
	"github.com/flosch/pongo2"
	"github.com/gofiber/fiber/v2"
)

type OpenaiChatCompletionRequest struct {
	Model       string          `json:"model"`
	Messages    []OpenaiMessage `json:"messages"`
	Temperature float64         `json:"temperature"`
}

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

type OpenaiChatCompletionResponse struct {
	ID      string         `json:"id"`
	Object  string         `json:"object"`
	Created int64          `json:"created"`
	Model   string         `json:"model"`
	Usage   OpenaiUsage    `json:"usage"`
	Choices []OpenaiChoice `json:"choices"`
}

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

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

func init() {
	var ModelInfosList = []ModelInfo{}

	modelInfo := ModelInfo{
		ID:          "gpt-3.5-turbo",
		Name:        "GPT-3.5",
		Icon:        "openai",
		MaxToken:    4096,
		InputPrice:  0.50 / 1000000,
		OutputPrice: 1.50 / 1000000,
	}
	ModelInfosList = append(ModelInfosList, modelInfo)
	ModelsInfos = append(ModelsInfos, modelInfo)

	modelInfo = ModelInfo{
		ID:          "gpt-4-turbo",
		Name:        "GPT-4",
		Icon:        "openai",
		MaxToken:    8192,
		InputPrice:  10.00 / 1000000,
		OutputPrice: 30.00 / 1000000,
	}
	ModelInfosList = append(ModelInfosList, modelInfo)
	ModelsInfos = append(ModelsInfos, modelInfo)

	companyInfo := CompanyInfo{
		ID:         "openai",
		Name:       "OpenAI",
		Icon:       "icons/openai.png",
		ModelInfos: ModelInfosList,
	}
	CompanyInfos = append(CompanyInfos, companyInfo)
}

func addOpenaiMessage(modelID string, selected bool) edgedb.UUID {
	Messages := getAllMessages()

	chatCompletion, err := RequestOpenai(modelID, Messages, 0.7)
	if err != nil {
		fmt.Println("Error:", err)
	} else if len(chatCompletion.Choices) == 0 {
		fmt.Println(chatCompletion)
		fmt.Println("No response from OpenAI")
		id := insertBotMessage("No response from OpenAI", selected, modelID)
		return id
	} else {
		Content := chatCompletion.Choices[0].Message.Content
		id := insertBotMessage(Content, selected, modelID)
		return id
	}
	return edgedb.UUID{}
}

func EdgeMessages2OpenaiMessages(messages []Message) []OpenaiMessage {
	openaiMessages := make([]OpenaiMessage, len(messages))
	for i, msg := range messages {
		var role string
		switch msg.Role {
		case "user":
			role = "user"
		case "bot":
			role = "assistant"
		default:
			role = "system"
		}
		openaiMessages[i] = OpenaiMessage{
			Role:    role,
			Content: msg.Content,
		}
	}
	return openaiMessages
}

func TestOpenaiKey(apiKey string) bool {
	url := "https://api.openai.com/v1/chat/completions"

	// Convert messages to OpenAI format
	openaiMessages := []OpenaiMessage{
		{
			Role:    "user",
			Content: "Hello",
		},
	}

	requestBody := OpenaiChatCompletionRequest{
		Model:       "gpt-3.5-turbo",
		Messages:    openaiMessages,
		Temperature: 0,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return false
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return false
	}

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

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

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false
	}

	var chatCompletionResponse OpenaiChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return false
	}
	if chatCompletionResponse.Choices == nil {
		return false
	}
	return false
}

func RequestOpenai(model string, messages []Message, temperature float64) (OpenaiChatCompletionResponse, error) {
	var apiKey string
	err := edgeClient.QuerySingle(edgeCtx, `
		with
			filtered_keys := (
				select Key {
					key
				} filter .company = <str>$0
			)
		select filtered_keys.key limit 1
	`, &apiKey, "openai")
	if err != nil {
		return OpenaiChatCompletionResponse{}, fmt.Errorf("error getting OpenAI API key: %w", err)
	}

	fmt.Println("OpenAI API key: ", apiKey)

	url := "https://api.openai.com/v1/chat/completions"

	// Convert messages to OpenAI format
	openaiMessages := EdgeMessages2OpenaiMessages(messages)

	requestBody := OpenaiChatCompletionRequest{
		Model:       model,
		Messages:    openaiMessages,
		Temperature: temperature,
	}

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

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

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

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

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

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

	var usedModelInfo ModelInfo
	for mi := range ModelsInfos {
		if ModelsInfos[mi].ID == model {
			usedModelInfo = ModelsInfos[mi]
		}
	}
	var inputCost float32 = float32(chatCompletionResponse.Usage.PromptTokens) * usedModelInfo.InputPrice
	var outputCost float32 = float32(chatCompletionResponse.Usage.CompletionTokens) * usedModelInfo.OutputPrice
	addUsage(inputCost, outputCost, chatCompletionResponse.Usage.PromptTokens, chatCompletionResponse.Usage.CompletionTokens, model)

	return chatCompletionResponse, nil
}

func addOpenaiKey(c *fiber.Ctx) error {
	key := c.FormValue("key")

	// Check if the key already exists
	var keyUUID edgedb.UUID
	err := edgeClient.QuerySingle(edgeCtx, `
	with
		filtered_keys := (
			select Key {
				id
			} filter .key = <str>$0 and .company = "openai"
		)
	select filtered_keys.key limit 1
	`, &keyUUID, key)
	if err == nil {
		fmt.Println("Error in edgedb.Query: in addOpenaiKey")
		return c.SendString("")
	}

	if !TestOpenaiKey(key) {
		fmt.Println("Invalid OpenAI API Key")
		NextMessages := []NextMessage{}
		nextMsg := NextMessage{
			Icon:    "bouvai2", // Assuming Icon is a field you want to include from Message
			Content: "<br>" + markdownToHTML("Invalid OpenAI API Key"),
			Hidden:  false, // Assuming Hidden is a field you want to include from Message
			Id:      "0",
			Name:    "JADE",
		}
		NextMessages = append(NextMessages, nextMsg)

		botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": 0})
		if err != nil {
			panic(err)
		}
		return c.SendString(botOut)
	}

	// Check if the company key already exists
	err = edgeClient.QuerySingle(edgeCtx, `
	with
		filtered_keys := (
			select Key {
				id
			} filter .company = "openai"
		)
	select filtered_keys.key limit 1
	`, &keyUUID, key)
	if err != nil {
		fmt.Println("Company key already exists")
		err = edgeClient.Execute(edgeCtx, `
		UPDATE Key filter .company = <str>$0 AND .key = <str>$1
		SET {
			key := <str>$1,
		}
		`, "openai", key)
		if err != nil {
			fmt.Println("Error in edgedb.QuerySingle: in addOpenaiKey")
			fmt.Println(err)
		}

		return c.SendString("")
	}

	fmt.Println("OpenAI API key: ", key)

	err = edgeClient.Execute(edgeCtx, `
	UPDATE global currentUser.setting
	SET {
		keys += (
			INSERT Key {
				company := <str>$0,
				key := <str>$1,
				name := <str>$2,
			}
		)
	}`, "openai", key, "OpenAI API Key")
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in addOpenaiKey")
		fmt.Println(err)
	}

	NextMessages := []NextMessage{}
	nextMsg := NextMessage{
		Icon:    "bouvai2", // Assuming Icon is a field you want to include from Message
		Content: "<br> Key added successfully!",
		Hidden:  false, // Assuming Hidden is a field you want to include from Message
		Id:      "0",
		Name:    "JADE",
	}
	NextMessages = append(NextMessages, nextMsg)

	botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": 0})
	if err != nil {
		panic(err)
	}

	return c.SendString(botOut)
}