package ai

import (
	"encoding/json"
	"fmt"
)

func repairChatCompletionBody(body []byte) ([]byte, error) {
	var req map[string]any
	if err := json.Unmarshal(body, &req); err != nil {
		return body, fmt.Errorf("repair chat completion body: %w", err)
	}

	rawMessages, ok := req["messages"].([]any)
	if !ok || len(rawMessages) == 0 {
		return body, nil
	}

	repaired := repairToolMessages(rawMessages)
	if sameMessages(rawMessages, repaired) {
		return body, nil
	}

	req["messages"] = repaired
	return json.Marshal(req)
}

func repairToolMessages(messages []any) []any {
	out := make([]any, 0, len(messages)+1)
	for _, raw := range messages {
		msg, ok := raw.(map[string]any)
		if !ok {
			out = append(out, raw)
			continue
		}

		if messageRole(msg) == "tool" && !previousMessageHasToolCalls(out) {
			toolCallID := messageToolCallID(msg)
			if toolCallID == "" {
				toolCallID = "call_repaired"
			}
			out = append(out, map[string]any{
				"role":    "assistant",
				"content": "",
				"tool_calls": []any{
					map[string]any{
						"id":   toolCallID,
						"type": "function",
						"function": map[string]any{
							"name":      "tool",
							"arguments": "{}",
						},
					},
				},
			})
		}

		out = append(out, msg)
	}
	return out
}

func previousMessageHasToolCalls(messages []any) bool {
	if len(messages) == 0 {
		return false
	}
	msg, ok := messages[len(messages)-1].(map[string]any)
	if !ok {
		return false
	}
	if messageRole(msg) != "assistant" {
		return false
	}
	toolCalls, ok := msg["tool_calls"].([]any)
	return ok && len(toolCalls) > 0
}

func messageRole(msg map[string]any) string {
	role, _ := msg["role"].(string)
	return role
}

func messageToolCallID(msg map[string]any) string {
	id, _ := msg["tool_call_id"].(string)
	return id
}

func sameMessages(before, after []any) bool {
	if len(before) != len(after) {
		return false
	}
	b, err := json.Marshal(before)
	if err != nil {
		return false
	}
	a, err := json.Marshal(after)
	if err != nil {
		return false
	}
	return string(b) == string(a)
}

func chatCompletionStreamRequested(body []byte) bool {
	var probe struct {
		Stream *bool `json:"stream"`
	}
	if err := json.Unmarshal(body, &probe); err != nil {
		return false
	}
	return probe.Stream != nil && *probe.Stream
}