package pipeline

import (
	"fmt"
	"strings"

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

func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []source.Chunk) ([]source.Chunk, error) {
	if len(chunks) == 0 {
		return nil, fmt.Errorf("chunks must not be empty")
	}

	sourceUnitIndexes := make(map[int]int, len(doc.Units))
	sourceUnits := make(map[int]source.SourceUnit, len(doc.Units))
	for index, unit := range doc.Units {
		sourceUnitIndexes[unit.ID] = index
		sourceUnits[unit.ID] = unit
	}

	canonicalChunks := make([]source.Chunk, 0, len(chunks))
	seenChunkIDs := make(map[string]struct{}, len(chunks))
	for chunkIndex, chunk := range chunks {
		if strings.TrimSpace(chunk.ID) == "" {
			return nil, fmt.Errorf("chunk[%d].id must not be empty", chunkIndex)
		}
		if _, ok := seenChunkIDs[chunk.ID]; ok {
			return nil, fmt.Errorf("chunk id %q is duplicated", chunk.ID)
		}
		seenChunkIDs[chunk.ID] = struct{}{}

		if chunk.SourceID != doc.ID {
			return nil, fmt.Errorf("chunk %q source_id %q does not match source document id %q", chunk.ID, chunk.SourceID, doc.ID)
		}
		if chunk.Index != chunkIndex {
			return nil, fmt.Errorf("chunk %q index %d does not match returned order %d", chunk.ID, chunk.Index, chunkIndex)
		}
		if len(chunk.Units) == 0 {
			return nil, fmt.Errorf("chunk %q units must not be empty", chunk.ID)
		}
		if len(chunk.Content) == 0 {
			return nil, fmt.Errorf("chunk %q content must not be empty", chunk.ID)
		}
		if strings.TrimSpace(chunk.MediaType) == "" {
			return nil, fmt.Errorf("chunk %q media_type must not be empty", chunk.ID)
		}
		if err := source.ValidateRef(doc, chunk.Ref); err != nil {
			return nil, fmt.Errorf("chunk %q ref: %w", chunk.ID, err)
		}

		seenUnitIDs := make(map[int]struct{}, len(chunk.Units))
		previousSourceIndex := -1
		canonicalUnits := make([]source.SourceUnit, 0, len(chunk.Units))
		for unitIndex, unit := range chunk.Units {
			if unit.ID <= 0 {
				return nil, fmt.Errorf("chunk %q unit[%d].id must be positive", chunk.ID, unitIndex)
			}
			if _, ok := seenUnitIDs[unit.ID]; ok {
				return nil, fmt.Errorf("chunk %q repeats source unit %d", chunk.ID, unit.ID)
			}
			seenUnitIDs[unit.ID] = struct{}{}

			sourceIndex, ok := sourceUnitIndexes[unit.ID]
			if !ok {
				return nil, fmt.Errorf("chunk %q source unit %d was not found in source document %q", chunk.ID, unit.ID, doc.ID)
			}
			if previousSourceIndex >= 0 && sourceIndex != previousSourceIndex+1 {
				return nil, fmt.Errorf("chunk %q source units must form a contiguous range in source document order", chunk.ID)
			}
			if unit.Ref != sourceUnits[unit.ID].Ref {
				return nil, fmt.Errorf("chunk %q source unit %d ref does not match source document", chunk.ID, unit.ID)
			}
			previousSourceIndex = sourceIndex
			canonicalUnits = append(canonicalUnits, cloneSourceUnit(sourceUnits[unit.ID]))
		}
		expectedRef := source.SourceRef{
			SourceID:    doc.ID,
			StartUnitID: canonicalUnits[0].Ref.StartUnitID,
			EndUnitID:   canonicalUnits[len(canonicalUnits)-1].Ref.EndUnitID,
		}
		if chunk.Ref != expectedRef {
			return nil, fmt.Errorf("chunk %q ref %#v does not match unit span %#v", chunk.ID, chunk.Ref, expectedRef)
		}

		canonicalChunks = append(canonicalChunks, source.Chunk{
			ID:        chunk.ID,
			SourceID:  chunk.SourceID,
			Index:     chunk.Index,
			Ref:       expectedRef,
			Content:   append([]byte(nil), chunk.Content...),
			MediaType: chunk.MediaType,
			Units:     canonicalUnits,
			Metadata:  cloneMetadata(chunk.Metadata),
		})
	}

	return canonicalChunks, nil
}

func cloneSourceUnit(unit source.SourceUnit) source.SourceUnit {
	return source.SourceUnit{
		ID:       unit.ID,
		Kind:     unit.Kind,
		Text:     unit.Text,
		Ref:      unit.Ref,
		Metadata: cloneMetadata(unit.Metadata),
	}
}