package scenes

import (
	"bytes"
	"fmt"
	"io/fs"
	"sort"
	"sync"

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

const scriptoriumPromptRoot = "assets/prompts"

func RegisterPromptAssets(registry *llm.AssetRegistry) error {
	sharedPromptFS, err := promptSharedFS()
	if err != nil {
		return fmt.Errorf("prepare shared prompt assets: %w", err)
	}
	promptFS, err := modulePromptFS(sharedPromptFS)
	if err != nil {
		return fmt.Errorf("prepare scene prompt assets: %w", err)
	}
	if err := registry.RegisterPromptFS(promptFS, scriptoriumPromptRoot); err != nil {
		return err
	}
	return registry.RegisterSchemaFS(embeddedAssets, "assets/schemas")
}

func promptSharedFS() (fs.FS, error) {
	registry := llm.NewAssetRegistry()
	if err := sharedassets.Register(registry); err != nil {
		return nil, err
	}
	return registry.PromptFS()
}

func promptInputs(req contracts.ChunkRequest) contracts.LLMInputSet {
	return contracts.LLMInputSet{
		"transcript": transcriptPromptInput(req.SourceInput),
		"roster":     referencePromptMaterial("roster", req.References.Slots["roster"]),
		"glossary":   referencePromptMaterial("glossary", req.References.Slots["glossary"]),
	}
}

func transcriptPromptInput(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 scriptoriumPromptMetadata() (string, error) {
	scriptoriumPromptHashOnce.Do(func() {
		parts := append([]llm.AssetHashPart{
			{FS: embeddedAssets, Path: "assets/prompts/dnd.scenes.yaml"},
			{FS: embeddedAssets, Path: "assets/prompts/task.md"},
			{FS: embeddedAssets, Path: "assets/prompts/instructions.md"},
		}, append(sharedassets.CommonHashParts(), sharedassets.ReferenceHashParts()...)...)
		scriptoriumPromptHash, scriptoriumPromptHashErr = llm.HashAssets(parts)
	})
	return scriptoriumPromptHash, scriptoriumPromptHashErr
}

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()
}

var (
	scriptoriumPromptHashOnce sync.Once
	scriptoriumPromptHash     string
	scriptoriumPromptHashErr  error
)