package main

import (
	"fmt"
	"log"
	"sync"

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

type ModelInfo struct {
	ID          string
	Name        string
	Icon        string
	MaxToken    int
	InputPrice  float32
	OutputPrice float32
}

type CompanyInfo struct {
	ID         string
	Name       string
	Icon       string
	ModelInfos []ModelInfo
}

var CompanyInfos []CompanyInfo
var ModelsInfos []ModelInfo

type MultipleModelsCompletionRequest struct {
	ModelIds []string
	Messages []Message
	Message  string
}

type BotContentMessage struct {
	Content string
	Hidden  bool
}

var lastSelectedModelIds []string

func addUsage(inputCost float32, outputCost float32, inputToken int32, outputToken int32, modelID string) {
	// Create a new usage
	err := edgeClient.Execute(edgeCtx, `
		INSERT Usage {
			input_cost := <float32>$0,
			output_cost := <float32>$1,
			input_token := <int32>$2,
			output_token := <int32>$3,
			model_id := <str>$4,
			user := global currentUser
		}
	`, inputCost, outputCost, inputToken, outputToken, modelID)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in addUsage")
		log.Fatal(err)
	}
}

func GenerateMultipleMessages(c *fiber.Ctx) error {
	// Create a wait group to synchronize the goroutines
	var wg sync.WaitGroup
	var InsertedIDs []edgedb.UUID

	// Add the length of lastSelectedModelIds goroutines to the wait group
	wg.Add(len(lastSelectedModelIds))

	for i := range lastSelectedModelIds {
		if model2Icon(lastSelectedModelIds[i]) == "openai" {
			go func() {
				defer wg.Done()
				response := addOpenaiMessage(lastSelectedModelIds[i], i == 0)
				InsertedIDs = append(InsertedIDs, response)
			}()
		} else if model2Icon(lastSelectedModelIds[i]) == "anthropic" {
			go func() {
				defer wg.Done()
				response := addAnthropicMessage(lastSelectedModelIds[i], i == 0)
				InsertedIDs = append(InsertedIDs, response)
			}()
		}
	}

	// Wait for both goroutines to finish
	wg.Wait()

	return c.SendString(generateChatHTML())
}

func RequestMultipleMessages(c *fiber.Ctx) error {
	message := c.FormValue("message")
	if chatString, commandExecuted := DetectCommand(message, c); commandExecuted {
		return c.SendString(chatString)
	}

	// Add an Area with the user message inside
	insertArea()
	messageID := insertUserMessage(message)

	selectedModelIds := []string{}
	for CompanyInfo := range CompanyInfos {
		for ModelInfo := range CompanyInfos[CompanyInfo].ModelInfos {
			out := c.FormValue("model-check-" + CompanyInfos[CompanyInfo].ModelInfos[ModelInfo].ID)
			if out != "" {
				selectedModelIds = append(selectedModelIds, CompanyInfos[CompanyInfo].ModelInfos[ModelInfo].ID)
			}
		}
	}
	lastSelectedModelIds = selectedModelIds

	out := ""

	messageOut, _ := userTmpl.Execute(pongo2.Context{"Content": markdownToHTML(message), "ID": messageID.String()})
	out += messageOut

	messageOut, _ = botTmpl.Execute(pongo2.Context{"IsPlaceholder": true, "SelectedModelIds": selectedModelIds, "Message": message})
	out += messageOut

	return c.SendString(out)
}