package stage

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"golang.org/x/sync/errgroup"

	"gitea.maximumdirect.net/eric/narratio/internal/adapters/whisperx"
	"gitea.maximumdirect.net/eric/narratio/internal/artifacts"
	"gitea.maximumdirect.net/eric/narratio/internal/manifest"
)

type transcribeStage struct{}

func (transcribeStage) Name() string { return "transcribe" }

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

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

	paths := sessionPathsForEnv(env, sessionID)
	reconciliation, err := reconcileTranscribeSpeakers(ctx, env, m)
	if err != nil {
		return nil, fmt.Errorf("transcribe: reconcile speakers: %w", err)
	}
	state := newTranscribeExecutionState(env, m, paths, reconciliation.current)
	if err := state.checkpointStartingCollection(ctx); err != nil {
		return nil, fmt.Errorf("transcribe: checkpoint reconciled speaker collection before execution: %w", err)
	}

	jobs := append([]transcribeReconciliationJob(nil), reconciliation.work...)
	if len(jobs) > 0 && env.WhisperX == nil {
		return nil, fmt.Errorf("transcribe: whisperx adapter is required")
	}
	if len(jobs) > 0 {
		runLayout, err := resolveRunStageLayout(env, m, paths, sessionID, "transcribe")
		if err != nil {
			return nil, fmt.Errorf("transcribe: resolve run-stage layout: %w", err)
		}
		for index := range jobs {
			jobs[index].runOutputPath, err = runLocalPathForCanonical(runLayout, paths, jobs[index].canonicalOutputPath)
			if err != nil {
				return nil, fmt.Errorf("transcribe: resolve run-local output path for %q: %w", jobs[index].speakerID, err)
			}
		}
	}

	concurrency := 1
	if env.Config.Pipeline.WhisperX.Concurrency != nil && *env.Config.Pipeline.WhisperX.Concurrency > 0 {
		concurrency = *env.Config.Pipeline.WhisperX.Concurrency
	}
	if len(jobs) > 0 && concurrency > len(jobs) {
		concurrency = len(jobs)
	}

	group, groupCtx := errgroup.WithContext(ctx)
	group.SetLimit(concurrency)
	for _, job := range jobs {
		job := job
		group.Go(func() error {
			if err := groupCtx.Err(); err != nil {
				return err
			}
			if err := state.checkpointFailure(); err != nil {
				return err
			}
			res, runErr := env.WhisperX.Transcribe(groupCtx, whisperx.TranscribeRequest{
				SpeakerID:               job.speakerID,
				AudioPath:               job.audioPath,
				OutputRawTranscriptPath: job.runOutputPath,
			})
			if runErr != nil {
				return fmt.Errorf("speaker %q (%s): %w", job.speakerID, filepath.Base(job.audioPath), runErr)
			}

			meta := map[string]any{
				"attempts":    res.Attempts,
				"http_status": res.HTTPStatus,
				"duration_ms": res.Duration.Milliseconds(),
				"output_path": job.runOutputPath,
			}

			if err := state.acceptSpeaker(ctx, job, res.OutputRawTranscriptPath, meta); err != nil {
				return fmt.Errorf("speaker %q: %w", job.speakerID, err)
			}
			return nil
		})
	}
	groupErr := group.Wait()
	if err := state.checkpointFailure(); err != nil {
		return nil, err
	}
	completedCount := state.completedCount()
	if err := ctx.Err(); err != nil {
		return nil, fmt.Errorf("transcribe: canceled after completing %d of %d audio jobs: %w", completedCount, len(jobs), err)
	}
	if groupErr != nil {
		return nil, groupErr
	}
	if completedCount != len(jobs) {
		return nil, fmt.Errorf("transcribe: completed %d of %d planned audio jobs", completedCount, len(jobs))
	}
	outputs, outputPaths, runOutputPaths, orderedPerFile, err := state.successDetails()
	if err != nil {
		return nil, fmt.Errorf("transcribe: assemble completed outputs: %w", err)
	}

	retries := 0
	if env.Config.Pipeline.WhisperX.Retries != nil {
		retries = *env.Config.Pipeline.WhisperX.Retries
	}

	return &StageResult{
		Outputs: outputs,
		Metadata: map[string]any{
			"stage":             "transcribe",
			"audio_files_count": len(reconciliation.jobs),
			"language":          env.Config.Pipeline.WhisperX.Language,
			"concurrency":       concurrency,
			"retries":           retries,
			"retry_delay":       env.Config.Pipeline.WhisperX.RetryDelay,
			"timeout":           env.Config.Pipeline.WhisperX.Timeout,
			"run_output_paths":  runOutputPaths,
			"output_paths":      outputPaths,
			"per_file":          orderedPerFile,
		},
	}, nil
}

