package main

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

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

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  int `json:"input_tokens"`
	OutputTokens int `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)

	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)

	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)

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

func addAnthropicMessage(modelID string, selected bool) string {
	Messages := getAllMessages()

	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")
		collection := mongoClient.Database("chat").Collection("messages")
		response, err := collection.InsertOne(context.Background(), bson.M{"message": "no response", "role": "bot", "date": time.Now(), "model": modelID, "selected": selected})
		if err != nil {
			fmt.Println(err)
		}
		return response.InsertedID.(primitive.ObjectID).Hex()
	} else {
		Content := chatCompletion.Content[0].Text
		collection := mongoClient.Database("chat").Collection("messages")
		response, err := collection.InsertOne(context.Background(), bson.M{"message": Content, "role": "bot", "date": time.Now(), "model": modelID, "selected": selected})
		if err != nil {
			fmt.Println(err)
		}
		return response.InsertedID.(primitive.ObjectID).Hex()
	}
	return ""
}

func Messages2AnthropicMessages(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 RequestAnthropic(model string, messages []Message, maxTokens int, temperature float64) (AnthropicChatCompletionResponse, error) {
	apiKey := "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" // TODO Use env variable
	url := "https://api.anthropic.com/v1/messages"

	AnthropicMessages := Messages2AnthropicMessages(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)
	}

	return chatCompletionResponse, nil
}