package spells

import (
	"fmt"
	"strings"

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

type promptData struct {
	SourceID   string
	HasChunk   bool
	ChunkID    string
	ChunkIndex int
	Units      []promptUnit
}

type promptUnit struct {
	ID       string
	Text     string
	Metadata []promptMetadata
}

type promptMetadata struct {
	Key   string
	Value string
}

var spellsPromptBundle = mustLoadPromptBundle()

func mustLoadPromptBundle() *prompt.Bundle {
	bundle, err := prompt.LoadBundle(embeddedAssets, prompt.Definition{
		PromptID:       PromptID,
		Version:        SchemaVersion,
		EmbeddedPath:   "assets/prompts",
		SystemPath:     "assets/prompts/system.md",
		UserPath:       "assets/prompts/user.md",
		ReferenceSlots: cloneReferenceSlots(referenceSlots),
	})
	if err != nil {
		panic(err)
	}
	return bundle
}

func buildPromptData(req contracts.ExtractionRequest) (promptData, error) {
	if req.Source == nil {
		return promptData{}, fmt.Errorf("dnd spells prompt: source must not be nil")
	}
	if req.Chunk == nil {
		return promptData{}, fmt.Errorf("dnd spells prompt: chunk must not be nil")
	}

	data := promptData{
		SourceID:   req.Source.ID,
		HasChunk:   true,
		ChunkID:    req.Chunk.ID,
		ChunkIndex: req.Chunk.Index,
		Units:      make([]promptUnit, 0, len(req.Chunk.Units)),
	}
	for _, unit := range req.Chunk.Units {
		data.Units = append(data.Units, promptUnit{
			ID:       unit.ID,
			Text:     unit.Text,
			Metadata: selectedMetadata(unit),
		})
	}
	return data, nil
}

func renderPrompt(req contracts.ExtractionRequest) (system string, user string, metadata prompt.Metadata, err error) {
	data, err := buildPromptData(req)
	if err != nil {
		return "", "", prompt.Metadata{}, err
	}
	system, user, metadata, err = spellsPromptBundle.RenderUserSystemWithReferences(data, req.References)
	if err != nil {
		return "", "", prompt.Metadata{}, fmt.Errorf("dnd spells prompt: %w", err)
	}
	return system, user, metadata, nil
}

func selectedMetadata(unit source.SourceUnit) []promptMetadata {
	if len(unit.Metadata) == 0 {
		return nil
	}

	keys := []string{"speaker", "start", "end"}
	metadata := make([]promptMetadata, 0, len(keys))
	for _, key := range keys {
		value, ok := unit.Metadata[key]
		if !ok {
			continue
		}
		rendered := strings.TrimSpace(fmt.Sprint(value))
		if rendered == "" {
			continue
		}
		metadata = append(metadata, promptMetadata{
			Key:   key,
			Value: rendered,
		})
	}
	return metadata
}