package pipeline

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"

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

type referenceTargetKey struct {
	StepID string
	LaneID string
	Stage  ModuleStage
}

func keyForReferenceTarget(target ResolvedReferenceTarget) referenceTargetKey {
	return referenceTargetKey{StepID: target.StepID, LaneID: target.LaneID, Stage: target.Stage}
}

func operationReferenceSet(input RunInput, target ResolvedReferenceTarget) contracts.ReferenceSet {
	if input.references != nil {
		if set, ok := input.references[keyForReferenceTarget(target)]; ok {
			return CloneReferenceSet(set)
		}
	}
	return CloneReferenceSet(target.ReferenceSet)
}

// buildStepReferenceSets resolves every generated binding for a step before
// any lane in that step is allowed to start. Each returned set is a fresh
// operation-time view; prepared reference sets are never modified.
func buildStepReferenceSets(input RunInput, step PreparedPipelineStep, outputs []contracts.SerializedOutput) (map[referenceTargetKey]contracts.ReferenceSet, []artifacts.ReferenceProvenance, error) {
	if len(step.lanes) == 0 {
		return nil, nil, nil
	}
	sets := make(map[referenceTargetKey]contracts.ReferenceSet)
	var provenance []artifacts.ReferenceProvenance
	canonical := make(map[string]contracts.ReferenceItem)
	for _, prepared := range step.lanes {
		lane := prepared.resolved
		for _, target := range []ResolvedReferenceTarget{lane.ExtractReferences, lane.MergeReferences, lane.NormalizeReferences} {
			generated := false
			resolved := CloneReferenceSet(target.ReferenceSet)
			for _, binding := range target.Bindings {
				if binding.Artifact == nil {
					continue
				}
				generated = true
				item, err := generatedReferenceItem(input, binding, outputs, canonical)
				if err != nil {
					return nil, nil, contextualHandoffError(input, target, binding, err)
				}
				slotName := strings.TrimSpace(binding.SlotName)
				slot, ok := resolved.Slots[slotName]
				if !ok {
					return nil, nil, contextualHandoffError(input, target, binding, fmt.Errorf("reference slot %q has no resolved declaration", slotName))
				}
				if len(slot.Items) > 0 {
					return nil, nil, contextualHandoffError(input, target, binding, fmt.Errorf("reference slot %q already contains materialized items", slotName))
				}
				if len(slot.Slot.AcceptedArtifactKinds) == 0 || !containsArtifactKind(slot.Slot.AcceptedArtifactKinds, item.ArtifactKind) {
					return nil, nil, contextualHandoffError(input, target, binding, fmt.Errorf("reference slot %q does not accept artifact kind %q", slotName, item.ArtifactKind))
				}
				if !referenceMediaTypeAccepted(item.MediaType, slot.Slot.AcceptedMediaTypes) {
					return nil, nil, contextualHandoffError(input, target, binding, fmt.Errorf("reference slot %q does not accept media type %q", slotName, item.MediaType))
				}
				if slot.Slot.MaxBytes > 0 && item.SizeBytes > slot.Slot.MaxBytes {
					return nil, nil, contextualHandoffError(input, target, binding, fmt.Errorf("reference slot %q is %d bytes, limit %d", slotName, item.SizeBytes, slot.Slot.MaxBytes))
				}
				slot.Items = []contracts.ReferenceItem{contracts.CloneReferenceItem(item)}
				resolved.Slots[slotName] = slot
			}
			if !generated {
				continue
			}
			key := keyForReferenceTarget(target)
			sets[key] = resolved
			for _, item := range generatedItems(resolved) {
				provenance = append(provenance, referenceProvenanceForItem(target, item))
			}
		}
	}
	return sets, provenance, nil
}

func containsArtifactKind(kinds []contracts.ArtifactKind, want contracts.ArtifactKind) bool {
	for _, kind := range kinds {
		if kind == want {
			return true
		}
	}
	return false
}

func contextualHandoffError(input RunInput, target ResolvedReferenceTarget, binding ReferenceBinding, err error) error {
	return fmt.Errorf("pipeline %q step %q lane %q %s generated dependency %q from %q/%q: %w", input.pipeline.ID, target.StepID, target.LaneID, target.Stage, binding.SlotName, binding.Artifact.Step, binding.Artifact.Lane, err)
}

