package stage

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/narratio/internal/adapters/storage"
	"gitea.maximumdirect.net/eric/narratio/internal/artifacts"
	"gitea.maximumdirect.net/eric/narratio/internal/audio"
	"gitea.maximumdirect.net/eric/narratio/internal/config"
	"gitea.maximumdirect.net/eric/narratio/internal/manifest"
	"gopkg.in/yaml.v3"
)

type prepareStage struct{}

func (prepareStage) Name() string { return "prepare" }

func (prepareStage) Declares() IODecl {
	return IODecl{
		Inputs: []artifacts.Ref{
			{Kind: "config", Category: "inputs", RelativePath: "campaign.yml"},
			{Kind: "config", Category: "inputs", RelativePath: "session.yml"},
			{Kind: "config", Category: "inputs", RelativePath: "pipeline.resolved.yml"},
			{Kind: "config", Category: "inputs", RelativePath: "speakers.yml"},
			{Kind: "config", Category: "inputs", RelativePath: "autocorrect.yml"},
			{Kind: "config", Category: "inputs", RelativePath: "glossary.yml"},
			{Kind: "audio", Category: "audio", RelativePath: "*.flac"},
		},
	}
}

func (prepareStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (*StageResult, error) {
	if env == nil || env.Config == nil {
		return nil, fmt.Errorf("prepare: stage environment config is required")
	}
	if env.ArtifactStore == nil {
		return nil, fmt.Errorf("prepare: artifact store is required")
	}
	if env.Config.Session == nil || env.Config.Pipeline == nil {
		return nil, fmt.Errorf("prepare: resolved config must include pipeline and session")
	}

	sessionID := strings.TrimSpace(m.SessionID)
	if sessionID == "" {
		sessionID = strings.TrimSpace(env.Config.Session.SessionID)
	}
	if sessionID == "" {
		return nil, fmt.Errorf("prepare: session id is required")
	}

	paths, err := ensureLayoutForEnv(env, sessionID)
	if err != nil {
		return nil, fmt.Errorf("prepare: ensure workdir layout: %w", err)
	}

	campaignSrc := env.Config.CampaignPath
	if err := requireFile(campaignSrc, "campaign.yml"); err != nil {
		return nil, fmt.Errorf("prepare: %w", err)
	}

	sessionSrc := env.Config.SessionPath
	if err := requireFile(sessionSrc, "session.yml"); err != nil {
		return nil, fmt.Errorf("prepare: %w", err)
	}
	sessionDir := filepath.Dir(sessionSrc)

	speakersInput := stableInputSource(env.Config.StableInputs.SpeakersFile, env.Config.Session.Inputs.SpeakersFile, sessionSrc)
	autocorrectInput := stableInputSource(env.Config.StableInputs.AutocorrectFile, env.Config.Session.Inputs.AutocorrectFile, sessionSrc)
	glossaryInput := stableInputSource(env.Config.StableInputs.GlossaryFile, env.Config.Session.Inputs.GlossaryFile, sessionSrc)

	speakersSrc, err := resolveConfigRelativePath(speakersInput)
	if err != nil {
		return nil, fmt.Errorf("prepare: speakers path: %w", err)
	}
	autocorrectSrc, err := resolveConfigRelativePath(autocorrectInput)
	if err != nil {
		return nil, fmt.Errorf("prepare: autocorrect path: %w", err)
	}
	glossarySrc, err := resolveConfigRelativePath(glossaryInput)
	if err != nil {
		return nil, fmt.Errorf("prepare: glossary path: %w", err)
	}

	for _, required := range []struct {
		path string
		name string
	}{
		{path: speakersSrc, name: "speakers.yml"},
		{path: autocorrectSrc, name: "autocorrect.yml"},
		{path: glossarySrc, name: "glossary.yml"},
	} {
		if err := requireFile(required.path, required.name); err != nil {
			return nil, fmt.Errorf("prepare: %w", err)
		}
	}

	resolvedLocalAudio, useS3Audio, err := resolveAudioInputs(sessionDir, env.Config.Session.Inputs)
	if err != nil {
		return nil, fmt.Errorf("prepare: resolve audio inputs: %w", err)
	}

	inputs := make([]manifest.InputRecord, 0, 6+len(resolvedLocalAudio))
	registerInput := func(kind, path, checksum string) {
		inputs = append(inputs, manifest.InputRecord{Kind: kind, Path: path, Checksum: checksum})
	}
	registerConfigInput := func(kind, path, checksum, source string) {
		inputs = append(inputs, manifest.InputRecord{Kind: kind, Path: path, Checksum: checksum, Source: source})
	}
	registerSessionConfigInput := func(path, checksum string) {
		source := env.Config.SessionSource
		if strings.TrimSpace(source.Source) == "" {
			source.Source = "session_config"
		}
		inputs = append(inputs, manifest.InputRecord{
			Kind:      "session_config",
			Path:      path,
			Checksum:  checksum,
			Source:    source.Source,
			S3Bucket:  source.S3Bucket,
			S3Key:     source.S3Key,
			S3Size:    source.S3Size,
			S3ETag:    source.S3ETag,
			SpoolPath: source.SpoolPath,
		})
	}

	campaignDst := filepath.Join(paths.InputsDir, "campaign.yml")
	campaignChecksum, err := copyFileIfChanged(env.ArtifactStore, campaignSrc, campaignDst)
	if err != nil {
		return nil, fmt.Errorf("prepare: materialize campaign.yml: %w", err)
	}
	registerConfigInput("campaign_config", campaignDst, campaignChecksum, "campaign_config")

	sessionDst := filepath.Join(paths.InputsDir, "session.yml")
	sessionChecksum, err := copyFileIfChanged(env.ArtifactStore, sessionSrc, sessionDst)
	if err != nil {
		return nil, fmt.Errorf("prepare: materialize session.yml: %w", err)
	}
	registerSessionConfigInput(sessionDst, sessionChecksum)

	pipelineResolvedBytes, err := renderResolvedPipeline(env.Config.Pipeline)
	if err != nil {
		return nil, fmt.Errorf("prepare: render pipeline.resolved.yml: %w", err)
	}
	pipelineDst := filepath.Join(paths.InputsDir, "pipeline.resolved.yml")
	pipelineChecksum, err := writeBytesIfChanged(env.ArtifactStore, pipelineDst, pipelineResolvedBytes)
	if err != nil {
		return nil, fmt.Errorf("prepare: materialize pipeline.resolved.yml: %w", err)
	}
	registerInput("pipeline_resolved", pipelineDst, pipelineChecksum)

	for _, cfgFile := range []struct {
		kind   string
		src    string
		dst    string
		source string
	}{
		{kind: "speakers", src: speakersSrc, dst: filepath.Join(paths.InputsDir, "speakers.yml"), source: speakersInput.Source},
		{kind: "autocorrect", src: autocorrectSrc, dst: filepath.Join(paths.InputsDir, "autocorrect.yml"), source: autocorrectInput.Source},
		{kind: "glossary", src: glossarySrc, dst: filepath.Join(paths.InputsDir, "glossary.yml"), source: glossaryInput.Source},
	} {
		checksum, err := copyFileIfChanged(env.ArtifactStore, cfgFile.src, cfgFile.dst)
		if err != nil {
			return nil, fmt.Errorf("prepare: materialize %s: %w", cfgFile.kind, err)
		}
		registerConfigInput(cfgFile.kind, cfgFile.dst, checksum, cfgFile.source)
	}

	var audioCacheStats s3AudioMaterializationStats
	if useS3Audio {
		stats, err := materializeS3AudioInputs(ctx, env, m, sessionID, &inputs)
		if err != nil {
			return nil, fmt.Errorf("prepare: materialize s3 audio: %w", err)
		}
		audioCacheStats = stats
	} else {
		if err := materializeLocalAudioInputs(env, paths, resolvedLocalAudio, registerInput); err != nil {
			return nil, fmt.Errorf("prepare: %w", err)
		}
	}

	previousRequirements := collectPreparePreviousRequirements(env.Config)
	var previousHydration *previousSessionHydrationResult
	if len(previousRequirements) > 0 {
		if err := clearManagedPreviousState(paths); err != nil {
			return nil, fmt.Errorf("prepare: clear previous-session cache: %w", err)
		}
		hydration, err := hydratePreviousSessionArtifacts(ctx, env, paths, previousRequirements)
		if err != nil {
			return nil, fmt.Errorf("prepare: hydrate previous-session artifacts: %w", err)
		}
		previousHydration = hydration
		inputs = append(inputs, hydration.Inputs...)
	}

	sort.Slice(inputs, func(i, j int) bool {
		if inputs[i].Kind != inputs[j].Kind {
			return inputs[i].Kind < inputs[j].Kind
		}
		return inputs[i].Path < inputs[j].Path
	})
	m.Inputs = inputs

	metadata := map[string]any{
		"prepared":             true,
		"stage":                "prepare",
		"inputs_count":         len(inputs),
		"audio_files_resolved": countAudioInputs(inputs),
	}
	if useS3Audio {
		metadata["audio_cache_hits"] = audioCacheStats.CacheHits
		metadata["audio_cache_misses"] = audioCacheStats.CacheMisses
		metadata["audio_s3_downloads"] = audioCacheStats.Downloads
	}
	if len(previousRequirements) > 0 {
		metadata["previous_requirements_count"] = len(previousRequirements)
		if previousHydration != nil {
			metadata["previous_artifacts_hydrated"] = append([]string(nil), previousHydration.Hydrated...)
			metadata["previous_artifacts_hydrated_count"] = len(previousHydration.Hydrated)
			metadata["previous_artifacts_missing_optional"] = append([]string(nil), previousHydration.SkippedMissing...)
			metadata["previous_artifacts_missing_optional_count"] = len(previousHydration.SkippedMissing)
			if strings.TrimSpace(previousHydration.PreviousRunID) != "" {
				metadata["previous_session_run_id"] = strings.TrimSpace(previousHydration.PreviousRunID)
			}
		}
	}

	return &StageResult{
		Metadata: metadata,
	}, nil
}

func renderResolvedPipeline(cfg *config.PipelineConfig) ([]byte, error) {
	return yaml.Marshal(cfg)
}

func stableInputSource(resolved config.ResolvedInputFile, fallbackPath, fallbackConfigPath string) config.ResolvedInputFile {
	if strings.TrimSpace(resolved.Path) != "" || strings.TrimSpace(resolved.ConfigPath) != "" || strings.TrimSpace(resolved.Source) != "" {
		if strings.TrimSpace(resolved.ConfigPath) == "" {
			resolved.ConfigPath = fallbackConfigPath
		}
		if strings.TrimSpace(resolved.Source) == "" {
			resolved.Source = "session_config"
		}
		return resolved
	}
	return config.ResolvedInputFile{
		Path:       fallbackPath,
		ConfigPath: fallbackConfigPath,
		Source:     "session_config",
	}
}

func resolveConfigRelativePath(input config.ResolvedInputFile) (string, error) {
	basePath := strings.TrimSpace(input.ConfigPath)
	if basePath == "" {
		return "", fmt.Errorf("source config path is required")
	}
	return resolvePath(filepath.Dir(basePath), input.Path)
}

func resolveAudioInputs(sessionDir string, inputs config.SessionInputsConfig) ([]string, bool, error) {
	hasLocal := strings.TrimSpace(inputs.AudioDir) != "" || len(inputs.AudioFiles) > 0
	if inputs.AudioS3 != nil {
		if hasLocal {
			return nil, false, fmt.Errorf("audio_dir/audio_files and audio_s3 are mutually exclusive")
		}
		return nil, true, nil
	}

	local, err := resolveLocalAudioFiles(sessionDir, inputs)
	if err != nil {
		return nil, false, err
	}
	return local, false, nil
}

func resolveLocalAudioFiles(sessionDir string, inputs config.SessionInputsConfig) ([]string, error) {
	if len(inputs.AudioFiles) > 0 {
		out := make([]string, 0, len(inputs.AudioFiles))
		for _, p := range inputs.AudioFiles {
			resolved, err := resolvePath(sessionDir, p)
			if err != nil {
				return nil, err
			}
			if !isFlac(resolved) {
				return nil, fmt.Errorf("audio file %q must have .flac extension", resolved)
			}
			if err := requireFile(resolved, "audio file"); err != nil {
				return nil, err
			}
			out = append(out, resolved)
		}
		sort.Strings(out)
		return out, nil
	}

	audioDir, err := resolvePath(sessionDir, inputs.AudioDir)
	if err != nil {
		return nil, err
	}
	entries, err := os.ReadDir(audioDir)
	if err != nil {
		return nil, fmt.Errorf("read audio directory %q: %w", audioDir, err)
	}

	out := make([]string, 0)
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		full := filepath.Join(audioDir, name)
		if !isFlac(full) {
			continue
		}
		if err := requireFile(full, "audio file"); err != nil {
			return nil, err
		}
		out = append(out, full)
	}

	if len(out) == 0 {
		return nil, fmt.Errorf("no .flac files found in audio directory %q", audioDir)
	}
	sort.Strings(out)
	return out, nil
}

