package main

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

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

type AnthropicChatCompletionRequest struct {
	Model       string             `json:"model"`
	Messages    []AnthropicMessage `json:"messages"`
	MaxTokens   int                `json:"max_tokens"`
	Temperature float64            `json:"temperature"`
}

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

type AnthropicChatCompletionResponse struct {
	ID           string                 `json:"id"`
	Content      []AnthropicContentItem `json:"content"`
	Model        string                 `json:"model"`
	StopReason   string                 `json:"stop_reason"`
	StopSequence string                 `json:"stop_sequence"`
	Usage        AnthropicUsage         `json:"usage"`
}

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

type AnthropicUsage struct {
	InputTokens  int32 `json:"input_tokens"`
	OutputTokens int32 `json:"output_tokens"`
}

func init() {
	var ModelInfosList = []ModelInfo{}

	modelInfo := ModelInfo{
		ID:          "claude-3-haiku-20240307",
		Name:        "Haiku",
		Icon:        "anthropic",
		MaxToken:    8192,
		InputPrice:  0.50 / 1000000,
		OutputPrice: 1.50 / 1000000,
	}
	ModelInfosList = append(ModelInfosList, modelInfo)
	ModelsInfos = append(ModelsInfos, modelInfo)

	modelInfo = ModelInfo{
		ID:          "claude-3-sonnet-20240229",
		Name:        "Sonnet",
		Icon:        "anthropic",
		MaxToken:    8192,
		InputPrice:  0.50 / 1000000,
		OutputPrice: 1.50 / 1000000,
	}
	ModelInfosList = append(ModelInfosList, modelInfo)
	ModelsInfos = append(ModelsInfos, modelInfo)

	modelInfo = ModelInfo{
		ID:          "claude-3-opus-20240229",
		Name:        "Opus",
		Icon:        "anthropic",
		MaxToken:    8192,
		InputPrice:  0.50 / 1000000,
		OutputPrice: 1.50 / 1000000,
	}
	ModelInfosList = append(ModelInfosList, modelInfo)
	ModelsInfos = append(ModelsInfos, modelInfo)

	companyInfo := CompanyInfo{
		ID:         "anthropic",
		Name:       "Anthropic",
		ModelInfos: ModelInfosList,
	}
	CompanyInfos = append(CompanyInfos, companyInfo)
}

func addAnthropicMessage(modelID string, selected bool) edgedb.UUID {
	Messages := getAllSelectedMessages()

	chatCompletion, err := RequestAnthropic(modelID, Messages, 2048, 0.7) // TODO CHange parameters
	if err != nil {
		// Print error
		fmt.Println("Error:", err)
	} else if len(chatCompletion.Content) == 0 {
		fmt.Println(chatCompletion)
		fmt.Println("No response from Anthropic")
		id := insertBotMessage("No response from Anthropic", selected, modelID)
		return id
	} else {
		Content := chatCompletion.Content[0].Text
		id := insertBotMessage(Content, selected, modelID)
		return id
	}
	return edgedb.UUID{}
}

func EdgeMessages2AnthropicMessages(messages []Message) []AnthropicMessage {
	AnthropicMessages := make([]AnthropicMessage, len(messages))
	for i, msg := range messages {
		var role string
		switch msg.Role {
		case "user":
			role = "user"
		case "bot":
			role = "assistant"
		default:
			role = "system"
		}
		AnthropicMessages[i] = AnthropicMessage{
			Role:    role,
			Content: msg.Content,
		}
	}
	return AnthropicMessages
}

func TestAnthropicKey(apiKey string) bool {
	url := "https://api.anthropic.com/v1/messages"

	AnthropicMessages := []AnthropicMessage{
		{
			Role:    "user",
			Content: "Hello",
		},
	}

	requestBody := AnthropicChatCompletionRequest{
		Model:       "claude-3-haiku-20240307",
		Messages:    AnthropicMessages,
		MaxTokens:   10,
		Temperature: 0,
	}

	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")
	req.Header.Set("anthropic-version", "2023-06-01")
	req.Header.Set("x-api-key", apiKey)

	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 AnthropicChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return false
	}
	if chatCompletionResponse.Content == nil {
		return false
	}
	return true
}

func RequestAnthropic(model string, messages []Message, maxTokens int, temperature float64) (AnthropicChatCompletionResponse, error) {
	var apiKey string
	err := edgeClient.QuerySingle(edgeCtx, `
		with
			filtered_keys := (
				select Key {
					key
				} filter .company = <str>$0
			)
		select filtered_keys.key limit 1
	`, &apiKey, "anthropic")
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error getting OpenAI API key: %w", err)
	}

	url := "https://api.anthropic.com/v1/messages"

	AnthropicMessages := EdgeMessages2AnthropicMessages(messages)

	requestBody := AnthropicChatCompletionRequest{
		Model:       model,
		Messages:    AnthropicMessages,
		MaxTokens:   maxTokens,
		Temperature: temperature,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error marshaling JSON: %w", err)
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error creating request: %w", err)
	}

	req.Header.Set("content-Type", "application/json")
	req.Header.Set("anthropic-version", "2023-06-01")
	req.Header.Set("x-api-key", apiKey)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error sending request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error reading response body: %w", err)
	}

	var chatCompletionResponse AnthropicChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err)
	}

	var usedModelInfo ModelInfo
	for mi := range ModelsInfos {
		if ModelsInfos[mi].ID == model {
			usedModelInfo = ModelsInfos[mi]
		}
	}
	var inputCost float32 = float32(chatCompletionResponse.Usage.InputTokens) * usedModelInfo.InputPrice
	var outputCost float32 = float32(chatCompletionResponse.Usage.OutputTokens) * usedModelInfo.OutputPrice
	addUsage(inputCost, outputCost, chatCompletionResponse.Usage.InputTokens, chatCompletionResponse.Usage.OutputTokens, model)

	return chatCompletionResponse, nil
}