package semanticreconcile

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

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

const (
	candidateInputName  = "candidates"
	transcriptInputName = "transcript"
	jsonMediaType       = "application/json"
)

var defaultLimits = Limits{
	ContextRadius:        2,
	MaximumCandidates:    128,
	MaximumMaterialBytes: 262144,
}

// Candidate is contextual source-backed input supplied by a typed consumer.
// Prepare does not retain or mutate Label or SourceRefs.
type Candidate struct {
	Label      string
	SourceRefs []source.SourceRef
}

// Limits bounds source context and serialized model input.
type Limits struct {
	ContextRadius        int
	MaximumCandidates    int
	MaximumMaterialBytes int
}

// DefaultLimits returns the core-owned production limits.
func DefaultLimits() Limits {
	return defaultLimits
}

// Validate rejects limits that cannot safely bound preparation.
func (limits Limits) Validate() error {
	if limits.ContextRadius < 0 {
		return fmt.Errorf("semantic reconciliation limits: context radius must not be negative")
	}
	if limits.MaximumCandidates <= 0 {
		return fmt.Errorf("semantic reconciliation limits: maximum candidates must be positive")
	}
	if limits.MaximumMaterialBytes <= 0 {
		return fmt.Errorf("semantic reconciliation limits: maximum bytes must be positive")
	}
	return nil
}

// Disposition describes whether prepared materials may be sent to a model.
type Disposition uint8

const (
	// Ready indicates that the result contains complete bounded materials.
	Ready Disposition = iota + 1
	// InsufficientCandidates indicates that fewer than two candidates were
	// eligible after source-reference validation.
	InsufficientCandidates
	// LimitExceeded indicates that a candidate or serialized-material bound was
	// exceeded and no request should be split or sent.
	LimitExceeded
)

// CandidateMapping relates one model-visible request-local ID to the
// corresponding zero-based position in the caller's candidate slice.
type CandidateMapping struct {
	CandidateID       int
	CandidatePosition int
}

// Preparation owns the visible candidate mapping and prompt materials.
type Preparation struct {
	disposition Disposition
	mappings    []CandidateMapping
	materials   contracts.LLMInputSet
}

// Disposition returns the preparation outcome.
func (preparation Preparation) Disposition() Disposition {
	return preparation.disposition
}

// CandidateMappings returns an owned copy in model-visible candidate order.
func (preparation Preparation) CandidateMappings() []CandidateMapping {
	return append([]CandidateMapping(nil), preparation.mappings...)
}

// Materials returns independently owned candidate and transcript materials.
// It is empty unless Disposition returns Ready.
func (preparation Preparation) Materials() contracts.LLMInputSet {
	return preparation.materials.Clone()
}

type sourceRange struct {
	StartUnitID int `json:"start_unit_id"`
	EndUnitID   int `json:"end_unit_id"`
}

type visibleCandidate struct {
	CandidateID int           `json:"candidate_id"`
	Label       string        `json:"label"`
	SourceRefs  []sourceRange `json:"source_refs"`
}

type candidateInput struct {
	Candidates []visibleCandidate `json:"candidates"`
}

type transcriptInput struct {
	Windows []transcriptWindow `json:"windows"`
}

type transcriptWindow struct {
	Units []transcriptUnit `json:"units"`
}

type transcriptUnit struct {
	ID       int            `json:"id"`
	Kind     string         `json:"kind"`
	Text     string         `json:"text"`
	Metadata map[string]any `json:"metadata,omitempty"`
	Cited    bool           `json:"cited"`
}

type sourceInterval struct {
	start int
	end   int
}

type preparedCandidate struct {
	position   int
	references []sourceRange
	intervals  []sourceInterval
}

