package stage

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

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

type mergeStage struct{}

func (mergeStage) Name() string { return "merge" }

func (mergeStage) Declares() IODecl {
	return IODecl{
		Inputs: []artifacts.Ref{
			{Kind: "transcript_raw", Category: "transcripts", RelativePath: "transcripts/raw/*.json"},
			{Kind: "speakers", Category: "inputs", RelativePath: "inputs/speakers.yml"},
			{Kind: "autocorrect", Category: "inputs", RelativePath: "inputs/autocorrect.yml"},
		},
		Outputs: []artifacts.Ref{
			{Kind: "transcript_base", Category: "transcripts", RelativePath: "transcripts/base.json"},
			{Kind: "seriatim_report", Category: "artifacts", RelativePath: "artifacts/seriatim.report.json"},
		},
	}
}

func (mergeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (*StageResult, error) {
	if env == nil || env.Config == nil {
		return nil, fmt.Errorf("merge: stage environment config is required")
	}
	if env.ArtifactStore == nil {
		return nil, fmt.Errorf("merge: artifact store is required")
	}
	if env.Config.Pipeline == nil || env.Config.Session == nil {
		return nil, fmt.Errorf("merge: resolved config must include pipeline and session")
	}
	if env.Seriatim == nil {
		return nil, fmt.Errorf("merge: seriatim adapter 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("merge: session id is required")
	}

	paths := sessionPathsForEnv(env, sessionID)
	runLayout, err := resolveRunStageLayout(env, m, paths, sessionID, "merge")
	if err != nil {
		return nil, fmt.Errorf("merge: resolve run-stage layout: %w", err)
	}

	inputs, err := discoverRawTranscripts(m, paths)
	if err != nil {
		return nil, fmt.Errorf("merge: resolve raw transcripts: %w", err)
	}
	if len(inputs) == 0 {
		return nil, fmt.Errorf("merge: no raw transcript inputs found")
	}
	for _, in := range inputs {
		if err := validateTranscriptJSONFile(in); err != nil {
			return nil, fmt.Errorf("merge: input transcript %q invalid: %w", in, err)
		}
	}

	speakersPath := filepath.Join(paths.InputsDir, "speakers.yml")
	if err := requireFile(speakersPath, "speakers.yml"); err != nil {
		return nil, fmt.Errorf("merge: %w", err)
	}
	autocorrectPath := filepath.Join(paths.InputsDir, "autocorrect.yml")
	if err := requireFile(autocorrectPath, "autocorrect.yml"); err != nil {
		return nil, fmt.Errorf("merge: %w", err)
	}

	canonicalMergedPath := filepath.Join(paths.TranscriptsDir, "base.json")
	mergedPath, err := runLocalPathForCanonical(runLayout, paths, canonicalMergedPath)
	if err != nil {
		return nil, fmt.Errorf("merge: resolve run-local merged transcript path: %w", err)
	}
	canonicalReportPath := filepath.Join(paths.ArtifactsDir, "seriatim.report.json")
	reportPath := canonicalReportPath
	if runLayout.Enabled {
		reportPath, err = runLocalPathForCanonical(runLayout, paths, canonicalReportPath)
		if err != nil {
			return nil, fmt.Errorf("merge: resolve run-local report path: %w", err)
		}
	}
	stdoutPath := filepath.Join(paths.LogsDir, "seriatim.stdout.log")
	stderrPath := filepath.Join(paths.LogsDir, "seriatim.stderr.log")
	genCfgPath := filepath.Join(paths.ConfigDir, "seriatim.generated.yml")
	if runLayout.Enabled {
		stdoutPath = filepath.Join(runLayout.LogsDir, "seriatim.stdout.log")
		stderrPath = filepath.Join(runLayout.LogsDir, "seriatim.stderr.log")
		genCfgPath = filepath.Join(runLayout.ConfigDir, "seriatim.generated.yml")
	}

	normalizedInputs, normalizeLogs, normalizeConfigs, normalizeMeta, err := normalizeMergeInputs(ctx, env, inputs, paths, runLayout)
	if err != nil {
		return nil, err
	}

	reportEnabled := env.Config.Pipeline.Seriatim.Report != nil && *env.Config.Pipeline.Seriatim.Report
	req := seriatim.MergeRequest{
		GeneratedConfigPath:        genCfgPath,
		InputTranscriptPaths:       normalizedInputs,
		OutputMergedTranscriptPath: mergedPath,
		ReportPath:                 "",
		SpeakersPath:               speakersPath,
		AutocorrectPath:            autocorrectPath,
		StdoutLogPath:              stdoutPath,
		StderrLogPath:              stderrPath,
	}
	if reportEnabled {
		req.ReportPath = reportPath
	}

	res, err := env.Seriatim.Run(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("merge: seriatim merge failed: %w", err)
	}

	finalMergedPath := mergedPath
	if strings.TrimSpace(res.MergedTranscriptPath) != "" {
		finalMergedPath = res.MergedTranscriptPath
	}
	if err := validateTranscriptJSONFile(finalMergedPath); err != nil {
		return nil, fmt.Errorf("merge: merged transcript %q invalid: %w", finalMergedPath, err)
	}

	finalReportPath := req.ReportPath
	if strings.TrimSpace(res.ReportPath) != "" {
		finalReportPath = res.ReportPath
	}
	if reportEnabled {
		if err := validateTranscriptJSONFile(finalReportPath); err != nil {
			return nil, fmt.Errorf("merge: report %q invalid: %w", finalReportPath, err)
		}
	}

	materializedMerged, err := materializeRunLocalOutput(env.ArtifactStore, finalMergedPath, canonicalMergedPath, artifacts.Ref{
		Kind:      "transcript_base",
		Category:  "transcripts",
		SessionID: sessionID,
	})
	if err != nil {
		return nil, fmt.Errorf("merge: materialize canonical base transcript: %w", err)
	}
	outputs := []artifacts.Ref{materializedMerged}
	if reportEnabled {
		materializedReport, err := materializeRunLocalOutput(env.ArtifactStore, finalReportPath, canonicalReportPath, artifacts.Ref{
			Kind:      "seriatim_report",
			Category:  "artifacts",
			SessionID: sessionID,
		})
		if err != nil {
			return nil, fmt.Errorf("merge: materialize canonical report: %w", err)
		}
		outputs = append(outputs, materializedReport)
	}

	coalesceGap := any(nil)
	if env.Config.Pipeline.Seriatim.CoalesceGap != nil {
		coalesceGap = *env.Config.Pipeline.Seriatim.CoalesceGap
	}
	reportCanonicalPath := ""
	if reportEnabled {
		reportCanonicalPath = canonicalReportPath
	}

	meta := map[string]any{
		"stage":                    "merge",
		"input_transcripts_count":  len(normalizedInputs),
		"input_transcript_paths":   inputs,
		"normalized_inputs_count":  len(normalizedInputs),
		"normalized_input_paths":   normalizedInputs,
		"normalize_inputs":         normalizeMeta,
		"output_schema":            env.Config.Pipeline.Seriatim.OutputSchema,
		"coalesce_gap":             coalesceGap,
		"report_enabled":           reportEnabled,
		"run_output_path":          finalMergedPath,
		"output_path":              canonicalMergedPath,
		"run_report_path":          finalReportPath,
		"report_path":              reportCanonicalPath,
		"timeout":                  env.Config.Pipeline.Seriatim.Timeout,
		"binary":                   env.Config.Pipeline.Seriatim.Binary,
		"adapter_duration_ms":      res.Duration.Milliseconds(),
		"adapter_exit_code":        res.ExitCode,
		"adapter_invoked_binary":   res.InvokedBinary,
		"adapter_output_schema":    res.OutputSchema,
		"generated_config_path":    genCfgPath,
		"stdout_log_path":          stdoutPath,
		"stderr_log_path":          stderrPath,
		"adapter_report_path":      res.ReportPath,
		"adapter_merged_out_path":  res.MergedTranscriptPath,
		"adapter_generated_config": res.GeneratedConfigPath,
	}
	if res.Metadata != nil {
		meta["adapter_metadata"] = res.Metadata
	}

	return &StageResult{
		Outputs:          outputs,
		Logs:             append(normalizeLogs, stdoutPath, stderrPath),
		GeneratedConfigs: append(normalizeConfigs, genCfgPath),
		Metadata:         meta,
	}, nil
}

type normalizeMergeInputMeta struct {
	InputPath         string `json:"input_path"`
	OutputPath        string `json:"output_path"`
	StdoutLogPath     string `json:"stdout_log_path"`
	StderrLogPath     string `json:"stderr_log_path"`
	GeneratedConfig   string `json:"generated_config_path"`
	DurationMs        int64  `json:"duration_ms"`
	ExitCode          int    `json:"exit_code"`
	InvokedBinary     string `json:"invoked_binary"`
	OutputSchema      string `json:"output_schema"`
	AdapterReportPath string `json:"adapter_report_path,omitempty"`
	AdapterOutputPath string `json:"adapter_output_path,omitempty"`
}

func normalizeMergeInputs(
	ctx context.Context,
	env *Env,
	rawInputs []string,
	paths artifacts.SessionPaths,
	runLayout runStageLayout,
) ([]string, []string, []string, []normalizeMergeInputMeta, error) {
	normalizedInputs := make([]string, 0, len(rawInputs))
	logs := make([]string, 0, len(rawInputs)*2)
	configs := make([]string, 0, len(rawInputs))
	meta := make([]normalizeMergeInputMeta, 0, len(rawInputs))
	normalizedDir := filepath.Join(paths.TranscriptsRawDir, "normalized")
	if runLayout.Enabled {
		normalizedDir = filepath.Join(runLayout.ScratchDir, "normalized")
	}
	if err := os.MkdirAll(normalizedDir, 0o755); err != nil {
		return nil, nil, nil, nil, fmt.Errorf("merge: ensure normalized transcripts directory %q: %w", normalizedDir, err)
	}

	var timeout time.Duration
	timeoutRaw := strings.TrimSpace(env.Config.Pipeline.Seriatim.Timeout)
	if timeoutRaw != "" {
		parsed, err := time.ParseDuration(timeoutRaw)
		if err != nil {
			return nil, nil, nil, nil, fmt.Errorf("merge: parse seriatim timeout %q: %w", env.Config.Pipeline.Seriatim.Timeout, err)
		}
		timeout = parsed
	}
	for _, input := range rawInputs {
		base := strings.TrimSuffix(filepath.Base(input), filepath.Ext(input))
		outPath := filepath.Join(normalizedDir, base+".normalized.json")
		stdoutPath := filepath.Join(paths.LogsDir, "seriatim.normalize."+base+".stdout.log")
		stderrPath := filepath.Join(paths.LogsDir, "seriatim.normalize."+base+".stderr.log")
		cfgPath := filepath.Join(paths.ConfigDir, "seriatim.normalize."+base+".generated.yml")
		if runLayout.Enabled {
			stdoutPath = filepath.Join(runLayout.LogsDir, "seriatim.normalize."+base+".stdout.log")
			stderrPath = filepath.Join(runLayout.LogsDir, "seriatim.normalize."+base+".stderr.log")
			cfgPath = filepath.Join(runLayout.ConfigDir, "seriatim.normalize."+base+".generated.yml")
		}

		req := seriatim.NormalizeRequest{
			Binary:               env.Config.Pipeline.Seriatim.Binary,
			InputTranscriptPath:  input,
			OutputNormalizedPath: outPath,
			OutputSchema:         env.Config.Pipeline.Seriatim.OutputSchema,
			ReportPath:           "",
			StdoutLogPath:        stdoutPath,
			StderrLogPath:        stderrPath,
			GeneratedConfigPath:  cfgPath,
			Timeout:              timeout,
		}
		res, err := env.Seriatim.Normalize(ctx, req)
		if err != nil {
			return nil, nil, nil, nil, fmt.Errorf("merge: normalize input %q failed: %w", input, err)
		}

		finalOutputPath := outPath
		if strings.TrimSpace(res.OutputNormalizedPath) != "" {
			finalOutputPath = strings.TrimSpace(res.OutputNormalizedPath)
		}
		if err := validateTranscriptJSONFile(finalOutputPath); err != nil {
			return nil, nil, nil, nil, fmt.Errorf("merge: normalized transcript %q invalid (input %q): %w", finalOutputPath, input, err)
		}

		normalizedInputs = append(normalizedInputs, finalOutputPath)
		logs = append(logs, stdoutPath, stderrPath)
		configs = append(configs, cfgPath)
		meta = append(meta, normalizeMergeInputMeta{
			InputPath:         input,
			OutputPath:        finalOutputPath,
			StdoutLogPath:     stdoutPath,
			StderrLogPath:     stderrPath,
			GeneratedConfig:   cfgPath,
			DurationMs:        res.Duration.Milliseconds(),
			ExitCode:          res.ExitCode,
			InvokedBinary:     res.InvokedBinary,
			OutputSchema:      res.OutputSchema,
			AdapterReportPath: res.ReportPath,
			AdapterOutputPath: res.OutputNormalizedPath,
		})
	}

	return normalizedInputs, logs, configs, meta, nil
}

func discoverRawTranscripts(m *manifest.Manifest, paths artifacts.SessionPaths) ([]string, error) {
	fromManifest := make([]string, 0)
	if m != nil && m.Stages != nil {
		if tr := m.Stages["transcribe"]; tr != nil {
			for _, out := range tr.Outputs {
				if out.Kind != "transcript_raw" {
					continue
				}
				p := strings.TrimSpace(out.LocalPath)
				if p == "" {
					continue
				}
				p = artifacts.ResolveSessionLocalPathForRead(paths, p)
				fromManifest = append(fromManifest, filepath.Clean(p))
			}
		}
	}

	if len(fromManifest) > 0 {
		deduped := dedupeAndSortPaths(fromManifest)
		return deduped, nil
	}

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

	out := make([]string, 0, len(entries))
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		p := filepath.Join(paths.TranscriptsRawDir, entry.Name())
		if !strings.EqualFold(filepath.Ext(p), ".json") {
			continue
		}
		out = append(out, p)
	}
	return dedupeAndSortPaths(out), nil
}

func dedupeAndSortPaths(paths []string) []string {
	if len(paths) == 0 {
		return nil
	}
	set := make(map[string]struct{}, len(paths))
	out := make([]string, 0, len(paths))
	for _, p := range paths {
		clean := filepath.Clean(strings.TrimSpace(p))
		if clean == "" {
			continue
		}
		if _, exists := set[clean]; exists {
			continue
		}
		set[clean] = struct{}{}
		out = append(out, clean)
	}
	sort.Strings(out)
	return out
}