package checkpoint

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"path"
	"sort"
	"strings"
	"time"

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

type FilesystemRecorder struct {
	root           string
	identityDigest string
	now            func() time.Time
}

func NewFilesystemRecorder(root string, identity Identity) (pipeline.CheckpointRecorder, error) {
	root = strings.TrimSpace(root)
	if root == "" {
		return pipeline.NoopCheckpointRecorder(), nil
	}
	relative, err := identity.RelativePath()
	if err != nil {
		return nil, err
	}
	target, err := fileio.SafePath(root, relative)
	if err != nil {
		return nil, err
	}
	return &FilesystemRecorder{root: target, identityDigest: identity.Digest, now: time.Now}, nil
}

func (r *FilesystemRecorder) SourceRunning(moduleKey string) error {
	manifest := r.newStageManifest(StageSource, StatusRunning)
	manifest.ModuleKey = moduleKey
	manifest.StartedAt = timePtr(r.timestamp())
	return r.writeManifest("source/manifest.json", SourceManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) SourceSucceeded(moduleKey string, doc *source.SourceDocument) error {
	if doc == nil {
		return fmt.Errorf("checkpoint source document must not be nil")
	}
	if err := r.writePayload("source/source-document.json", sourceDocumentEnvelope{Document: cloneSourceDocument(*doc)}); err != nil {
		return err
	}
	manifest := r.newStageManifest(StageSource, StatusSucceeded)
	manifest.ModuleKey = moduleKey
	manifest.OutputDigests = checkpointFingerprints(digestFingerprints("source_document", doc.Digest))
	manifest.CompletedAt = timePtr(r.timestamp())
	return r.writeManifest("source/manifest.json", SourceManifest{
		StageManifest: manifest,
		SourceID:      doc.ID,
	})
}

func (r *FilesystemRecorder) SourceFailed(moduleKey string, err error) error {
	manifest := r.newStageManifest(StageSource, StatusFailed)
	manifest.ModuleKey = moduleKey
	manifest.CompletedAt = timePtr(r.timestamp())
	manifest.Metadata = errorMetadata(err)
	return r.writeManifest("source/manifest.json", SourceManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) ExtractRunning(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
	return r.ExtractRunningForStep("", laneID, moduleKey, dependencies)
}

func (r *FilesystemRecorder) ExtractRunningForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
	manifest := r.laneManifest(StageExtract, StatusRunning, stepID, laneID, moduleKey, dependencies)
	manifest.StartedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("extract", stepID, laneID), ExtractLaneManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) ExtractSucceeded(laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint, outputs []pipeline.CheckpointArtifact, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error {
	return r.ExtractSucceededForStep("", laneID, moduleKey, dependencies, outputs, rejected, warnings)
}

func (r *FilesystemRecorder) ExtractSucceededForStep(stepID, laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint, outputs []pipeline.CheckpointArtifact, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error {
	payload := artifactExtractEnvelope{Outputs: artifactCheckpointEnvelopes(outputs), Rejected: cloneRejectedOutputs(rejected), Warnings: cloneWarnings(warnings)}
	if err := r.writePayload(lanePayloadPath("extract", stepID, laneID, "outputs.json"), payload); err != nil {
		return err
	}
	manifest := r.laneManifest(StageExtract, statusForRejected(rejected), stepID, laneID, moduleKey, dependencies)
	manifest.OutputDigests = checkpointFingerprints(artifactOutputDigests(outputs))
	manifest.ValidationStatus = validationStatusString(warnings, rejected)
	manifest.Rejections = rejectionSummaries(rejected)
	manifest.CompletedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("extract", stepID, laneID), ExtractLaneManifest{StageManifest: manifest, ChunkCount: len(outputs) + len(rejected), OutputCount: len(outputs)})
}

func (r *FilesystemRecorder) ExtractFailed(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, err error) error {
	return r.ExtractFailedForStep("", laneID, moduleKey, dependencies, err)
}

func (r *FilesystemRecorder) ExtractFailedForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, err error) error {
	manifest := r.laneManifest(StageExtract, StatusFailed, stepID, laneID, moduleKey, dependencies)
	manifest.CompletedAt = timePtr(r.timestamp())
	manifest.Metadata = errorMetadata(err)
	return r.writeManifest(laneManifestPath("extract", stepID, laneID), ExtractLaneManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) MergeRunning(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
	return r.MergeRunningForStep("", laneID, moduleKey, dependencies)
}

func (r *FilesystemRecorder) MergeRunningForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
	manifest := r.laneManifest(StageMerge, StatusRunning, stepID, laneID, moduleKey, dependencies)
	manifest.StartedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("merge", stepID, laneID), MergeLaneManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) MergeSucceeded(laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint, output pipeline.CheckpointArtifact, warnings []contracts.Warning) error {
	return r.MergeSucceededForStep("", laneID, moduleKey, dependencies, output, warnings)
}

func (r *FilesystemRecorder) MergeSucceededForStep(stepID, laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint, output pipeline.CheckpointArtifact, warnings []contracts.Warning) error {
	if err := r.writePayload(lanePayloadPath("merge", stepID, laneID, "output.json"), artifactSingleEnvelope{Output: artifactCheckpointEnvelopeFromOutput(output), Warnings: cloneWarnings(warnings)}); err != nil {
		return err
	}
	manifest := r.laneManifest(StageMerge, StatusSucceeded, stepID, laneID, moduleKey, dependencies)
	manifest.OutputDigests = checkpointFingerprints(artifactOutputDigests([]pipeline.CheckpointArtifact{output}))
	manifest.ValidationStatus = validationStatusString(warnings, nil)
	manifest.CompletedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("merge", stepID, laneID), MergeLaneManifest{StageManifest: manifest, InputCount: len(dependencies)})
}

