package main

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

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

func ChatPageHandler(c *fiber.Ctx) error {
	authCookie := c.Cookies("jade-edgedb-auth-token", "")

	if authCookie != "" && !checkIfLogin() {
		edgeClient = edgeClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": authCookie})
	}

	return c.Render("chat", fiber.Map{"IsLogin": checkIfLogin(), "HaveKey": checkIfHaveKey()}, "layouts/main")
}

func DeleteMessageHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")

	messageUUID, err := edgedb.ParseUUID(messageId)
	if err != nil {
		panic(err)
	}

	// Delete all messages
	err = edgeClient.Execute(edgeCtx, `
	WITH
		messageArea := (SELECT Message FILTER .id = <uuid>$0).area
	DELETE Area
		FILTER .position >= messageArea.position AND .conversation = messageArea.conversation;
	`, messageUUID)
	if err != nil {
		panic(err)
	}

	return c.SendString(generateChatHTML())
}

func LoadChatHandler(c *fiber.Ctx) error {
	deleteLLMtoDelete()

	if checkIfLogin() {
		if getCurrentUserKeys() == nil {
			return c.SendString(generateEnterKeyChatHTML())
		}
		return c.SendString(generateChatHTML())
	} else {
		return c.SendString(generateWelcomeChatHTML())
	}
}

type TemplateMessage struct {
	Icon    string
	Content string
	Hidden  bool
	Id      string
	Name    string
	Model   string
	ModelID string
}

func generateChatHTML() string {
	// Get the messages from the database
	// Maybe redo that to be area by area because look like shit rn. It come from early stage of dev. It work tho soooo...
	Messages := getAllMessages()

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	var templateMessages []TemplateMessage

	for i, message := range Messages {
		if message.Role == "user" {
			htmlContent := markdownToHTML(message.Content)
			userOut, err := userTmpl.Execute(pongo2.Context{"Content": htmlContent, "ID": message.ID.String()})
			if err != nil {
				panic(err)
			}
			htmlString += userOut
			// Reset NextMessages when a user message is encountered
			templateMessages = []TemplateMessage{}
		} else {
			// For bot messages, add them to NextMessages with only the needed fields
			templateMessage := TemplateMessage{
				Icon:    message.LLM.Model.Company.Icon, // Assuming Icon is a field you want to include from Message
				Content: markdownToHTML(message.Content),
				Hidden:  !message.Selected, // Assuming Hidden is a field you want to include from Message
				Id:      message.ID.String(),
				Name:    message.LLM.Name,
				Model:   message.LLM.Model.Name,
				ModelID: message.LLM.Model.ModelID,
			}
			templateMessages = append(templateMessages, templateMessage)

			// Check if the next message is not a bot or if it's the last message
			if i+1 == len(Messages) || Messages[i+1].Role != "bot" {
				sort.Slice(templateMessages, func(i, j int) bool {
					if !templateMessages[i].Hidden && templateMessages[j].Hidden {
						return true
					}
					if templateMessages[i].Hidden && !templateMessages[j].Hidden {
						return false
					}
					return true
				})

				botOut, err := botTmpl.Execute(pongo2.Context{"Messages": templateMessages, "ConversationAreaId": i})
				if err != nil {
					panic(err)
				}
				htmlString += botOut
				htmlString += "<div style='height: 10px;'></div>"
			}
		}
	}

	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

func GetUserMessageHandler(c *fiber.Ctx) error {
	id := c.FormValue("id")

	messageUUID, _ := edgedb.ParseUUID(id)

	var selectedMessage Message
	err := edgeClient.QuerySingle(context.Background(), `
		SELECT Message {
			content
		}
		FILTER
			.id = <uuid>$0;
	`, &selectedMessage, messageUUID)
	if err != nil {
		panic(err)
	}

	out, err := userTmpl.Execute(pongo2.Context{"Content": markdownToHTML(selectedMessage.Content), "ID": id})
	if err != nil {
		panic(err)
	}

	return c.SendString(out)
}

func GetMessageContentHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")
	onlyContent := c.FormValue("onlyContent") // To init the text area of the edit message form

	messageUUID, _ := edgedb.ParseUUID(messageId)

	var selectedMessage Message
	err := edgeClient.QuerySingle(context.Background(), `
		SELECT Message {
			content,
			llm : {
				name,
				modelInfo : {
					modelID,
					name,
				}
			}
		}
		FILTER
			.id = <uuid>$0;
	`, &selectedMessage, messageUUID)
	if err != nil {
		panic(err)
	}

	if onlyContent == "true" {
		return c.SendString(markdownToHTML(selectedMessage.Content))
	}

	out := "<div class='message-header'>"
	out += "<p>"
	out += "<strong>" + selectedMessage.LLM.Name + "</strong> <small>" + selectedMessage.LLM.Model.ModelID + "</small>"
	out += "	</p>"
	out += "</div>"
	out += "<div class='message-body'>"
	out += "	<ct class='content'>"
	out += markdownToHTML(selectedMessage.Content)
	out += "	</ct>"
	out += "</div>"

	// Update the selected value of messages in the database
	err = edgeClient.Execute(edgeCtx, `
	WITH m := (SELECT Message FILTER .id = <uuid>$0)
	UPDATE Message
	FILTER .area = m.area
	SET {selected := false};
	`, messageUUID)
	if err != nil {
		panic(err)
	}

	_ = edgeClient.Execute(edgeCtx, `
	UPDATE Message
	FILTER .id = <uuid>$0
	SET {selected := true};
	`, messageUUID)

	return c.SendString(out)
}

