package llm

import (
	"bytes"
	"fmt"
	"path/filepath"
	"sort"

	ck "go.maximumdirect.net/configkit"
	pk "go.maximumdirect.net/promptkit"
	"gopkg.in/yaml.v3"
)

func capturedSourcePaths(engine *pk.Engine) ([]string, error) {
	inspection, err := engine.InspectConfiguration()
	if err != nil {
		return nil, fmt.Errorf("inspect configuration: %w", err)
	}

	sources := inspection.Sources()
	hostByID := make(map[string]string, len(sources))
	paths := make(map[string]struct{}, len(sources))
	for _, source := range sources {
		if source.Kind == ck.SourceKindFile && source.Path != "" {
			hostByID[source.ID] = source.Path
			paths[source.Path] = struct{}{}
		}
	}

	if len(hostByID) == 0 {
		return []string{}, nil
	}

	type promptList []struct {
		Messages []struct {
			ContentFile *string `yaml:"content_file"`
		} `yaml:"messages"`
	}
	var effective struct {
		Prompts     promptList `yaml:"prompts"`
		Collections map[string]struct {
			Prompts promptList `yaml:"prompts"`
		} `yaml:"collections"`
	}
	var projection bytes.Buffer
	if err := inspection.WriteEffective(&projection); err != nil {
		return nil, fmt.Errorf("write effective PromptKit configuration: %w", err)
	}
	if err := yaml.Unmarshal(projection.Bytes(), &effective); err != nil {
		return nil, fmt.Errorf("decode effective PromptKit message paths: %w", err)
	}
	metadata := inspection.Metadata()
	capture := func(base ck.Path, prompts promptList) {
		for i, prompt := range prompts {
			messages := base.Index(i).Key("messages")
			for j, message := range prompt.Messages {
				if message.ContentFile == nil {
					continue
				}
				field, ok := metadata.Lookup(messages.Index(j).Key("content_file"))
				if !ok || len(field.Origin.Locations) != 1 {
					continue
				}
				if host := hostByID[field.Origin.Locations[0].SourceID]; host != "" {
					paths[filepath.Join(filepath.Dir(host), filepath.FromSlash(*message.ContentFile))] = struct{}{}
				}
			}
		}
	}
	capture(ck.Keys("prompts"), effective.Prompts)
	for name, collection := range effective.Collections {
		capture(ck.Keys("collections", name, "prompts"), collection.Prompts)
	}

	result := make([]string, 0, len(paths))
	for path := range paths {
		result = append(result, path)
	}
	sort.Strings(result)
	return result, nil
}