func materializeLocalAudioInputs(env *Env, paths artifacts.SessionPaths, resolvedAudio []string, registerInput func(kind, path, checksum string)) error {
	copiedByDest := map[string]string{}
	for _, src := range resolvedAudio {
		base := filepath.Base(src)
		if prev, exists := copiedByDest[base]; exists && prev != src {
			return fmt.Errorf("duplicate audio basename %q from %q and %q", base, prev, src)
		}
		copiedByDest[base] = src

		dst := filepath.Join(paths.AudioDir, base)
		checksum, err := copyFileIfChanged(env.ArtifactStore, src, dst)
		if err != nil {
			return fmt.Errorf("materialize audio %q: %w", base, err)
		}
		registerInput("audio", dst, checksum)
	}
	return nil
}

type s3AudioMaterializationStats struct {
	CacheHits   int
	CacheMisses int
	Downloads   int
}

func materializeS3AudioInputs(ctx context.Context, env *Env, m *manifest.Manifest, sessionID string, inputs *[]manifest.InputRecord) (s3AudioMaterializationStats, error) {
	if env.ObjectStore == nil {
		return s3AudioMaterializationStats{}, fmt.Errorf("s3 audio input requires object store backend")
	}
	if env.Config == nil || env.Config.Pipeline == nil || env.Config.Session == nil || env.Config.Pipeline.Storage.S3 == nil || env.Config.Session.Inputs.AudioS3 == nil {
		return s3AudioMaterializationStats{}, fmt.Errorf("s3 audio input requires pipeline.storage.s3 and session.inputs.audio_s3 configuration")
	}

	campaign := strings.TrimSpace(env.Config.Session.Campaign)
	if campaign == "" {
		return s3AudioMaterializationStats{}, fmt.Errorf("session campaign is required for s3 audio input")
	}
	runID := strings.TrimSpace(m.RunID)
	if runID == "" {
		return s3AudioMaterializationStats{}, fmt.Errorf("run id is required for s3 audio input")
	}

	sessionPrefix := artifacts.S3SessionPrefix(env.Config.Pipeline.Storage.S3.RootPrefix, campaign, sessionID)
	audioPrefix := artifacts.S3AudioPrefix(sessionPrefix, env.Config.Session.Inputs.AudioS3.Prefix)
	objects, err := env.ObjectStore.List(ctx, audioPrefix)
	if err != nil {
		return s3AudioMaterializationStats{}, fmt.Errorf("list s3 audio objects under %q: %w", audioPrefix, err)
	}

	audioObjects := make([]storage.ObjectInfo, 0, len(objects))
	for _, obj := range objects {
		key := strings.TrimSpace(obj.Key)
		if key == "" || strings.HasSuffix(key, "/") {
			continue
		}
		if !isFlac(key) {
			continue
		}
		audioObjects = append(audioObjects, obj)
	}
	sort.Slice(audioObjects, func(i, j int) bool {
		return audioObjects[i].Key < audioObjects[j].Key
	})
	if len(audioObjects) == 0 {
		return s3AudioMaterializationStats{}, fmt.Errorf("no .flac files found under s3 audio prefix %q", audioPrefix)
	}

	spoolAudioDir := strings.TrimSpace(m.LocalSpoolDir)
	if spoolAudioDir == "" {
		spoolAudioDir = artifacts.SessionSpoolAudioDir(env.Config.Pipeline.Spool.Root, campaign, sessionID, runID)
	}
	workAudioDir := filepath.Join(pathsWorkDirForManifest(env, m, sessionID), "audio")

	if err := os.MkdirAll(spoolAudioDir, 0o755); err != nil {
		return s3AudioMaterializationStats{}, fmt.Errorf("create spool audio directory %q: %w", spoolAudioDir, err)
	}
	if err := os.MkdirAll(workAudioDir, 0o755); err != nil {
		return s3AudioMaterializationStats{}, fmt.Errorf("create work audio directory %q: %w", workAudioDir, err)
	}

	seenBase := map[string]string{}
	stats := s3AudioMaterializationStats{}
	cacheEnabled := env.Config.Pipeline.Cache.S3Audio == nil || *env.Config.Pipeline.Cache.S3Audio
	for _, obj := range audioObjects {
		base := path.Base(obj.Key)
		if prev, exists := seenBase[base]; exists && prev != obj.Key {
			return s3AudioMaterializationStats{}, fmt.Errorf("duplicate s3 audio basename %q from %q and %q", base, prev, obj.Key)
		}
		seenBase[base] = obj.Key

		spoolPath := filepath.Join(spoolAudioDir, base)
		workPath := filepath.Join(workAudioDir, base)
		result, err := audio.MaterializeS3Audio(ctx, audio.S3MaterializeRequest{
			Store:        env.ObjectStore,
			Object:       obj,
			Bucket:       strings.TrimSpace(env.Config.Pipeline.Storage.S3.Bucket),
			CacheRoot:    strings.TrimSpace(env.Config.Pipeline.Cache.Root),
			CacheEnabled: cacheEnabled,
			SpoolPath:    spoolPath,
			DestPath:     workPath,
		})
		if err != nil {
			return s3AudioMaterializationStats{}, err
		}
		if result.CacheHit {
			stats.CacheHits++
		} else {
			stats.CacheMisses++
		}
		if result.Downloaded {
			stats.Downloads++
		}

		*inputs = append(*inputs, manifest.InputRecord{
			Kind:      "audio",
			Path:      workPath,
			Checksum:  result.Checksum,
			Source:    "s3",
			S3Bucket:  strings.TrimSpace(env.Config.Pipeline.Storage.S3.Bucket),
			S3Key:     obj.Key,
			S3Size:    obj.Size,
			S3ETag:    obj.ETag,
			SpoolPath: result.SpoolPath,
			CachePath: result.CachePath,
		})
	}

	return stats, nil
}