func (r *FilesystemRecorder) MergeRejected(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, rejected contracts.RejectedOutput) error {
	return r.MergeRejectedForStep("", laneID, moduleKey, dependencies, rejected)
}

func (r *FilesystemRecorder) MergeRejectedForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, rejected contracts.RejectedOutput) error {
	manifest := r.laneManifest(StageMerge, StatusSucceededWithRejections, stepID, laneID, moduleKey, dependencies)
	manifest.ValidationStatus = "rejected"
	manifest.Rejections = rejectionSummaries([]contracts.RejectedOutput{rejected})
	manifest.CompletedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("merge", stepID, laneID), MergeLaneManifest{StageManifest: manifest, InputCount: len(dependencies)})
}

func (r *FilesystemRecorder) MergeFailed(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, err error) error {
	return r.MergeFailedForStep("", laneID, moduleKey, dependencies, err)
}

func (r *FilesystemRecorder) MergeFailedForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, err error) error {
	manifest := r.laneManifest(StageMerge, StatusFailed, stepID, laneID, moduleKey, dependencies)
	manifest.CompletedAt = timePtr(r.timestamp())
	manifest.Metadata = errorMetadata(err)
	return r.writeManifest(laneManifestPath("merge", stepID, laneID), MergeLaneManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) NormalizeRunning(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
	return r.NormalizeRunningForStep("", laneID, moduleKey, dependencies)
}

func (r *FilesystemRecorder) NormalizeRunningForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) error {
	manifest := r.laneManifest(StageNormalize, StatusRunning, stepID, laneID, moduleKey, dependencies)
	manifest.StartedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("normalize", stepID, laneID), NormalizeLaneManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) NormalizeSucceeded(laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint, output pipeline.CheckpointArtifact, warnings []contracts.Warning) error {
	return r.NormalizeSucceededForStep("", laneID, moduleKey, dependencies, output, warnings)
}

func (r *FilesystemRecorder) NormalizeSucceededForStep(stepID, laneID, moduleKey string, dependencies []pipeline.CheckpointFingerprint, output pipeline.CheckpointArtifact, warnings []contracts.Warning) error {
	if err := r.writePayload(lanePayloadPath("normalize", stepID, laneID, "output.json"), artifactSingleEnvelope{Output: artifactCheckpointEnvelopeFromOutput(output), Warnings: cloneWarnings(warnings)}); err != nil {
		return err
	}
	manifest := r.laneManifest(StageNormalize, StatusSucceeded, stepID, laneID, moduleKey, dependencies)
	manifest.OutputDigests = checkpointFingerprints(artifactOutputDigests([]pipeline.CheckpointArtifact{output}))
	manifest.ValidationStatus = validationStatusString(warnings, nil)
	manifest.CompletedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("normalize", stepID, laneID), NormalizeLaneManifest{StageManifest: manifest, InputCount: len(dependencies)})
}