func generatedReferenceItem(input RunInput, binding ReferenceBinding, outputs []contracts.SerializedOutput, cache map[string]contracts.ReferenceItem) (contracts.ReferenceItem, error) {
	selector := binding.Artifact
	if selector == nil {
		return contracts.ReferenceItem{}, fmt.Errorf("generated reference selector must not be nil")
	}
	stepID := strings.TrimSpace(selector.Step)
	laneID := strings.TrimSpace(selector.Lane)
	cacheKey := stepID + "\x00" + laneID
	if item, ok := cache[cacheKey]; ok {
		item.SlotName = strings.TrimSpace(binding.SlotName)
		item.BindingSource = strings.TrimSpace(binding.BindingSource)
		return contracts.CloneReferenceItem(item), nil
	}
	matches := make([]contracts.SerializedOutput, 0, 1)
	for _, output := range outputs {
		if strings.TrimSpace(output.StepID) == stepID && strings.TrimSpace(output.LaneID) == laneID {
			matches = append(matches, output)
		}
	}
	if len(matches) == 0 {
		return contracts.ReferenceItem{}, fmt.Errorf("producer has no accepted normalized output")
	}
	if len(matches) > 1 {
		return contracts.ReferenceItem{}, fmt.Errorf("producer has %d accepted normalized outputs; exactly one is required", len(matches))
	}
	producer, ok := findResolvedLane(input.pipeline, stepID, laneID)
	if !ok {
		return contracts.ReferenceItem{}, fmt.Errorf("producer lane is not present in the resolved pipeline")
	}
	if input.Prepared == nil || input.Prepared.artifactCodecs == nil {
		return contracts.ReferenceItem{}, fmt.Errorf("artifact codec registry is unavailable")
	}
	serialized := contracts.CloneSerializedArtifact(matches[0].Artifact)
	if serialized.Kind != producer.ArtifactKind {
		return contracts.ReferenceItem{}, fmt.Errorf("producer artifact kind %q does not match resolved kind %q", serialized.Kind, producer.ArtifactKind)
	}
	value, err := input.Prepared.artifactCodecs.Decode(serialized)
	if err != nil {
		return contracts.ReferenceItem{}, fmt.Errorf("decode producer artifact: %w", err)
	}
	canonical, err := input.Prepared.artifactCodecs.Encode(producer.ArtifactKind, value)
	if err != nil {
		return contracts.ReferenceItem{}, fmt.Errorf("canonicalize producer artifact: %w", err)
	}
	expected, ok := input.Prepared.artifactCodecs.Spec(producer.ArtifactKind)
	if !ok {
		return contracts.ReferenceItem{}, fmt.Errorf("producer artifact codec %q is not registered", producer.ArtifactKind)
	}
	if canonical.Kind != expected.Kind || canonical.Schema.ID != expected.Schema.ID || canonical.Schema.Name != expected.Schema.Name || canonical.Schema.Version != expected.Schema.Version || contracts.DigestArtifactSchema(canonical.Schema) != expected.SchemaDigest || canonical.MediaType != expected.MediaType {
		return contracts.ReferenceItem{}, fmt.Errorf("canonical producer artifact does not match registered codec identity")
	}
	item := contracts.ReferenceItem{
		SlotName:       strings.TrimSpace(binding.SlotName),
		MediaType:      canonical.MediaType,
		Content:        append([]byte(nil), canonical.Content...),
		Digest:         referenceDigest(canonical.Content),
		Origin:         contracts.ReferenceOrigin{Type: "generated"},
		SizeBytes:      int64(len(canonical.Content)),
		BindingSource:  strings.TrimSpace(binding.BindingSource),
		ArtifactKind:   canonical.Kind,
		ArtifactSchema: contracts.CloneArtifactSchema(canonical.Schema),
		Producer: contracts.ReferenceProducer{
			PipelineID: input.pipeline.ID,
			StepID:     stepID,
			LaneID:     laneID,
			ModuleKey:  producer.Normalize.Module,
		},
	}
	cacheItem := contracts.CloneReferenceItem(item)
	cacheItem.SlotName = ""
	cacheItem.BindingSource = ""
	cache[cacheKey] = cacheItem
	return item, nil
}

