package kimi_linear

import (
	"encoding/json"
	"strings"

	"makarna/pkg/chat"
)

type kimiRenderer struct{}

func (r *kimiRenderer) Render(messages []chat.Message, opts chat.Options) (string, error) {
	var sb strings.Builder

	if len(opts.Tools) > 0 {
		b, err := json.Marshal(opts.Tools)
		if err != nil {
			return "", err
		}
		sb.WriteString("<|im_system|>tool_declare<|im_middle|")
		sb.Write(b)
		sb.WriteString("<|im_end|>")
	}

	for i := 0; i < len(messages); i++ {
		m := messages[i]
		roleName := m.Role
		if roleName == "" {
			roleName = "user"
		}

		switch m.Role {
		case "user":
			sb.WriteString("<|im_user|")
			sb.WriteString(roleName)
			sb.WriteString("<|im_middle|>")
		case "assistant":
			sb.WriteString("<|im_assistant|")
			sb.WriteString(roleName)
			sb.WriteString("<|im_middle|>")
		case "tool":
			sb.WriteString("<|im_system|")
			sb.WriteString(roleName)
			sb.WriteString("<|im_middle|>")
		default:
			sb.WriteString("<|im_system|")
			sb.WriteString(roleName)
			sb.WriteString("<|im_middle|>")
		}

		if m.Role == "assistant" && len(m.ToolCalls) > 0 {
			sb.WriteString(m.Content)
			sb.WriteString("<|tool_calls_section_begin|>")
			for _, tc := range m.ToolCalls {
				id := tc.Name
				sb.WriteString("<|tool_call_begin|>")
				sb.WriteString(id)
				sb.WriteString("<|tool_call_argument_begin|>")
				if len(tc.Arguments) == 0 {
					sb.WriteString("{}")
				} else {
					sb.Write(tc.Arguments)
				}
				sb.WriteString("<|tool_call_end|>")
			}
			sb.WriteString("<|tool_calls_section_end|>")
		} else if m.Role == "tool" {
			sb.WriteString(m.Content)
		} else {
			sb.WriteString(m.Content)
		}

		sb.WriteString("<|im_end|>")
	}

	if opts.AddGenerationPrompt {
		sb.WriteString("<|im_assistant|>assistant<|im_middle|>")
	}

	return sb.String(), nil
}

func init() {
	chat.Register("KimiLinearForCausalLM", &kimiRenderer{})
	chat.Register("kimi_linear", &kimiRenderer{})
	chat.Register("kimi", &kimiRenderer{})
}