func countAudioInputs(inputs []manifest.InputRecord) int {
	count := 0
	for _, in := range inputs {
		if in.Kind == "audio" {
			count++
		}
	}
	return count
}

func collectPreparePreviousRequirements(cfg *config.Config) []artifacts.PreviousArtifactRequirement {
	if cfg == nil || cfg.Pipeline == nil || cfg.Pipeline.Scriptorium == nil {
		return nil
	}
	return artifacts.CollectPreviousArtifactRequirements(cfg.Pipeline.Scriptorium.Artifacts)
}

func clearManagedPreviousState(paths artifacts.SessionPaths) error {
	previousDir := filepath.Clean(paths.PreviousDir)
	sessionRoot := filepath.Clean(paths.Root)
	if strings.TrimSpace(previousDir) == "" || strings.TrimSpace(sessionRoot) == "" {
		return fmt.Errorf("previous/session root paths are required")
	}
	if previousDir == sessionRoot {
		return fmt.Errorf("refusing to clear session root as previous cache: %q", previousDir)
	}
	prefix := sessionRoot + string(filepath.Separator)
	if !strings.HasPrefix(previousDir, prefix) {
		return fmt.Errorf("refusing to clear path outside session root: %q", previousDir)
	}
	if filepath.Base(previousDir) != config.PathPreviousDirSegment {
		return fmt.Errorf("refusing to clear non-previous path %q", previousDir)
	}
	if err := os.RemoveAll(previousDir); err != nil {
		return err
	}
	return os.MkdirAll(previousDir, 0o755)
}