func (r *FilesystemRecorder) NormalizeRejected(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, rejected contracts.RejectedOutput) error {
	return r.NormalizeRejectedForStep("", laneID, moduleKey, dependencies, rejected)
}

func (r *FilesystemRecorder) NormalizeRejectedForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, rejected contracts.RejectedOutput) error {
	manifest := r.laneManifest(StageNormalize, StatusSucceededWithRejections, stepID, laneID, moduleKey, dependencies)
	manifest.ValidationStatus = "rejected"
	manifest.Rejections = rejectionSummaries([]contracts.RejectedOutput{rejected})
	manifest.CompletedAt = timePtr(r.timestamp())
	return r.writeManifest(laneManifestPath("normalize", stepID, laneID), NormalizeLaneManifest{StageManifest: manifest, InputCount: len(dependencies)})
}

func (r *FilesystemRecorder) NormalizeFailed(laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, err error) error {
	return r.NormalizeFailedForStep("", laneID, moduleKey, dependencies, err)
}

func (r *FilesystemRecorder) NormalizeFailedForStep(stepID, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint, err error) error {
	manifest := r.laneManifest(StageNormalize, StatusFailed, stepID, laneID, moduleKey, dependencies)
	manifest.CompletedAt = timePtr(r.timestamp())
	manifest.Metadata = errorMetadata(err)
	return r.writeManifest(laneManifestPath("normalize", stepID, laneID), NormalizeLaneManifest{StageManifest: manifest})
}

func (r *FilesystemRecorder) writeManifest(name string, payload any) error {
	return r.writeJSON(name, payload)
}

func (r *FilesystemRecorder) writePayload(name string, payload any) error {
	return r.writeJSON(name, payload)
}

func (r *FilesystemRecorder) writeJSON(name string, payload any) error {
	if r == nil || strings.TrimSpace(r.root) == "" {
		return nil
	}
	return fileio.WriteJSON(r.root, name, payload, 0o700, 0o600)
}

func (r *FilesystemRecorder) timestamp() time.Time {
	if r == nil || r.now == nil {
		return time.Now().UTC()
	}
	return r.now().UTC()
}

func (r *FilesystemRecorder) newStageManifest(stage StageName, status StageStatus) StageManifest {
	manifest := NewStageManifest(stage, status)
	if strings.TrimSpace(r.identityDigest) != "" {
		manifest.Metadata = map[string]string{"checkpoint_identity_digest": r.identityDigest}
	}
	return manifest
}

func (r *FilesystemRecorder) laneManifest(stage StageName, status StageStatus, stepID string, laneID string, moduleKey string, dependencies []pipeline.CheckpointFingerprint) StageManifest {
	manifest := r.newStageManifest(stage, status)
	manifest.StepID = stepID
	manifest.LaneID = laneID
	manifest.ModuleKey = moduleKey
	manifest.DependencyFingerprints = checkpointFingerprints(dependencies)
	return manifest
}

type sourceDocumentEnvelope struct {
	Document source.SourceDocument `json:"document"`
}

type binaryEnvelope struct {
	ContentBase64 string              `json:"content_base64,omitempty"`
	ContentDigest string              `json:"content_digest,omitempty"`
	MediaType     string              `json:"media_type,omitempty"`
	Metadata      map[string]any      `json:"metadata,omitempty"`
	Warnings      []contracts.Warning `json:"warnings,omitempty"`
}

