package pipeline

import (
	"fmt"
	"sort"
	"strings"

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

type ModuleStage string

const (
	StageInput     ModuleStage = "input"
	StageChunk     ModuleStage = "chunk"
	StageExtract   ModuleStage = "extract"
	StageMerge     ModuleStage = "merge"
	StageNormalize ModuleStage = "normalize"
	StageValidate  ModuleStage = "validate"
	StageOutput    ModuleStage = "output"
)

type ModuleSpec struct {
	Key            string
	Stage          ModuleStage
	ExecutionClass contracts.ExecutionClass
	ArtifactKind   contracts.ArtifactKind
	Provides       []string
	Requires       []string
	ReferenceSlots []contracts.ReferenceSlot
}

func defaultModuleSpec(key string, stage ModuleStage) ModuleSpec {
	return ModuleSpec{
		Key:            key,
		Stage:          stage,
		ExecutionClass: contracts.ExecutionClassDeterministic,
	}
}

func normalizeModuleSpec(spec ModuleSpec) ModuleSpec {
	executionClass := contracts.ExecutionClass(strings.TrimSpace(string(spec.ExecutionClass)))
	return ModuleSpec{
		Key:            strings.TrimSpace(spec.Key),
		Stage:          spec.Stage,
		ExecutionClass: executionClass,
		ArtifactKind:   normalizeArtifactKind(spec.ArtifactKind),
		Provides:       normalizeCapabilities(spec.Provides),
		Requires:       normalizeCapabilities(spec.Requires),
		ReferenceSlots: normalizeReferenceSlots(spec.ReferenceSlots),
	}
}

func normalizeCapabilities(values []string) []string {
	if len(values) == 0 {
		return nil
	}

	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		normalized := strings.TrimSpace(value)
		if normalized == "" {
			continue
		}
		seen[normalized] = struct{}{}
	}
	if len(seen) == 0 {
		return nil
	}

	capabilities := make([]string, 0, len(seen))
	for value := range seen {
		capabilities = append(capabilities, value)
	}
	sort.Strings(capabilities)
	return capabilities
}

func cloneModuleSpec(spec ModuleSpec) ModuleSpec {
	return ModuleSpec{
		Key:            spec.Key,
		Stage:          spec.Stage,
		ExecutionClass: spec.ExecutionClass,
		ArtifactKind:   spec.ArtifactKind,
		Provides:       append([]string(nil), spec.Provides...),
		Requires:       append([]string(nil), spec.Requires...),
		ReferenceSlots: contracts.CloneReferenceSlots(spec.ReferenceSlots),
	}
}

func validateModuleSpec(kind string, expectedStage ModuleStage, spec ModuleSpec) error {
	if spec.Key == "" {
		return fmt.Errorf("%s key must not be empty", kind)
	}
	if spec.Stage != expectedStage {
		return fmt.Errorf("%s %q must use %q stage, got %q", kind, spec.Key, expectedStage, spec.Stage)
	}
	if spec.ExecutionClass == "" {
		return fmt.Errorf("%s %q execution class must not be empty", kind, spec.Key)
	}
	if spec.ExecutionClass != contracts.ExecutionClassDeterministic && spec.ExecutionClass != contracts.ExecutionClassLLMBacked {
		return fmt.Errorf("%s %q has unsupported execution class %q", kind, spec.Key, spec.ExecutionClass)
	}
	if spec.ArtifactKind != "" && spec.Stage != StageExtract && spec.Stage != StageMerge && spec.Stage != StageNormalize {
		return fmt.Errorf("%s %q must not declare an artifact kind", kind, spec.Key)
	}
	if !referenceSlotStage(spec.Stage) && len(spec.ReferenceSlots) > 0 {
		return fmt.Errorf("%s %q must not declare reference slots", kind, spec.Key)
	}
	if err := validateReferenceSlots(spec.ReferenceSlots); err != nil {
		return fmt.Errorf("%s %q reference slots: %w", kind, spec.Key, err)
	}
	return nil
}

func referenceSlotStage(stage ModuleStage) bool {
	return stage == StageChunk || stage == StageExtract || stage == StageMerge || stage == StageNormalize
}

func sortedRegistryKeys[C any](constructors map[string]C) []string {
	if len(constructors) == 0 {
		return nil
	}

	keys := make([]string, 0, len(constructors))
	for key := range constructors {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}

func sortArtifactKinds(kinds []contracts.ArtifactKind) {
	sort.Slice(kinds, func(i, j int) bool { return kinds[i] < kinds[j] })
}

func normalizeReferenceSlots(slots []contracts.ReferenceSlot) []contracts.ReferenceSlot {
	if len(slots) == 0 {
		return nil
	}

	normalized := make([]contracts.ReferenceSlot, 0, len(slots))
	for _, slot := range slots {
		slot.Name = strings.TrimSpace(slot.Name)
		slot.Description = strings.TrimSpace(slot.Description)
		slot.AcceptedMediaTypes = normalizeStringSet(slot.AcceptedMediaTypes)
		slot.AcceptedArtifactKinds = normalizeArtifactKinds(slot.AcceptedArtifactKinds)
		normalized = append(normalized, slot)
	}
	sort.SliceStable(normalized, func(i, j int) bool {
		return normalized[i].Name < normalized[j].Name
	})
	return normalized
}

func normalizeArtifactKinds(values []contracts.ArtifactKind) []contracts.ArtifactKind {
	if len(values) == 0 {
		return nil
	}
	seen := make(map[contracts.ArtifactKind]struct{}, len(values))
	containsEmpty := false
	for _, value := range values {
		value = normalizeArtifactKind(value)
		if value != "" {
			seen[value] = struct{}{}
		} else {
			containsEmpty = true
		}
	}
	if len(seen) == 0 && !containsEmpty {
		return nil
	}
	result := make([]contracts.ArtifactKind, 0, len(seen))
	for value := range seen {
		result = append(result, value)
	}
	if containsEmpty {
		result = append(result, "")
	}
	sort.Slice(result, func(i, j int) bool { return result[i] < result[j] })
	return result
}

func normalizeStringSet(values []string) []string {
	if len(values) == 0 {
		return nil
	}
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		normalized := strings.TrimSpace(value)
		if normalized == "" {
			continue
		}
		seen[normalized] = struct{}{}
	}
	if len(seen) == 0 {
		return nil
	}

	out := make([]string, 0, len(seen))
	for value := range seen {
		out = append(out, value)
	}
	sort.Strings(out)
	return out
}

func validateReferenceSlots(slots []contracts.ReferenceSlot) error {
	seen := make(map[string]struct{}, len(slots))
	for i, slot := range slots {
		if slot.Name == "" {
			return fmt.Errorf("slot[%d].name must not be empty", i)
		}
		if _, ok := seen[slot.Name]; ok {
			return fmt.Errorf("slot name %q is duplicated", slot.Name)
		}
		seen[slot.Name] = struct{}{}
		if slot.MaxBytes < 0 {
			return fmt.Errorf("slot %q max_bytes must not be negative", slot.Name)
		}
		for _, kind := range slot.AcceptedArtifactKinds {
			if normalizeArtifactKind(kind) == "" {
				return fmt.Errorf("slot %q accepted artifact kind must not be empty", slot.Name)
			}
		}
	}
	return nil
}