func findResolvedLane(pipeline ResolvedPipeline, stepID, laneID string) (ResolvedArtifactLane, bool) {
	for _, step := range pipeline.Steps {
		if step.ID != stepID {
			continue
		}
		for _, lane := range step.ArtifactLanes {
			if lane.ID == laneID {
				return lane, true
			}
		}
	}
	return ResolvedArtifactLane{}, false
}

func generatedItems(set contracts.ReferenceSet) []contracts.ReferenceItem {
	var items []contracts.ReferenceItem
	keys := make([]string, 0, len(set.Slots))
	for key := range set.Slots {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	for _, key := range keys {
		for _, item := range set.Slots[key].Items {
			if item.Origin.Type == "generated" {
				items = append(items, item)
			}
		}
	}
	return items
}

func referenceProvenanceForItem(target ResolvedReferenceTarget, item contracts.ReferenceItem) artifacts.ReferenceProvenance {
	return artifacts.ReferenceProvenance{
		Stage:            string(target.Stage),
		StepID:           target.StepID,
		LaneID:           target.LaneID,
		SlotName:         item.SlotName,
		OriginType:       item.Origin.Type,
		OriginURI:        item.Origin.URI,
		Digest:           item.Digest,
		MediaType:        item.MediaType,
		SizeBytes:        item.SizeBytes,
		BindingSource:    item.BindingSource,
		ArtifactKind:     string(item.ArtifactKind),
		SchemaID:         item.ArtifactSchema.ID,
		SchemaName:       item.ArtifactSchema.Name,
		SchemaVersion:    item.ArtifactSchema.Version,
		SchemaDigest:     contracts.DigestArtifactSchema(item.ArtifactSchema),
		ProducerPipeline: item.Producer.PipelineID,
		ProducerStep:     item.Producer.StepID,
		ProducerLane:     item.Producer.LaneID,
		ProducerModule:   item.Producer.ModuleKey,
	}
}

type generatedReferenceFingerprintIdentity struct {
	ProducerPipeline string `json:"producer_pipeline_id"`
	ProducerStep     string `json:"producer_step_id"`
	ProducerLane     string `json:"producer_lane_id"`
	ProducerModule   string `json:"producer_module_key"`
	ArtifactKind     string `json:"artifact_kind"`
	SchemaID         string `json:"schema_id"`
	SchemaName       string `json:"schema_name"`
	SchemaVersion    string `json:"schema_version"`
	SchemaDigest     string `json:"schema_digest"`
	MediaType        string `json:"media_type"`
	ContentDigest    string `json:"content_digest"`
}

// generatedReferenceDependencies returns the canonical semantic dependency
// identity that receiving operation checkpoints must include.
func generatedReferenceDependencies(set contracts.ReferenceSet) []CheckpointFingerprint {
	var fingerprints []CheckpointFingerprint
	keys := make([]string, 0, len(set.Slots))
	for key := range set.Slots {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	for _, slotName := range keys {
		for index, item := range set.Slots[slotName].Items {
			if item.Origin.Type != "generated" {
				continue
			}
			identity := generatedReferenceFingerprintIdentity{
				ProducerPipeline: item.Producer.PipelineID,
				ProducerStep:     item.Producer.StepID,
				ProducerLane:     item.Producer.LaneID,
				ProducerModule:   item.Producer.ModuleKey,
				ArtifactKind:     string(item.ArtifactKind),
				SchemaID:         item.ArtifactSchema.ID,
				SchemaName:       item.ArtifactSchema.Name,
				SchemaVersion:    item.ArtifactSchema.Version,
				SchemaDigest:     contracts.DigestArtifactSchema(item.ArtifactSchema),
				MediaType:        item.MediaType,
				ContentDigest:    item.Digest,
			}
			encoded, err := json.Marshal(identity)
			if err != nil {
				continue
			}
			sum := sha256.Sum256(encoded)
			fingerprints = append(fingerprints, CheckpointFingerprint{
				Name:  fmt.Sprintf("generated-reference:%s:%d", slotName, index),
				Value: "sha256:" + hex.EncodeToString(sum[:]),
			})
		}
	}
	return normalizeCheckpointFingerprints(fingerprints)
}