type artifactCheckpointEnvelope struct {
	LaneID       string                   `json:"lane_id"`
	ModuleKey    string                   `json:"module_key"`
	SourceID     string                   `json:"source_id,omitempty"`
	ChunkID      string                   `json:"chunk_id,omitempty"`
	ChunkIndex   int                      `json:"chunk_index,omitempty"`
	ChunkRef     source.SourceRef         `json:"chunk_ref,omitempty"`
	Kind         contracts.ArtifactKind   `json:"artifact_kind"`
	Schema       contracts.ArtifactSchema `json:"schema"`
	SchemaDigest string                   `json:"schema_digest"`
	Content      binaryEnvelope           `json:"content"`
}
type artifactExtractEnvelope struct {
	Outputs  []artifactCheckpointEnvelope `json:"outputs"`
	Rejected []contracts.RejectedOutput   `json:"rejected,omitempty"`
	Warnings []contracts.Warning          `json:"warnings,omitempty"`
}
type artifactSingleEnvelope struct {
	Output   artifactCheckpointEnvelope `json:"output"`
	Warnings []contracts.Warning        `json:"warnings,omitempty"`
}

func artifactCheckpointEnvelopeFromOutput(output pipeline.CheckpointArtifact) artifactCheckpointEnvelope {
	schema := contracts.CloneArtifactSchema(output.Artifact.Schema)
	schema.JSONSchema = nil
	return artifactCheckpointEnvelope{LaneID: output.LaneID, ModuleKey: output.ModuleKey, SourceID: output.SourceID, ChunkID: output.ChunkID, ChunkIndex: output.ChunkIndex, ChunkRef: output.ChunkRef, Kind: output.Artifact.Kind, Schema: schema, SchemaDigest: output.SchemaDigest, Content: binaryEnvelopeFromContent(output.Artifact.Content, output.Artifact.MediaType, output.Artifact.Metadata, nil)}
}
func artifactCheckpointEnvelopes(outputs []pipeline.CheckpointArtifact) []artifactCheckpointEnvelope {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]artifactCheckpointEnvelope, 0, len(outputs))
	for _, v := range outputs {
		out = append(out, artifactCheckpointEnvelopeFromOutput(v))
	}
	return out
}
func artifactOutputDigests(outputs []pipeline.CheckpointArtifact) []pipeline.CheckpointFingerprint {
	values := make([]pipeline.CheckpointFingerprint, 0, len(outputs))
	for i, v := range outputs {
		values = append(values, pipeline.CheckpointFingerprint{Name: fmt.Sprintf("artifact[%d]", i), Value: contentDigest(v.Artifact.Content)})
	}
	return normalizeFingerprints(values)
}

func binaryEnvelopeFromContent(content []byte, mediaType string, metadata map[string]any, warnings []contracts.Warning) binaryEnvelope {
	return binaryEnvelope{
		ContentBase64: base64.StdEncoding.EncodeToString(content),
		ContentDigest: contentDigest(content),
		MediaType:     mediaType,
		Metadata:      cloneMetadata(metadata),
		Warnings:      cloneWarnings(warnings),
	}
}

func cloneSourceDocument(doc source.SourceDocument) source.SourceDocument {
	doc.Units = cloneSourceUnits(doc.Units)
	doc.Metadata = cloneMetadata(doc.Metadata)
	return doc
}

func cloneSourceUnits(units []source.SourceUnit) []source.SourceUnit {
	if len(units) == 0 {
		return nil
	}
	out := make([]source.SourceUnit, 0, len(units))
	for _, unit := range units {
		out = append(out, source.SourceUnit{
			ID:       unit.ID,
			Kind:     unit.Kind,
			Text:     unit.Text,
			Ref:      unit.Ref,
			Metadata: cloneMetadata(unit.Metadata),
		})
	}
	return out
}

func cloneWarnings(warnings []contracts.Warning) []contracts.Warning {
	if len(warnings) == 0 {
		return nil
	}
	return append([]contracts.Warning(nil), warnings...)
}

func cloneRejectedOutputs(rejected []contracts.RejectedOutput) []contracts.RejectedOutput {
	if len(rejected) == 0 {
		return nil
	}
	return append([]contracts.RejectedOutput(nil), rejected...)
}

