package proxy

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"unicode/utf16"

	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
)

const fingerprintSalt = "59cf53e54c78"

func computeFingerprint(firstUserMessage string, version string) string {
	indices := []int{4, 7, 20}
	runes := utf16.Encode([]rune(firstUserMessage))
	var chars string
	for _, i := range indices {
		if i < len(runes) {
			chars += string(rune(runes[i]))
		} else {
			chars += "0"
		}
	}
	input := fingerprintSalt + chars + version
	hash := sha256.Sum256([]byte(input))
	return hex.EncodeToString(hash[:])[:3]
}

func extractFirstUserMessage(body []byte) string {
	messages := gjson.GetBytes(body, "messages")
	if !messages.Exists() || !messages.IsArray() {
		return ""
	}
	for _, msg := range messages.Array() {
		if msg.Get("role").String() != "user" {
			continue
		}
		content := msg.Get("content")
		if content.Type == gjson.String {
			return content.String()
		}
		if content.IsArray() {
			for _, block := range content.Array() {
				if block.Get("type").String() == "text" {
					return block.Get("text").String()
				}
			}
		}
		break
	}
	return ""
}

func buildBillingHeader(body []byte, version string) string {
	userMsg := extractFirstUserMessage(body)
	fp := computeFingerprint(userMsg, version)
	return fmt.Sprintf("x-anthropic-billing-header: cc_version=%s.%s; cc_entrypoint=cli; cch=00000;", version, fp)
}

func injectBillingHeader(body []byte, version string) []byte {
	header := buildBillingHeader(body, version)
	billingBlock := map[string]interface{}{"type": "text", "text": header}
	billingJSON, _ := json.Marshal(billingBlock)

	existing := gjson.GetBytes(body, "system")
	if !existing.Exists() {
		body, _ = sjson.SetRawBytes(body, "system", []byte("["+string(billingJSON)+"]"))
		return body
	}

	if existing.IsArray() {
		items := make([]json.RawMessage, 0, len(existing.Array())+1)
		items = append(items, billingJSON)
		for _, item := range existing.Array() {
			items = append(items, json.RawMessage(item.Raw))
		}
		systemJSON, _ := json.Marshal(items)
		body, _ = sjson.SetRawBytes(body, "system", systemJSON)
		return body
	}

	origText := existing.String()
	origBlock := map[string]string{"type": "text", "text": origText}
	origJSON, _ := json.Marshal(origBlock)
	body, _ = sjson.SetRawBytes(body, "system", []byte("["+string(billingJSON)+","+string(origJSON)+"]"))
	return body
}