package llm

import (
	"bytes"
	"encoding/binary"
	"math"
)

type ChatMessage struct {
	Role     string
	Contents []ChatMessageContent
}

type ChatMessageContent struct {
	Type    string // text, image, audio, video
	Content string
}

type ChatConfig struct {
	defaultConfig *ChatConfig
	Model         string
	Ratio         float64
	MaxTokens     int
	Temperature   float64
	TopP          float64
	Tools         map[string]any
}

func (chatConfig *ChatConfig) SetDefault(config *ChatConfig) {
	chatConfig.defaultConfig = config
}

func (chatConfig *ChatConfig) GetModel() string {
	if chatConfig.Model == "" && chatConfig.defaultConfig != nil {
		return chatConfig.defaultConfig.Model
	}
	return chatConfig.Model
}

func (chatConfig *ChatConfig) GetMaxTokens() int {
	if chatConfig.MaxTokens == 0 && chatConfig.defaultConfig != nil {
		return chatConfig.defaultConfig.MaxTokens
	}
	return chatConfig.MaxTokens
}

func (chatConfig *ChatConfig) GetTemperature() float64 {
	if chatConfig.Temperature == 0 && chatConfig.defaultConfig != nil {
		return chatConfig.defaultConfig.Temperature
	}
	return chatConfig.Temperature
}

func (chatConfig *ChatConfig) GetTopP() float64 {
	if chatConfig.TopP == 0 && chatConfig.defaultConfig != nil {
		return chatConfig.defaultConfig.TopP
	}
	return chatConfig.TopP
}

func (chatConfig *ChatConfig) GetTools() map[string]any {
	if chatConfig.Tools == nil && chatConfig.defaultConfig != nil {
		return chatConfig.defaultConfig.Tools
	}
	return chatConfig.Tools
}

type Usage struct {
	AskTokens    int64
	AnswerTokens int64
	TotalTokens  int64
	UsedTime     int64
}

type MessagesMaker struct {
	list []ChatMessage
}

func Messages() *MessagesMaker {
	return &MessagesMaker{
		list: make([]ChatMessage, 0),
	}
}

func (m *MessagesMaker) Make() []ChatMessage {
	return m.list
}

func (m *MessagesMaker) User(contents ...ChatMessageContent) *MessagesMaker {
	m.list = append(m.list, ChatMessage{
		Role:     RoleUser,
		Contents: contents,
	})
	return m
}

func (m *MessagesMaker) Assistant(contents ...ChatMessageContent) *MessagesMaker {
	m.list = append(m.list, ChatMessage{
		Role:     RoleAssistant,
		Contents: contents,
	})
	return m
}

func (m *MessagesMaker) System(contents ...ChatMessageContent) *MessagesMaker {
	m.list = append(m.list, ChatMessage{
		Role:     RoleSystem,
		Contents: contents,
	})
	return m
}

func (m *MessagesMaker) Tool(contents ...ChatMessageContent) *MessagesMaker {
	m.list = append(m.list, ChatMessage{
		Role:     RoleTool,
		Contents: contents,
	})
	return m
}

func (m *MessagesMaker) Text(text string) *MessagesMaker {
	if len(m.list) > 0 {
		lastIndex := len(m.list) - 1
		m.list[lastIndex].Contents = append(m.list[lastIndex].Contents, ChatMessageContent{
			Type:    TypeText,
			Content: text,
		})
	}
	return m
}

func (m *MessagesMaker) Image(text string) *MessagesMaker {
	if len(m.list) > 0 {
		lastIndex := len(m.list) - 1
		m.list[lastIndex].Contents = append(m.list[lastIndex].Contents, ChatMessageContent{
			Type:    TypeText,
			Content: text,
		})
	}
	return m
}

func (m *MessagesMaker) Video(text string) *MessagesMaker {
	if len(m.list) > 0 {
		lastIndex := len(m.list) - 1
		m.list[lastIndex].Contents = append(m.list[lastIndex].Contents, ChatMessageContent{
			Type:    TypeText,
			Content: text,
		})
	}
	return m
}

func bin2float64(in []byte) []float64 {
	buf := bytes.NewBuffer(in)
	out := make([]float64, len(in)/4)
	for i := 0; i < len(out); i++ {
		var f float32
		_ = binary.Read(buf, binary.LittleEndian, &f)
		out[i] = float64(f)
	}
	return out
}

func Similarity(buf1, buf2 []byte) float64 {
	a := bin2float64(buf1)
	b := bin2float64(buf2)
	if len(a) != len(b) {
		return 0
	}

	var dotProduct, magnitudeA, magnitudeB float64
	for i := 0; i < len(a); i++ {
		dotProduct += a[i] * b[i]
		magnitudeA += a[i] * a[i]
		magnitudeB += b[i] * b[i]
	}

	magnitudeA = math.Sqrt(magnitudeA)
	magnitudeB = math.Sqrt(magnitudeB)

	if magnitudeA == 0 || magnitudeB == 0 {
		return 0
	}

	return dotProduct / (magnitudeA * magnitudeB)
}