func cloneMetadata(metadata map[string]any) map[string]any {
	if len(metadata) == 0 {
		return nil
	}
	out := make(map[string]any, len(metadata))
	for key, value := range metadata {
		out[key] = value
	}
	return out
}

func digestFingerprints(name string, digest string) []pipeline.CheckpointFingerprint {
	digest = strings.TrimSpace(digest)
	if digest == "" {
		return nil
	}
	return []pipeline.CheckpointFingerprint{{Name: name, Value: digest}}
}

func checkpointFingerprints(values []pipeline.CheckpointFingerprint) []Fingerprint {
	normalized := normalizeFingerprints(values)
	if len(normalized) == 0 {
		return nil
	}
	out := make([]Fingerprint, 0, len(normalized))
	for _, value := range normalized {
		out = append(out, Fingerprint{Name: value.Name, Value: value.Value})
	}
	return out
}

func normalizeFingerprints(values []pipeline.CheckpointFingerprint) []pipeline.CheckpointFingerprint {
	if len(values) == 0 {
		return nil
	}
	byName := make(map[string]string, len(values))
	for _, value := range values {
		name := strings.TrimSpace(value.Name)
		fingerprint := strings.TrimSpace(value.Value)
		if name == "" || fingerprint == "" {
			continue
		}
		byName[name] = fingerprint
	}
	if len(byName) == 0 {
		return nil
	}
	names := make([]string, 0, len(byName))
	for name := range byName {
		names = append(names, name)
	}
	sort.Strings(names)
	out := make([]pipeline.CheckpointFingerprint, 0, len(names))
	for _, name := range names {
		out = append(out, pipeline.CheckpointFingerprint{Name: name, Value: byName[name]})
	}
	return out
}

func rejectionSummaries(rejected []contracts.RejectedOutput) []RejectionSummary {
	if len(rejected) == 0 {
		return nil
	}
	type key struct {
		validatorName string
		reasonCode    string
		message       string
	}
	counts := make(map[key]int, len(rejected))
	for _, item := range rejected {
		k := key{validatorName: item.ValidatorName, reasonCode: item.ReasonCode, message: item.Message}
		counts[k]++
	}
	keys := make([]key, 0, len(counts))
	for k := range counts {
		keys = append(keys, k)
	}
	sort.Slice(keys, func(i, j int) bool {
		if keys[i].validatorName != keys[j].validatorName {
			return keys[i].validatorName < keys[j].validatorName
		}
		if keys[i].reasonCode != keys[j].reasonCode {
			return keys[i].reasonCode < keys[j].reasonCode
		}
		return keys[i].message < keys[j].message
	})
	out := make([]RejectionSummary, 0, len(keys))
	for _, k := range keys {
		out = append(out, RejectionSummary{
			ValidatorName: k.validatorName,
			ReasonCode:    k.reasonCode,
			Message:       k.message,
			Count:         counts[k],
		})
	}
	return out
}

func statusForRejected(rejected []contracts.RejectedOutput) StageStatus {
	if len(rejected) > 0 {
		return StatusSucceededWithRejections
	}
	return StatusSucceeded
}

func validationStatusString(warnings []contracts.Warning, rejected []contracts.RejectedOutput) string {
	if len(rejected) > 0 {
		return "rejected"
	}
	if len(warnings) > 0 {
		return "approved_with_warnings"
	}
	return "approved"
}

func errorMetadata(err error) map[string]string {
	if err == nil {
		return nil
	}
	return map[string]string{"error": err.Error()}
}

func laneManifestPath(stage string, stepID string, laneID string) string {
	return lanePayloadPath(stage, stepID, laneID, "manifest.json")
}

func lanePayloadPath(stage string, stepID string, laneID string, file string) string {
	if strings.TrimSpace(stepID) == "" {
		return path.Join(stage, fileio.EncodePathComponent(laneID), file)
	}
	return path.Join(stage, fileio.EncodePathComponent(stepID), fileio.EncodePathComponent(laneID), file)
}

func contentDigest(content []byte) string {
	sum := sha256.Sum256(content)
	return "sha256:" + hex.EncodeToString(sum[:])
}

func timePtr(t time.Time) *time.Time {
	return &t
}