type transcribeJob struct {
	speakerID           string
	audioPath           string
	canonicalOutputPath string
	runOutputPath       string
}

func buildTranscribeJobs(m *manifest.Manifest, paths artifacts.SessionPaths) ([]transcribeJob, error) {
	audioFiles, err := discoverPreparedAudio(m, paths)
	if err != nil {
		return nil, err
	}
	jobs := make([]transcribeJob, 0, len(audioFiles))
	seenSpeaker := map[string]string{}
	for _, audioPath := range audioFiles {
		base := strings.TrimSuffix(filepath.Base(audioPath), filepath.Ext(audioPath))
		if base == "" {
			return nil, fmt.Errorf("could not derive speaker id from %q", audioPath)
		}
		if previous, exists := seenSpeaker[base]; exists {
			return nil, fmt.Errorf("duplicate speaker/audio basename %q from %q and %q", base, previous, audioPath)
		}
		seenSpeaker[base] = audioPath
		jobs = append(jobs, transcribeJob{
			speakerID:           base,
			audioPath:           audioPath,
			canonicalOutputPath: filepath.Join(paths.TranscriptsRawDir, base+".json"),
		})
	}
	return jobs, nil
}

func discoverPreparedAudio(m *manifest.Manifest, paths artifacts.SessionPaths) ([]string, error) {
	fromManifest := make([]string, 0)
	if m != nil {
		for _, in := range m.Inputs {
			if in.Kind != "audio" {
				continue
			}
			p := strings.TrimSpace(in.Path)
			if p == "" {
				continue
			}
			fromManifest = append(fromManifest, artifacts.ResolveSessionLocalPathForRead(paths, p))
		}
	}
	if len(fromManifest) > 0 {
		sort.Strings(fromManifest)
		if err := validateAudioFiles(fromManifest); err != nil {
			return nil, err
		}
		return fromManifest, nil
	}

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

	out := make([]string, 0)
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		p := filepath.Join(paths.AudioDir, entry.Name())
		if !isFlac(p) {
			continue
		}
		out = append(out, p)
	}
	sort.Strings(out)
	if err := validateAudioFiles(out); err != nil {
		return nil, err
	}
	return out, nil
}

func validateAudioFiles(paths []string) error {
	if len(paths) == 0 {
		return nil
	}
	seen := map[string]struct{}{}
	for _, p := range paths {
		clean := filepath.Clean(strings.TrimSpace(p))
		if clean == "" {
			return fmt.Errorf("audio path is required")
		}
		if !isFlac(clean) {
			return fmt.Errorf("audio file %q must have .flac extension", clean)
		}
		info, err := os.Stat(clean)
		if err != nil {
			return fmt.Errorf("audio file %q not found: %w", clean, err)
		}
		if !info.Mode().IsRegular() {
			return fmt.Errorf("audio file %q is not a regular file", clean)
		}
		if _, exists := seen[clean]; exists {
			return fmt.Errorf("duplicate audio file path %q", clean)
		}
		seen[clean] = struct{}{}
	}
	return nil
}

func validateTranscriptJSONFile(path string) error {
	if err := requireFile(path, "raw transcript"); err != nil {
		return err
	}
	data, err := readExternalResult(path, "whisperx transcript result")
	if err != nil {
		return err
	}
	var payload any
	if err := json.Unmarshal(data, &payload); err != nil {
		return fmt.Errorf("decode json: %w", err)
	}
	return nil
}