package main

import (
	"context"
	"time"

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

var edgeCtx context.Context
var edgeClient *edgedb.Client

type User struct {
	ID      edgedb.UUID `edgedb:"id"`
	Setting Setting     `edgedb:"setting"`
}

type Key struct {
	ID      edgedb.UUID `edgedb:"id"`
	Name    string      `edgedb:"name"`
	Company CompanyInfo `edgedb:"company"`
	Key     string      `edgedb:"key"`
	Date    time.Time   `edgedb:"date"`
}

type Setting struct {
	ID           edgedb.UUID        `edgedb:"id"`
	Keys         []Key              `edgedb:"keys"`
	DefaultModel edgedb.OptionalStr `edgedb:"default_model"`
}

type Conversation struct {
	ID   edgedb.UUID `edgedb:"id"`
	Name string      `edgedb:"name"`
	Date time.Time   `edgedb:"date"`
	User User        `edgedb:"user"`
}

type Area struct {
	ID       edgedb.UUID  `edgedb:"id"`
	Position int64        `edgedb:"position"`
	Conv     Conversation `edgedb:"conversation"`
}

type Message struct {
	ID       edgedb.UUID  `edgedb:"id"`
	Content  string       `edgedb:"content"`
	Role     string       `edgedb:"role"`
	Selected bool         `edgedb:"selected"`
	Date     time.Time    `edgedb:"date"`
	Area     Area         `edgedb:"area"`
	Conv     Conversation `edgedb:"conversation"`
	LLM      LLM          `edgedb:"llm"`
}

type Usage struct {
	ID          edgedb.UUID `edgedb:"id"`
	ModelID     string      `edgedb:"model_id"`
	Date        time.Time   `edgedb:"date"`
	InputCost   float32     `edgedb:"input_cost"`
	OutputCost  float32     `edgedb:"output_cost"`
	InputToken  int32       `edgedb:"input_token"`
	OutputToken int32       `edgedb:"output_token"`
}

type LLM struct {
	ID          edgedb.UUID `edgedb:"id"`
	Name        string      `edgedb:"name"`
	Context     string      `edgedb:"context"`
	Temperature float32     `edgedb:"temperature"`
	Model       ModelInfo   `edgedb:"modelInfo"`
}

type ModelInfo struct {
	ID          edgedb.UUID `edgedb:"id"`
	Name        string      `edgedb:"name"`
	MaxToken    int32       `edgedb:"maxToken"`
	InputPrice  float32     `edgedb:"inputPrice"`
	OutputPrice float32     `edgedb:"outputPrice"`
	ModelID     string      `edgedb:"modelID"`
	Company     CompanyInfo `edgedb:"company"`
}

type CompanyInfo struct {
	ID   edgedb.UUID `edgedb:"id"`
	Name string      `edgedb:"name"`
	Icon string      `edgedb:"icon"`
}

func init() {
	var ctx = context.Background()
	client, err := edgedb.CreateClient(ctx, edgedb.Options{})
	if err != nil {
		panic(err)
	}

	edgeCtx = ctx
	edgeClient = client
}

func getLastArea() edgedb.UUID {
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
		select Area
		filter .conversation.name = 'Default' AND .conversation.user = global currentUser
		order by .position desc
		limit 1
	`, &inserted)
	if err != nil {
		panic(err)
	}
	return inserted.id
}

func checkIfLogin() bool {
	var result User
	err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)
	return err == nil
}

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 {
		panic(err)
	}
}

func insertNewConversation() edgedb.UUID {
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
		INSERT Conversation {
			name := 'Default',
			user := global currentUser
		}
	`, &inserted)
	if err != nil {
		panic(err)
	}
	return inserted.id
}

func insertArea() (edgedb.UUID, int64) {
	// If the Default conversation doesn't exist, create it.
	var convExists bool
	edgeClient.QuerySingle(edgeCtx, `
	select exists (
		select Conversation
		filter .name = 'Default' AND .user = global currentUser
	  );
	`, &convExists)
	if !convExists {
		insertNewConversation()
	}
	// Insert a new area.
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
	WITH
		positionVar := count((SELECT Area FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser)) + 1
	INSERT Area {
			position := positionVar,
			conversation := (
				SELECT Conversation
				FILTER .name = 'Default' AND .user = global currentUser
				LIMIT 1
			)
		}
	`, &inserted)
	if err != nil {
		panic(err)
	}

	var positionSet struct{ position int64 }
	err = edgeClient.QuerySingle(edgeCtx, `
		SELECT Area {
			position
		}
		FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser
		ORDER BY .position desc
		LIMIT 1
	`, &positionSet)
	if err != nil {
		panic(err)
	}

	return inserted.id, positionSet.position
}

func insertUserMessage(content string) edgedb.UUID {
	// Insert a new user.
	lastAreaID := getLastArea()
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
        INSERT Message {
			role := <str>$0,
			content := <str>$1,
			area := (
				SELECT Area
				FILTER .id = <uuid>$2
			),
			conversation := (
				SELECT Conversation
				FILTER .name = 'Default' AND .user = global currentUser
				LIMIT 1
			),
			selected := true,
			llm := (
				SELECT LLM FILTER .id = <uuid>"a32c43ec-12fc-11ef-9dc9-b38e0de8bff0"
			)
        }
    `, &inserted, "user", content, lastAreaID)
	if err != nil {
		panic(err)
	}
	return inserted.id
}

func insertBotMessage(content string, selected bool, llmUUID edgedb.UUID) edgedb.UUID {
	lastAreaID := getLastArea()
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
		INSERT Message {
			role := <str>$0,
			content := <str>$2,
			selected := <bool>$3,
			conversation := (
				SELECT Area
				FILTER .id = <uuid>$4
				LIMIT 1
			).conversation,
			area := (
				SELECT Area
				FILTER .id = <uuid>$4
				LIMIT 1
			),
			llm := (
				SELECT LLM
				FILTER .id = <uuid>$1
				LIMIT 1
			)
		}
	`, &inserted, "bot", llmUUID, content, selected, lastAreaID)
	if err != nil {
		panic(err)
	}
	return inserted.id
}

func getAllMessages() []Message {
	// If no CurrentUser, return an empty array
	if !checkIfLogin() {
		return []Message{}
	}

	var messages []Message

	err := edgeClient.Query(edgeCtx, `
	SELECT Message {
		id,
		selected,
		role,
		content,
		date,
		llm : {
			name,
			modelInfo : {
				modelID,
				name,
				company : {
					icon
				}
			}
		}
	  } 
	FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser
	ORDER BY .date ASC
    `, &messages)
	if err != nil {
		panic(err)
	}

	return messages
}

func getAllSelectedMessages() []Message {
	// If no CurrentUser, return an empty array
	if !checkIfLogin() {
		return []Message{}
	}

	var messages []Message

	err := edgeClient.Query(edgeCtx, `
	SELECT Message {
		id,
		selected,
		role,
		content,
		date,
		llm : {
			modelInfo : {
				modelID,
				company : {
					icon
				}
			}
		}
	  } FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser AND .selected = true
	  ORDER BY .date ASC
    `, &messages)
	if err != nil {
		panic(err)
	}

	return messages
}

func getCurrentUserKeys() []Key {
	var result []Key
	err := edgeClient.Query(edgeCtx, "SELECT global currentUser.setting.keys", &result)
	if err != nil {
		panic(err)
	}
	return result
}

func checkIfHaveKey() bool {
	keys := getCurrentUserKeys()
	return len(keys) > 0
}