func generateWelcomeChatHTML() string {
	welcomeMessage := `To start using JADE, please login.`

	loginButton := `
	<a class="button is-primary is-small" href="/signin">
		Log in
	</a>`

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	NextMessages := []TemplateMessage{}
	nextMsg := TemplateMessage{
		Icon:    "icons/bouvai2.png", // Assuming Icon is a field you want to include from Message
		Content: "<br>" + markdownToHTML(welcomeMessage) + loginButton,
		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, "NotClickable": true})
	if err != nil {
		panic(err)
	}
	htmlString += botOut
	htmlString += "<div style='height: 10px;'></div>"
	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

func generatePricingTableChatHTML() string {
	stripeTable := `
	<stripe-pricing-table pricing-table-id="prctbl_1PJAxDP2nW0okNQyY0Q3mbg4"
	publishable-key="pk_live_51OxXuWP2nW0okNQyme1qdwbL535jbMmM1uIUi6U5zcvEUUwKraktmpCzudXNdPSTxlHpw2FbCtxpwbyFFcasQ7aj000tJJGpWW">
	</stripe-pricing-table>`

	closeBtn := `
	<div class="is-flex is-justify-content-flex-end">
	<a class="button is-small is-danger is-outlined" hx-get="/loadChat" hx-target="#chat-container" hx-swap="outerHTML"
	hx-trigger="click">
	<span class="icon">
	<i class="fa-solid fa-xmark"></i>
</span>
	</a>
	</div>`

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	NextMessages := []TemplateMessage{}
	nextMsg := TemplateMessage{
		Icon:    "icons/bouvai2.png", // Assuming Icon is a field you want to include from Message
		Content: "<br>" + stripeTable + closeBtn,
		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, "NotClickable": true, "notFlex": true})
	if err != nil {
		panic(err)
	}
	htmlString += botOut
	htmlString += "<div style='height: 10px;'></div>"
	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

func generateEnterKeyChatHTML() string {
	welcomeMessage := `To start using JADE, please enter at least one key in the settings.`

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	NextMessages := []TemplateMessage{}
	nextMsg := TemplateMessage{
		Icon:    "icons/bouvai2.png", // Assuming Icon is a field you want to include from Message
		Content: "<br>" + markdownToHTML(welcomeMessage),
		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, "NotClickable": true})
	if err != nil {
		panic(err)
	}
	htmlString += botOut
	htmlString += "<div style='height: 10px;'></div>"
	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