func pathsWorkDirForManifest(env *Env, m *manifest.Manifest, sessionID string) string {
	if env == nil || env.Config == nil || env.Config.Pipeline == nil || env.Config.Session == nil {
		return ""
	}
	campaign := strings.TrimSpace(env.Config.Session.Campaign)
	if campaign == "" {
		return ""
	}
	if m != nil && strings.TrimSpace(m.LocalWorkDir) != "" {
		return strings.TrimSpace(m.LocalWorkDir)
	}
	runID := ""
	if m != nil {
		runID = strings.TrimSpace(m.RunID)
	}
	if runID != "" {
		return artifacts.SessionRunRootForCampaign(env.Config.Pipeline.Workspace.Root, campaign, sessionID, runID)
	}
	return artifacts.SessionWorkDirForCampaign(env.Config.Pipeline.Workspace.Root, campaign, sessionID)
}

func resolvePath(baseDir, p string) (string, error) {
	trimmed := strings.TrimSpace(p)
	if trimmed == "" {
		return "", fmt.Errorf("path is required")
	}
	if filepath.IsAbs(trimmed) {
		return filepath.Clean(trimmed), nil
	}
	return filepath.Clean(filepath.Join(baseDir, trimmed)), nil
}

func requireFile(path string, label string) error {
	if strings.TrimSpace(path) == "" {
		return fmt.Errorf("%s path is required", label)
	}
	info, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("%s %q not found: %w", label, path, err)
	}
	if info.IsDir() {
		return fmt.Errorf("%s %q is a directory", label, path)
	}
	return nil
}

