package shared

import (
	"bytes"
	"fmt"
	"sort"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func PromptInputs(sourceInput contracts.LLMInputMaterial, references contracts.ReferenceSet) contracts.LLMInputSet {
	partySlot := references.Slots["party"]
	if len(partySlot.Items) == 0 {
		partySlot = references.Slots["roster"]
	}
	return contracts.LLMInputSet{
		"transcript": TranscriptPromptMaterial(sourceInput),
		"players":    ReferencePromptMaterial("players", references.Slots["players"]),
		"party":      ReferencePromptMaterial("party", partySlot),
		"glossary":   ReferencePromptMaterial("glossary", references.Slots["glossary"]),
	}
}

func TranscriptPromptMaterial(material contracts.LLMInputMaterial) contracts.LLMInputMaterial {
	out := material.Clone()
	out.Name = "transcript"
	return out
}

func ReferencePromptMaterial(name string, slot contracts.ResolvedReferenceSlot) contracts.LLMInputMaterial {
	body := ReferencePromptInput(slot)
	digest := ""
	originURI := ""
	if len(slot.Items) == 1 {
		digest = slot.Items[0].Digest
		originURI = slot.Items[0].Origin.URI
	}
	return contracts.NewLLMInputMaterial(name, "text/plain", body, digest, originURI)
}

func ReferencePromptInput(slot contracts.ResolvedReferenceSlot) []byte {
	if len(slot.Items) == 0 {
		return []byte(" ")
	}
	items := append([]contracts.ReferenceItem(nil), slot.Items...)
	sort.SliceStable(items, func(i, j int) bool {
		if items[i].Origin.URI != items[j].Origin.URI {
			return items[i].Origin.URI < items[j].Origin.URI
		}
		if items[i].Digest != items[j].Digest {
			return items[i].Digest < items[j].Digest
		}
		return string(items[i].Content) < string(items[j].Content)
	})
	if len(items) == 1 {
		return append([]byte(nil), items[0].Content...)
	}

	var b bytes.Buffer
	for i, item := range items {
		if i > 0 {
			b.WriteString("\n\n")
		}
		fmt.Fprintf(&b, "Reference %d\n", i+1)
		if item.Origin.Type != "" {
			fmt.Fprintf(&b, "Origin-Type: %s\n", item.Origin.Type)
		}
		if item.Origin.URI != "" {
			fmt.Fprintf(&b, "Origin-URI: %s\n", item.Origin.URI)
		}
		if item.Digest != "" {
			fmt.Fprintf(&b, "Digest: %s\n", item.Digest)
		}
		if item.MediaType != "" {
			fmt.Fprintf(&b, "Media-Type: %s\n", item.MediaType)
		}
		if item.SizeBytes > 0 {
			fmt.Fprintf(&b, "Size-Bytes: %d\n", item.SizeBytes)
		}
		b.WriteString("\n")
		b.Write(item.Content)
	}
	return b.Bytes()
}