// Buton actions
func GetEditMessageFormHandler(c *fiber.Ctx) error {
	id := c.FormValue("id")
	idUUID, _ := edgedb.ParseUUID(id)

	var message Message
	err := edgeClient.QuerySingle(context.Background(), `
		SELECT Message { content }
		FILTER .id = <uuid>$0;
	`, &message, idUUID)
	if err != nil {
		panic(err)
	}

	// Calculate the number of rows based on the length of the content
	rows := len(strings.Split(message.Content, "\n"))
	if rows < 10 {
		rows = 10
	}

	tmpl := pongo2.Must(pongo2.FromFile("views/partials/message-edit-form.html"))
	out, err := tmpl.Execute(pongo2.Context{"Content": message.Content, "ID": id, "Rows": rows})
	if err != nil {
		panic(err)
	}
	return c.SendString(out)
}

func RedoMessageHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")
	messageUUID, _ := edgedb.ParseUUID(messageId)

	var selectedLLMIds []string
	err := json.Unmarshal([]byte(c.FormValue("selectedLLMIds")), &selectedLLMIds)
	if err != nil {
		// Handle the error
		panic(err)
	}

	var message Message
	err = edgeClient.QuerySingle(context.Background(), `
		SELECT Message { content }
		FILTER .id = <uuid>$0;
	`, &message, messageUUID)
	if err != nil {
		panic(err)
	}

	// Delete messages
	err = edgeClient.Execute(edgeCtx, `
	WITH
		messageArea := (SELECT Message FILTER .id = <uuid>$0).area
	DELETE Area
		FILTER .position > messageArea.position AND .conversation = messageArea.conversation AND .conversation.user = global currentUser;
	`, messageUUID)
	if err != nil {
		panic(err)
	}

	return c.SendString(GeneratePlaceholderHTML(message.Content, selectedLLMIds, false))
}

func ClearChatHandler(c *fiber.Ctx) error {
	// Delete the default conversation
	err := edgeClient.Execute(edgeCtx, `
	DELETE Conversation
	FILTER .user = global currentUser AND .name = "Default";
	`)
	if err != nil {
		panic(err)
	}

	return c.SendString(generateChatHTML())
}

// Popover stuff
func LoadUsageKPIHandler(c *fiber.Ctx) error {
	if !checkIfLogin() || !checkIfHaveKey() {
		return c.SendString("")
	}

	InputDateID := c.FormValue("month", time.Now().Format("01-2006"))
	offset := c.FormValue("offset")
	IsActive := false

	var InputDate time.Time
	InputDate, err := time.Parse("01-2006", InputDateID)
	if err != nil {
		panic(err)
	}

	if offset == "-1" {
		InputDate = InputDate.AddDate(0, -1, 0)
		IsActive = true
	} else if offset == "1" {
		InputDate = InputDate.AddDate(0, 1, 0)
		IsActive = true
	}

	type UsageKPI struct {
		Key struct {
			ModelID string `edgedb:"model_id"`
		} `edgedb:"key"`
		TotalCost  float32 `edgedb:"total_cost"`
		TotalCount int64   `edgedb:"total_count"`
	}

	var usages []UsageKPI
	err = edgeClient.Query(edgeCtx, `
	WITH 
		U := (
			SELECT Usage 
			FILTER .user = global currentUser and .date >= <datetime>$0 AND .date < <datetime>$1
		),
		grouped := (
			GROUP U {
				model_id,
				input_cost,
				output_cost,
			} BY .model_id
		)
	SELECT grouped {
		key := .key { model_id },
		total_count := count(.elements),
		total_cost := sum(.elements.input_cost) + sum(.elements.output_cost),
	} FILTER .total_count > 0 ORDER BY .total_cost DESC
	`, &usages, InputDate, InputDate.AddDate(0, 1, 0))
	if err != nil {
		fmt.Println(err)
		panic(err)
	}

	BeautifullDate := InputDate.Format("Jan 2006")

	var (
		TotalCount int64
		TotalCost  float32
	)
	for _, usage := range usages {
		TotalCost += usage.TotalCost
		TotalCount += usage.TotalCount
	}

	fmt.Println(TotalCost, TotalCount)

	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-usage.html")).Execute(pongo2.Context{
		"usages":     usages,
		"TotalCost":  TotalCost,
		"TotalCount": TotalCount,
		"Date":       BeautifullDate,
		"DateID":     InputDate.Format("01-2006"),
		"IsActive":   IsActive,
	})
	if err != nil {
		panic(err)
	}

	return c.SendString(out)
}

