package main

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

	"github.com/gofiber/fiber/v2"
)

type GoogleRequestMessage struct {
	Role  string       `json:"role"`
	Parts []GooglePart `json:"parts"`
}

type GooglePart struct {
	Text string `json:"text"`
}

type GoogleChatCompletionRequest struct {
	Messages []GoogleRequestMessage `json:"contents"`
}

type GoogleChatCompletionResponse struct {
	Candidates    []GoogleCandidate   `json:"candidates"`
	UsageMetadata GoogleUsageMetadata `json:"usageMetadata"`
}

type GoogleCandidate struct {
	Content struct {
		Parts []GooglePart `json:"parts"`
		Role  string       `json:"role"`
	} `json:"content"`
	FinishReason  string               `json:"finishReason"`
	Index         int                  `json:"index"`
	SafetyRatings []GoogleSafetyRating `json:"safetyRatings"`
}

type GoogleSafetyRating struct {
	Category    string `json:"category"`
	Probability string `json:"probability"`
}

type GoogleUsageMetadata struct {
	PromptTokenCount     int32 `json:"promptTokenCount"`
	CandidatesTokenCount int32 `json:"candidatesTokenCount"`
	TotalTokenCount      int32 `json:"totalTokenCount"`
}

var GoogleErrorCodes map[string]string

// TODO: Update
func init() {
	GoogleErrorCodes = make(map[string]string)
	GoogleErrorCodes["401"] = "Invalid Authentication - Ensure that the API key is still valid."
	GoogleErrorCodes["403"] = "Accessing the API from an unsupported country, region, or territory."
	GoogleErrorCodes["429"] = "Rate limit reached for requests - You are sending requests too quickly."
	GoogleErrorCodes["429"] = "You have run out of credits or hit your maximum monthly spend - Buy more credits or learn how to increase your limits."
	GoogleErrorCodes["500"] = "Issue on Provider servers - Retry your request after a brief wait and contact the provider if the issue persists."
	GoogleErrorCodes["503"] = "Servers are experiencing high traffic - Please retry your requests after a brief wait."
}

func TestGoogleKey(apiKey string) bool {
	url := "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key=" + apiKey

	googlePart := GooglePart{
		Text: "Hello",
	}
	// Convert messages to OpenAI format
	googleMessages := []GoogleRequestMessage{
		{
			Role:  "user",
			Parts: []GooglePart{googlePart},
		},
	}

	requestBody := GoogleChatCompletionRequest{
		Messages: googleMessages,
	}

	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")

	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 GoogleChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return false
	}
	if chatCompletionResponse.UsageMetadata.CandidatesTokenCount == 0 {
		return false
	}

	return true
}

func RequestGoogle(c *fiber.Ctx, llm LLM, messages []Message) string {
	model := llm.Model.ModelID

	// TODO: Use those parameters
	// temperature := float64(llm.Temperature)
	// context := llm.Context
	//maxTokens := int(llm.MaxToken)

	var apiKey string
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		with
			filtered_keys := (
				select Key {
					key
				} filter .company.name = <str>$0 AND .<keys[is Setting].<setting[is User] = global currentUser
			)
		select filtered_keys.key limit 1
	`, &apiKey, "google")
	if err != nil {
		return "JADE internal error: 03-00-0000. Please contact the support."
	}

	url := "https://generativelanguage.googleapis.com/v1beta/models/" + model + ":generateContent?key=" + apiKey

	// Do an array of GoogleRequestMessage with the messages in a for loop
	googleMessages := []GoogleRequestMessage{}
	for _, message := range messages {
		if message.Role == "user" {
			googleMessages = append(googleMessages, GoogleRequestMessage{
				Role:  "user",
				Parts: []GooglePart{{Text: message.Content}}, // Changed something here, to test
			})
		} else {
			googleMessages = append(googleMessages, GoogleRequestMessage{
				Role:  "model",
				Parts: []GooglePart{{Text: message.Content}},
			})
		}
	}

	requestBody := GoogleChatCompletionRequest{
		Messages: googleMessages,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return "JADE internal error: 03-01-0001. Please contact the support."
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return "JADE internal error: 03-02-0002. Please contact the support."
	}

	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "JADE internal error: 03-02-0003. Please contact the support."
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "JADE internal error: 03-01-0004. Please contact the support."
	}

	for key, value := range GoogleErrorCodes {
		if strings.Contains(resp.Status, key) {
			return value
		}
	}

	var chatCompletionResponse GoogleChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return "JADE internal error: 03-01-0005. Please contact the support."
	}

	var usedModelInfo ModelInfo
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		SELECT ModelInfo {
			inputPrice,
			outputPrice
		}
		FILTER .modelID = <str>$0
		LIMIT 1
	`, &usedModelInfo, model)
	if err != nil {
		return "JADE internal error: 03-00-0006. Please contact the support."
	}

	var inputCost float32 = float32(chatCompletionResponse.UsageMetadata.PromptTokenCount) * usedModelInfo.InputPrice
	var outputCost float32 = float32(chatCompletionResponse.UsageMetadata.CandidatesTokenCount) * usedModelInfo.OutputPrice
	addUsage(c, inputCost, outputCost, chatCompletionResponse.UsageMetadata.PromptTokenCount, chatCompletionResponse.UsageMetadata.CandidatesTokenCount, model)

	if len(chatCompletionResponse.Candidates) == 0 {
		logMissingErrorCode.Println("Google -", resp.Status, "-", string(body))
		return "JADE internal error: 03-03-0007. Please contact the support."
	}

	return chatCompletionResponse.Candidates[0].Content.Parts[0].Text
}