narratio/internal/stage/trim.go

package stage

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

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

type trimStage struct{}

func (trimStage) Name() string { return "trim" }

func (trimStage) Declares() IODecl {
	return IODecl{
		Inputs: []artifacts.Ref{
			{Kind: "transcript_normalized", Category: "transcripts", RelativePath: "transcripts/normalized.json"},
		},
		Outputs: []artifacts.Ref{
			{Kind: "transcript_trimmed", Category: "transcripts", RelativePath: "transcripts/trimmed.json"},
			{Kind: "session_bounds", Category: "artifacts", RelativePath: "artifacts/session_bounds.json"},
		},
	}
}

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

	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("trim: session id is required")
	}

	paths := env.ArtifactStore.SessionPaths(sessionID)
	normalizedPath, normalizedSource, err := discoverNormalizedTranscript(m, paths)
	if err != nil {
		return nil, fmt.Errorf("trim: resolve normalized transcript: %w", err)
	}
	if normalizedPath == "" {
		return nil, fmt.Errorf("trim: normalized transcript input is required; run normalize stage first")
	}
	if err := validateProcessedTranscriptOutput(normalizedPath); err != nil {
		return nil, fmt.Errorf("trim: normalized transcript %q invalid: %w", normalizedPath, err)
	}

	trimCfg := env.Config.Pipeline.Trim
	enabled := trimCfg != nil && trimCfg.Enabled

	trimmedPath, err := resolveTrimmedOutputPath(paths, trimCfg)
	if err != nil {
		return nil, fmt.Errorf("trim: resolve trimmed output path: %w", err)
	}

	logPaths := []string{}
	generatedConfigs := []string{}
	metadata := map[string]any{
		"stage":                        "trim",
		"trim_enabled":                 enabled,
		"normalized_transcript_path":   normalizedPath,
		"normalized_transcript_source": normalizedSource,
		"trimmed_output_path":          trimmedPath,
	}

	if !enabled {
		if err := copyTranscript(env.ArtifactStore, normalizedPath, trimmedPath); err != nil {
			return nil, fmt.Errorf("trim: copy normalized transcript to trimmed output: %w", err)
		}
		if err := validateProcessedTranscriptOutput(trimmedPath); err != nil {
			return nil, fmt.Errorf("trim: copied trimmed transcript %q invalid: %w", trimmedPath, err)
		}
		metadata["trim_action"] = "copy_disabled"
		return &StageResult{
			Outputs: []artifacts.Ref{{
				Kind:         "transcript_trimmed",
				Category:     "transcripts",
				SessionID:    sessionID,
				AbsolutePath: trimmedPath,
			}},
			Metadata: metadata,
		}, nil
	}

	if env.Scriptorium == nil {
		return nil, fmt.Errorf("trim: scriptorium adapter is required when trim is enabled")
	}
	if env.Seriatim == nil {
		return nil, fmt.Errorf("trim: seriatim adapter is required when trim is enabled")
	}
	if env.Config.Pipeline.Scriptorium == nil {
		return nil, fmt.Errorf("trim: pipeline.scriptorium is required when trim is enabled")
	}

	boundsCfg := trimCfg.Bounds
	boundsOutputPath, err := resolveScriptoriumOutputPath(paths, boundsCfg.OutputPath)
	if err != nil {
		return nil, fmt.Errorf("trim: resolve bounds output path: %w", err)
	}
	boundsStdoutLogPath := filepath.Join(paths.LogsDir, "scriptorium.bounds.stdout.log")
	boundsStderrLogPath := filepath.Join(paths.LogsDir, "scriptorium.bounds.stderr.log")
	boundsGeneratedConfigPath := filepath.Join(paths.ConfigDir, "scriptorium.bounds.generated.yml")
	boundsTimeout, err := resolveScriptoriumTimeout(env.Config.Pipeline.Scriptorium.Timeout, boundsCfg.Timeout)
	if err != nil {
		return nil, fmt.Errorf("trim: resolve bounds timeout: %w", err)
	}

	inputPaths := map[string]string{
		boundsCfg.TranscriptInputName: normalizedPath,
	}
	vars := map[string]string{}

	metadata["bounds_prompt_id"] = boundsCfg.PromptID
	metadata["bounds_profile_id"] = boundsCfg.ProfileID
	metadata["bounds_output_path"] = boundsOutputPath
	metadata["bounds_timeout"] = boundsTimeout.String()
	metadata["bounds_input_name"] = boundsCfg.TranscriptInputName
	metadata["bounds_input_path"] = normalizedPath
	metadata["bounds_render_debug_enabled"] = boundsCfg.RenderDebug

	renderOutputPath := ""
	if boundsCfg.RenderDebug {
		renderOutputPath, err = resolveScriptoriumOutputPath(paths, boundsCfg.RenderOutputPath)
		if err != nil {
			return nil, fmt.Errorf("trim: resolve bounds render output path: %w", err)
		}
		renderStdoutLogPath := filepath.Join(paths.LogsDir, "scriptorium.bounds.render.stdout.log")
		renderStderrLogPath := filepath.Join(paths.LogsDir, "scriptorium.bounds.render.stderr.log")
		renderGeneratedConfigPath := filepath.Join(paths.ConfigDir, "scriptorium.bounds.render.generated.yml")

		renderReq := scriptorium.RenderArtifactRequest{
			Binary:              env.Config.Pipeline.Scriptorium.Binary,
			ConfigPath:          env.Config.Pipeline.Scriptorium.ConfigPath,
			PromptID:            boundsCfg.PromptID,
			ProfileID:           boundsCfg.ProfileID,
			InputPaths:          inputPaths,
			Vars:                vars,
			OutputPath:          renderOutputPath,
			StdoutLogPath:       renderStdoutLogPath,
			StderrLogPath:       renderStderrLogPath,
			GeneratedConfigPath: renderGeneratedConfigPath,
			Timeout:             boundsTimeout,
		}
		renderRes, renderErr := env.Scriptorium.RenderArtifact(ctx, renderReq)
		if renderErr != nil {
			return nil, fmt.Errorf("trim: scriptorium bounds render failed: %w", renderErr)
		}
		if renderRes.ValidationFailed {
			return nil, fmt.Errorf("trim: scriptorium bounds render returned validation_failed=true")
		}
		finalRenderOutputPath := coalesceString(renderRes.OutputPath, renderReq.OutputPath)
		if err := requireNonEmptyFile(finalRenderOutputPath, "bounds render output"); err != nil {
			return nil, fmt.Errorf("trim: %w", err)
		}
		if err := validateJSONFile(finalRenderOutputPath); err != nil {
			return nil, fmt.Errorf("trim: bounds render diagnostics %q invalid json: %w", finalRenderOutputPath, err)
		}

		logPaths = append(logPaths, renderStdoutLogPath, renderStderrLogPath)
		generatedConfigs = append(generatedConfigs, renderGeneratedConfigPath)
		metadata["bounds_render_output_path"] = finalRenderOutputPath
		metadata["bounds_render_stdout_log_path"] = renderStdoutLogPath
		metadata["bounds_render_stderr_log_path"] = renderStderrLogPath
		metadata["bounds_render_generated_config_path"] = renderGeneratedConfigPath
		metadata["bounds_render_adapter_exit_code"] = renderRes.ExitCode
		metadata["bounds_render_adapter_duration_ms"] = renderRes.Duration.Milliseconds()
		metadata["bounds_render_adapter_command_mode"] = renderRes.CommandMode
		if renderRes.Metadata != nil {
			metadata["bounds_render_adapter_metadata"] = renderRes.Metadata
		}
	}

	boundsReq := scriptorium.RunArtifactRequest{
		Binary:              env.Config.Pipeline.Scriptorium.Binary,
		ConfigPath:          env.Config.Pipeline.Scriptorium.ConfigPath,
		PromptID:            boundsCfg.PromptID,
		ProfileID:           boundsCfg.ProfileID,
		InputPaths:          inputPaths,
		Vars:                vars,
		OutputPath:          boundsOutputPath,
		StdoutLogPath:       boundsStdoutLogPath,
		StderrLogPath:       boundsStderrLogPath,
		GeneratedConfigPath: boundsGeneratedConfigPath,
		Timeout:             boundsTimeout,
	}

	boundsRes, boundsRunErr := env.Scriptorium.RunArtifact(ctx, boundsReq)
	if boundsRunErr != nil {
		if boundsRes.ValidationFailed {
			return nil, fmt.Errorf(
				"trim: scriptorium bounds validation failed (prompt_id=%q, output_path=%q, exit_code=%d, stdout_log=%q, stderr_log=%q): %w",
				boundsReq.PromptID,
				coalesceString(boundsRes.OutputPath, boundsReq.OutputPath),
				boundsRes.ExitCode,
				coalesceString(boundsRes.StdoutLogPath, boundsReq.StdoutLogPath),
				coalesceString(boundsRes.StderrLogPath, boundsReq.StderrLogPath),
				boundsRunErr,
			)
		}
		return nil, fmt.Errorf("trim: scriptorium bounds generation failed: %w", boundsRunErr)
	}
	if boundsRes.ValidationFailed {
		return nil, fmt.Errorf("trim: scriptorium bounds run returned validation_failed=true")
	}

	finalBoundsOutputPath := coalesceString(boundsRes.OutputPath, boundsReq.OutputPath)
	if err := requireNonEmptyFile(finalBoundsOutputPath, "session bounds output"); err != nil {
		return nil, fmt.Errorf("trim: %w", err)
	}

	boundsPayload, err := contracts.ParseSessionBoundsFile(finalBoundsOutputPath)
	if err != nil {
		return nil, fmt.Errorf("trim: parse bounds output %q: %w", finalBoundsOutputPath, err)
	}
	if err := contracts.ValidateSessionBoundsAgainstTranscript(boundsPayload, normalizedPath); err != nil {
		return nil, fmt.Errorf("trim: validate bounds output %q against transcript %q: %w", finalBoundsOutputPath, normalizedPath, err)
	}

	keepSelector, copyUnchanged, err := contracts.BuildSeriatimKeepSelector(boundsPayload)
	if err != nil {
		return nil, fmt.Errorf("trim: build keep selector from bounds output: %w", err)
	}
	trimAction := strings.ToLower(strings.TrimSpace(boundsPayload.TrimAction))
	if trimAction == "" {
		trimAction = "trim"
	}

	logPaths = append(logPaths, boundsStdoutLogPath, boundsStderrLogPath)
	generatedConfigs = append(generatedConfigs, boundsGeneratedConfigPath)
	metadata["bounds_confidence"] = boundsPayload.Confidence
	metadata["trim_action"] = trimAction
	metadata["start_segment_id"] = boundsPayload.StartSegmentID
	metadata["end_segment_id"] = boundsPayload.EndSegmentID
	metadata["warnings"] = boundsPayload.Warnings
	metadata["keep_selector"] = keepSelector
	metadata["bounds_output_path"] = finalBoundsOutputPath
	metadata["bounds_stdout_log_path"] = boundsStdoutLogPath
	metadata["bounds_stderr_log_path"] = boundsStderrLogPath
	metadata["bounds_generated_config_path"] = boundsGeneratedConfigPath
	metadata["bounds_adapter_exit_code"] = boundsRes.ExitCode
	metadata["bounds_adapter_duration_ms"] = boundsRes.Duration.Milliseconds()
	metadata["bounds_adapter_command_mode"] = boundsRes.CommandMode
	metadata["bounds_adapter_prompt_id"] = boundsRes.PromptID
	metadata["bounds_adapter_profile_id"] = boundsRes.ProfileID
	metadata["bounds_adapter_output_path"] = boundsRes.OutputPath
	metadata["bounds_adapter_generated_config"] = boundsRes.GeneratedConfigPath
	metadata["bounds_adapter_stdout_log_path"] = boundsRes.StdoutLogPath
	metadata["bounds_adapter_stderr_log_path"] = boundsRes.StderrLogPath
	if boundsRes.Metadata != nil {
		metadata["bounds_adapter_metadata"] = boundsRes.Metadata
	}

	if copyUnchanged {
		if err := copyTranscript(env.ArtifactStore, normalizedPath, trimmedPath); err != nil {
			return nil, fmt.Errorf("trim: copy normalized transcript to trimmed output: %w", err)
		}
	} else {
		trimStdoutLogPath := filepath.Join(paths.LogsDir, "seriatim.trim.stdout.log")
		trimStderrLogPath := filepath.Join(paths.LogsDir, "seriatim.trim.stderr.log")
		trimGeneratedConfigPath := filepath.Join(paths.ConfigDir, "seriatim.trim.generated.yml")
		trimTimeout, err := resolveTrimSeriatimTimeout(env.Config.Pipeline.Seriatim.Timeout)
		if err != nil {
			return nil, fmt.Errorf("trim: resolve seriatim timeout: %w", err)
		}
		trimReq := seriatim.TrimRequest{
			Binary:              env.Config.Pipeline.Seriatim.Binary,
			InputTranscriptPath: normalizedPath,
			OutputTrimmedPath:   trimmedPath,
			KeepSelector:        keepSelector,
			StdoutLogPath:       trimStdoutLogPath,
			StderrLogPath:       trimStderrLogPath,
			GeneratedConfigPath: trimGeneratedConfigPath,
			Timeout:             trimTimeout,
		}
		trimRes, trimErr := env.Seriatim.Trim(ctx, trimReq)
		if trimErr != nil {
			return nil, fmt.Errorf("trim: seriatim trim failed: %w", trimErr)
		}
		logPaths = append(logPaths, trimStdoutLogPath, trimStderrLogPath)
		generatedConfigs = append(generatedConfigs, trimGeneratedConfigPath)
		metadata["seriatim_trim_stdout_log_path"] = trimStdoutLogPath
		metadata["seriatim_trim_stderr_log_path"] = trimStderrLogPath
		metadata["seriatim_trim_generated_config_path"] = trimGeneratedConfigPath
		metadata["seriatim_trim_exit_code"] = trimRes.ExitCode
		metadata["seriatim_trim_duration_ms"] = trimRes.Duration.Milliseconds()
		metadata["seriatim_trim_invoked_binary"] = trimRes.InvokedBinary
		metadata["seriatim_trim_keep_selector"] = trimRes.KeepSelector
		if trimRes.Metadata != nil {
			metadata["seriatim_trim_metadata"] = trimRes.Metadata
		}
	}

	if err := validateProcessedTranscriptOutput(trimmedPath); err != nil {
		return nil, fmt.Errorf("trim: trimmed transcript %q invalid: %w", trimmedPath, err)
	}

	outputs := []artifacts.Ref{
		{
			Kind:         "transcript_trimmed",
			Category:     "transcripts",
			SessionID:    sessionID,
			AbsolutePath: trimmedPath,
		},
		{
			Kind:         "session_bounds",
			Category:     "artifacts",
			SessionID:    sessionID,
			AbsolutePath: finalBoundsOutputPath,
		},
	}

	return &StageResult{
		Outputs:          outputs,
		Logs:             logPaths,
		GeneratedConfigs: generatedConfigs,
		Metadata:         metadata,
	}, nil
}

func resolveTrimmedOutputPath(paths artifacts.SessionPaths, cfg *config.TrimConfig) (string, error) {
	configured := ""
	if cfg != nil {
		configured = strings.TrimSpace(cfg.OutputPath)
	}
	if configured == "" {
		return filepath.Join(paths.TranscriptsDir, "trimmed.json"), nil
	}
	return resolveScriptoriumOutputPath(paths, configured)
}

func copyTranscript(store artifacts.Store, src, dst string) error {
	data, err := os.ReadFile(src)
	if err != nil {
		return fmt.Errorf("read source transcript: %w", err)
	}
	if err := store.WriteFileAtomic(dst, data, 0o644); err != nil {
		return fmt.Errorf("write destination transcript: %w", err)
	}
	return nil
}

func resolveTrimSeriatimTimeout(raw string) (time.Duration, error) {
	trimmed := strings.TrimSpace(raw)
	if trimmed == "" {
		return 0, nil
	}
	d, err := time.ParseDuration(trimmed)
	if err != nil {
		return 0, fmt.Errorf("parse duration %q: %w", raw, err)
	}
	if d <= 0 {
		return 0, fmt.Errorf("duration must be > 0")
	}
	return d, nil
}