func LoadKeysHandler(c *fiber.Ctx) error {
	if !checkIfLogin() {
		return c.SendString("")
	}
	openaiExists, anthropicExists, mistralExists, groqExists, gooseaiExists, googleExists := getExistingKeys()

	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-keys.html")).Execute(pongo2.Context{
		"IsLogin":         checkIfLogin(),
		"OpenaiExists":    openaiExists,
		"AnthropicExists": anthropicExists,
		"MistralExists":   mistralExists,
		"GroqExists":      groqExists,
		"GooseaiExists":   gooseaiExists,
		"GoogleExists":    googleExists,
		"AnyExists":       openaiExists || anthropicExists || mistralExists || groqExists || gooseaiExists || googleExists,
	})
	if err != nil {
		panic(err)
	}
	return c.SendString(out)
}

func GenerateModelPopoverHTML(refresh bool) string {
	openaiExists, anthropicExists, mistralExists, groqExists, gooseaiExists, googleExists := getExistingKeys()

	var llms []LLM
	err := edgeClient.Query(context.Background(), `
		SELECT LLM {
			id,
			name,
			context,
			temperature,
			modelInfo : {
				modelID,
				name,
				company : {
					name,
					icon
				}
			}
		}
		FILTER .user = global currentUser AND .name != 'none' AND .to_delete = false 
	`, &llms)
	if err != nil {
		panic(err)
	}

	modelInfos := GetAvailableModels()

	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-models.html")).Execute(pongo2.Context{
		"IsLogin":         checkIfLogin(),
		"OpenaiExists":    openaiExists,
		"AnthropicExists": anthropicExists,
		"MistralExists":   mistralExists,
		"GroqExists":      groqExists,
		"GooseaiExists":   gooseaiExists,
		"GoogleExists":    googleExists,
		"AnyExists":       openaiExists || anthropicExists || mistralExists || groqExists || gooseaiExists || googleExists,
		"LLMs":            llms,
		"ModelInfos":      modelInfos,
		"DeleteUpdate":    refresh,
	})
	if err != nil {
		panic(err)
	}
	return out
}

func LoadModelSelectionHandler(c *fiber.Ctx) error {
	if !checkIfLogin() || !checkIfHaveKey() {
		return c.SendString("")
	}
	return c.SendString(GenerateModelPopoverHTML(false))
}

func LoadSettingsHandler(c *fiber.Ctx) error {
	if !checkIfLogin() {
		return c.SendString("")
	}

	openaiExists, anthropicExists, mistralExists, groqExists, gooseaiExists, googleExists := getExistingKeys()

	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-settings.html")).Execute(pongo2.Context{
		"IsLogin":         checkIfLogin(),
		"OpenaiExists":    openaiExists,
		"AnthropicExists": anthropicExists,
		"MistralExists":   mistralExists,
		"GroqExists":      groqExists,
		"GooseaiExists":   gooseaiExists,
		"GoogleExists":    googleExists,
		"AnyExists":       openaiExists || anthropicExists || mistralExists || groqExists || gooseaiExists || googleExists,
	})
	if err != nil {
		panic(err)
	}
	return c.SendString(out)
}