func isFlac(path string) bool {
	return strings.EqualFold(filepath.Ext(path), ".flac")
}

func copyFileIfChanged(store artifacts.Store, src, dst string) (string, error) {
	srcChecksum, err := store.Checksum(src)
	if err != nil {
		return "", err
	}

	exists, err := store.Exists(dst)
	if err != nil {
		return "", err
	}
	if exists {
		dstChecksum, err := store.Checksum(dst)
		if err != nil {
			return "", err
		}
		if dstChecksum == srcChecksum {
			return srcChecksum, nil
		}
	}

	data, err := os.ReadFile(src)
	if err != nil {
		return "", err
	}
	if err := store.WriteFileAtomic(dst, data, 0o644); err != nil {
		return "", err
	}
	return srcChecksum, nil
}

func writeBytesIfChanged(store artifacts.Store, dst string, data []byte) (string, error) {
	digest := sha256.Sum256(data)
	targetChecksum := hex.EncodeToString(digest[:])

	exists, err := store.Exists(dst)
	if err != nil {
		return "", err
	}
	if exists {
		dstChecksum, err := store.Checksum(dst)
		if err != nil {
			return "", err
		}
		if dstChecksum == targetChecksum {
			return targetChecksum, nil
		}
	}

	if err := store.WriteFileAtomic(dst, data, 0o644); err != nil {
		return "", err
	}
	return targetChecksum, nil
}