// Prepare validates candidates and constructs bounded, source-ordered model
// inputs. Deterministic skip conditions are represented by the returned
// disposition rather than an error.
func Prepare(document *source.SourceDocument, candidates []Candidate, limits Limits) (Preparation, error) {
	if err := limits.Validate(); err != nil {
		return Preparation{}, err
	}

	documentIndex := source.NewDocumentIndex(document)
	prepared := make([]preparedCandidate, 0, len(candidates))
	for candidatePosition, candidate := range candidates {
		references, intervals, valid := prepareReferences(documentIndex, candidate.SourceRefs)
		if !valid {
			continue
		}
		prepared = append(prepared, preparedCandidate{
			position:   candidatePosition,
			references: references,
			intervals:  intervals,
		})
	}

	result := Preparation{
		disposition: InsufficientCandidates,
		mappings:    make([]CandidateMapping, len(prepared)),
	}
	views := make([]visibleCandidate, len(prepared))
	for index, candidate := range prepared {
		candidateID := index + 1
		result.mappings[index] = CandidateMapping{
			CandidateID:       candidateID,
			CandidatePosition: candidate.position,
		}
		views[index] = visibleCandidate{
			CandidateID: candidateID,
			Label:       candidates[candidate.position].Label,
			SourceRefs:  cloneSourceRanges(candidate.references),
		}
	}
	if len(prepared) < 2 {
		return result, nil
	}
	if len(prepared) > limits.MaximumCandidates {
		result.disposition = LimitExceeded
		return result, nil
	}

	candidateContent, withinLimit, err := marshalCandidateInput(views, limits.MaximumMaterialBytes)
	if err != nil {
		return Preparation{}, fmt.Errorf("prepare semantic reconciliation: encode candidate material: %w", err)
	}
	if !withinLimit {
		result.disposition = LimitExceeded
		return result, nil
	}

	contextIntervals := make([]sourceInterval, 0)
	citedIntervals := make([]sourceInterval, 0)
	for _, candidate := range prepared {
		for _, interval := range candidate.intervals {
			citedIntervals = append(citedIntervals, interval)
			contextIntervals = append(contextIntervals, sourceInterval{
				start: max(0, interval.start-limits.ContextRadius),
				end:   min(len(document.Units)-1, interval.end+limits.ContextRadius),
			})
		}
	}
	transcriptContent, withinLimit, err := marshalTranscriptInput(
		document.Units,
		coalesceIntervals(contextIntervals),
		coalesceIntervals(citedIntervals),
		limits.MaximumMaterialBytes-len(candidateContent),
	)
	if err != nil {
		return Preparation{}, fmt.Errorf("prepare semantic reconciliation: build transcript material: invalid source metadata")
	}
	if !withinLimit {
		result.disposition = LimitExceeded
		return result, nil
	}

	result.disposition = Ready
	result.materials = contracts.LLMInputSet{
		candidateInputName:  newInputMaterial(candidateInputName, candidateContent),
		transcriptInputName: newInputMaterial(transcriptInputName, transcriptContent),
	}
	return result, nil
}

func prepareReferences(index source.DocumentIndex, references []source.SourceRef) ([]sourceRange, []sourceInterval, bool) {
	if len(references) == 0 {
		return nil, nil, false
	}
	type referencedInterval struct {
		reference sourceRange
		interval  sourceInterval
	}
	prepared := make([]referencedInterval, 0, len(references))
	for _, reference := range references {
		if err := index.ValidateRef(reference); err != nil {
			return nil, nil, false
		}
		start, _ := index.Position(reference.StartUnitID)
		end, _ := index.Position(reference.EndUnitID)
		prepared = append(prepared, referencedInterval{
			reference: sourceRange{StartUnitID: reference.StartUnitID, EndUnitID: reference.EndUnitID},
			interval:  sourceInterval{start: start, end: end},
		})
	}
	sort.Slice(prepared, func(left, right int) bool {
		if prepared[left].interval.start != prepared[right].interval.start {
			return prepared[left].interval.start < prepared[right].interval.start
		}
		return prepared[left].interval.end < prepared[right].interval.end
	})

	canonicalReferences := make([]sourceRange, 0, len(prepared))
	intervals := make([]sourceInterval, 0, len(prepared))
	for _, item := range prepared {
		if len(canonicalReferences) > 0 && canonicalReferences[len(canonicalReferences)-1] == item.reference {
			continue
		}
		canonicalReferences = append(canonicalReferences, item.reference)
		intervals = append(intervals, item.interval)
	}
	return canonicalReferences, intervals, true
}

func cloneSourceRanges(ranges []sourceRange) []sourceRange {
	if len(ranges) == 0 {
		return []sourceRange{}
	}
	return append([]sourceRange(nil), ranges...)
}

