// I guess it should be some kind of package with different part for different Company
// I will maybe do it in the future, rn, I don't have time to learn that and I like it like that
// It do need more time and try and error to do all of them but they are fully independant
// And this is simple, I don't need to trick my mind to undersant that some part are share, ect
// If I want to see how I go from the message to the response, I can. That what I want

// If you are wondering how it work:
// User send message -> Generate HTML of one user message and one bot placeholder ----
// -> Send HTML and append it to the chat container -> The placeholder do a load HTMX request to GenerateMultipleMessagesHandler ----
// -> Make multiple request in parallel to all APIs -> Send one SSE event per message receive.
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

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

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

var lastSelectedLLMs []LLM

func GeneratePlaceholderHandler(c *fiber.Ctx) error {
	// Step 1 I create a User message and send it as output with a placeholder
	// that will make a request to GenerateMultipleMessagesHandler when loading
	message := c.FormValue("message", "")
	var selectedLLMIds []string
	err := json.Unmarshal([]byte(c.FormValue("selectedLLMIds")), &selectedLLMIds)
	if err != nil {
		fmt.Println("Error unmarshalling selected LLM IDs")
		panic(err)
	}

	return c.SendString(GeneratePlaceholderHTML(c, message, selectedLLMIds, true))
}

func GeneratePlaceholderHTML(c *fiber.Ctx, message string, selectedLLMIds []string, with_user_message bool) string {
	var selectedLLMs []LLM
	var selectedLLM LLM

	for _, id := range selectedLLMIds {
		idUUID, _ := edgedb.ParseUUID(id)
		err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(context.Background(), `
			SELECT LLM {
				id,
				name,
				context,
				temperature,
        max_tokens,
				custom_endpoint : {
					id,
					endpoint,
					key
				},
				modelInfo : {
					modelID,
					company : {
						icon,
						name
					}
				}
			}
			FILTER
				.id = <uuid>$0;
		`, &selectedLLM, idUUID)
		if err != nil {
			fmt.Println("Error getting LLM")
			panic(err)
		}
		selectedLLMs = append(selectedLLMs, selectedLLM)
	}
	lastSelectedLLMs = selectedLLMs

	_, position := insertArea(c)

	var user User
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(context.Background(), `
		SELECT global currentUser { id } LIMIT 1
	`, &user)
	if err != nil {
		fmt.Println("Error getting user")
		panic(err)
	}

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

	messageOut, _ := botTmpl.Execute(pongo2.Context{
		"IsPlaceholder":      true,
		"SelectedLLMs":       selectedLLMs,
		"ConversationAreaId": position + 1,
		"userID":             user.ID.String(),
	})
	out += messageOut

	// defer sendEvent("hide-placeholder", "")

	return out
}

func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
	// Step 2 generate multiple messages
	// And send them one by one using events
	insertArea(c)
	selectedLLMs := lastSelectedLLMs

	var user User
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(context.Background(), `
		SELECT global currentUser { id } LIMIT 1
	`, &user)
	if err != nil {
		fmt.Println("Error getting user")
		panic(err)
	}

	// Create a wait group to synchronize the goroutines
	var wg sync.WaitGroup

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

	// Create a channel to receive the index of the first completed goroutine
	firstDone := make(chan int, 1)

	for i := range selectedLLMs {
		idx := i
		go func() {
			defer wg.Done()

			// Create a context with a 1-minute timeout
			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
			defer cancel() // Ensure the context is cancelled to free resources

			// Determine which message function to call based on the model
			var addMessageFunc func(c *fiber.Ctx, selectedLLM LLM, selected bool) edgedb.UUID
			switch selectedLLMs[idx].Model.Company.Name {
			case "openai":
				addMessageFunc = addOpenaiMessage
			case "anthropic":
				addMessageFunc = addAnthropicMessage
			case "mistral":
				addMessageFunc = addMistralMessage
			case "groq":
				addMessageFunc = addGroqMessage
			case "gooseai":
				addMessageFunc = addGooseaiMessage
			case "huggingface":
				addMessageFunc = addHuggingfaceMessage
			case "google":
				addMessageFunc = addGoogleMessage
			case "perplexity":
				addMessageFunc = addPerplexityMessage
			case "fireworks":
				addMessageFunc = addFireworkMessage
      case "nim":
        addMessageFunc = addNimMessage
			}

			var messageID edgedb.UUID
			if addMessageFunc != nil {
				messageID = addMessageFunc(c, selectedLLMs[idx], idx == 0)
			}

			var message Message
			err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
				SELECT Message {
					id,
					content,
					area : {
						id,
						position
					},
					llm : {
						modelInfo : {
							modelID,
							name,
							company : {
								icon,
							}
						}
					}
				}
				FILTER .id = <uuid>$0;
			`, &message, messageID)
			if err != nil {
				fmt.Println("Error getting message for the placeholder. The function addProviderMessage seem to not return any message ID.")
				panic(err)
			}

			// Check if the context's deadline is exceeded
			select {
			case <-ctx.Done():
				// The context's deadline was exceeded
				fmt.Printf("Goroutine %d timed out\n", idx)
			default:
				// Send the index of the completed goroutine to the firstDone channel
				select {
				case firstDone <- idx:
					// Generate the HTML content

					outIcon := `<img src="` + selectedLLMs[idx].Model.Company.Icon + `" alt="User Image" id="selectedIcon-` + fmt.Sprintf("%d", message.Area.Position) + `">`

					go func() {
						// I do a ping because of sse size limit. Do see if it's possible to do without it TODO
						sendEvent(
							user.ID.String(),
							"swapContent-"+fmt.Sprintf("%d", message.Area.Position),
							`<hx hx-get="/messageContent?id=`+message.ID.String()+`" hx-trigger="load" hx-swap="outerHTML"></hx>`,
						)
						sendEvent(
							user.ID.String(),
							"swapSelectionBtn-"+selectedLLMs[idx].ID.String(),
							`<hx hx-get="/selectionBtn?id=`+message.ID.String()+`" hx-trigger="load" hx-swap="outerHTML"></hx>`,
						)
						sendEvent(
							user.ID.String(),
							"swapIcon-"+fmt.Sprintf("%d", message.Area.Position),
							outIcon,
						)
					}()
				default:
					// Send Content event
					go func() {
						sendEvent(
							user.ID.String(),
							"swapSelectionBtn-"+selectedLLMs[idx].ID.String(),
							`<hx hx-get="/selectionBtn?id=`+message.ID.String()+`" hx-trigger="load" hx-swap="outerHTML"></hx>`,
						)
					}()
				}
			}
		}()
	}

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

	return c.SendString("")
}

func addUsage(c *fiber.Ctx, inputCost float32, outputCost float32, inputToken int32, outputToken int32, modelID string) {
	// Create a new usage
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).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 adding usage")
		panic(err)
	}
}