func coalesceIntervals(intervals []sourceInterval) []sourceInterval {
	if len(intervals) == 0 {
		return nil
	}
	ordered := append([]sourceInterval(nil), intervals...)
	sort.Slice(ordered, func(left, right int) bool {
		if ordered[left].start != ordered[right].start {
			return ordered[left].start < ordered[right].start
		}
		return ordered[left].end < ordered[right].end
	})

	coalesced := make([]sourceInterval, 0, len(ordered))
	for _, interval := range ordered {
		if len(coalesced) == 0 || interval.start > coalesced[len(coalesced)-1].end+1 {
			coalesced = append(coalesced, interval)
			continue
		}
		if interval.end > coalesced[len(coalesced)-1].end {
			coalesced[len(coalesced)-1].end = interval.end
		}
	}
	return coalesced
}

func marshalCandidateInput(candidates []visibleCandidate, maximumBytes int) ([]byte, bool, error) {
	content := make([]byte, 0, min(maximumBytes, 4096))
	var withinLimit bool
	content, withinLimit = appendWithinLimit(content, maximumBytes, []byte(`{"candidates":[`))
	if !withinLimit {
		return nil, false, nil
	}
	for index, candidate := range candidates {
		encoded, err := json.Marshal(candidate)
		if err != nil {
			return nil, false, err
		}
		separator := []byte(nil)
		if index > 0 {
			separator = []byte(",")
		}
		content, withinLimit = appendWithinLimit(content, maximumBytes, separator, encoded)
		if !withinLimit {
			return nil, false, nil
		}
	}
	content, withinLimit = appendWithinLimit(content, maximumBytes, []byte("]}"))
	return content, withinLimit, nil
}

// marshalTranscriptInput retains at most maximumBytes while visiting source
// units in order. It deliberately serializes one unit at a time so an oversized
// request does not require a document-sized transcript copy before rejection.
func marshalTranscriptInput(units []source.SourceUnit, contextIntervals, citedIntervals []sourceInterval, maximumBytes int) ([]byte, bool, error) {
	content := make([]byte, 0, min(maximumBytes, 4096))
	content, withinLimit := appendWithinLimit(content, maximumBytes, []byte(`{"windows":[`))
	if !withinLimit {
		return nil, false, nil
	}

	citedIndex := 0
	for windowIndex, interval := range contextIntervals {
		separator := []byte(nil)
		if windowIndex > 0 {
			separator = []byte(",")
		}
		content, withinLimit = appendWithinLimit(content, maximumBytes, separator, []byte(`{"units":[`))
		if !withinLimit {
			return nil, false, nil
		}
		for position := interval.start; position <= interval.end; position++ {
			for citedIndex < len(citedIntervals) && citedIntervals[citedIndex].end < position {
				citedIndex++
			}
			cited := citedIndex < len(citedIntervals) && citedIntervals[citedIndex].start <= position
			unit := units[position]
			metadata, err := source.CloneMetadata(unit.Metadata)
			if err != nil {
				return nil, false, err
			}
			encoded, err := json.Marshal(transcriptUnit{
				ID: unit.ID, Kind: unit.Kind, Text: unit.Text, Metadata: metadata, Cited: cited,
			})
			if err != nil {
				return nil, false, err
			}
			separator = nil
			if position > interval.start {
				separator = []byte(",")
			}
			content, withinLimit = appendWithinLimit(content, maximumBytes, separator, encoded)
			if !withinLimit {
				return nil, false, nil
			}
		}
		content, withinLimit = appendWithinLimit(content, maximumBytes, []byte("]}"))
		if !withinLimit {
			return nil, false, nil
		}
	}
	content, withinLimit = appendWithinLimit(content, maximumBytes, []byte("]}"))
	return content, withinLimit, nil
}

func appendWithinLimit(content []byte, maximumBytes int, parts ...[]byte) ([]byte, bool) {
	for _, part := range parts {
		if len(content) > maximumBytes || len(part) > maximumBytes-len(content) {
			return content, false
		}
		content = append(content, part...)
	}
	return content, true
}

func newInputMaterial(name string, content []byte) contracts.LLMInputMaterial {
	digest := sha256.Sum256(content)
	return contracts.NewLLMInputMaterial(
		name,
		jsonMediaType,
		content,
		"sha256:"+hex.EncodeToString(digest[:]),
		"",
	)
}