narratio/internal/stage/polish.go

package stage

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

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

type polishStage struct{}

func (polishStage) Name() string { return "polish" }

func (polishStage) Declares() IODecl {
	return IODecl{
		Inputs: []artifacts.Ref{
			{Kind: "transcript_base", Category: "transcripts", RelativePath: "transcripts/base.json"},
			{Kind: "glossary", Category: "inputs", RelativePath: "inputs/glossary.yml"},
		},
		Outputs: []artifacts.Ref{
			{Kind: "transcript_polished", Category: "transcripts", RelativePath: "transcripts/polished.json"},
			{Kind: "audita_report", Category: "artifacts", RelativePath: "artifacts/audita.report.json"},
		},
	}
}

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

	paths := sessionPathsForEnv(env, sessionID)
	runLayout, err := resolveRunStageLayout(env, m, paths, sessionID, "polish")
	if err != nil {
		return nil, fmt.Errorf("polish: resolve run-stage layout: %w", err)
	}
	mergedPath, source, err := discoverMergedTranscript(m, paths)
	if err != nil {
		return nil, fmt.Errorf("polish: resolve merged transcript: %w", err)
	}
	if mergedPath == "" {
		return nil, fmt.Errorf("polish: merged transcript input is required")
	}
	if err := validateTranscriptJSONFile(mergedPath); err != nil {
		return nil, fmt.Errorf("polish: merged transcript %q invalid: %w", mergedPath, err)
	}

	glossaryPath := filepath.Join(paths.InputsDir, "glossary.yml")
	if err := requireFile(glossaryPath, "glossary.yml"); err != nil {
		return nil, fmt.Errorf("polish: %w", err)
	}

	canonicalProcessedPath := filepath.Join(paths.TranscriptsDir, "polished.json")
	processedPath, err := runLocalPathForCanonical(runLayout, paths, canonicalProcessedPath)
	if err != nil {
		return nil, fmt.Errorf("polish: resolve run-local processed transcript path: %w", err)
	}
	canonicalReportPath := filepath.Join(paths.ArtifactsDir, "audita.report.json")
	reportPath := canonicalReportPath
	if runLayout.Enabled {
		reportPath, err = runLocalPathForCanonical(runLayout, paths, canonicalReportPath)
		if err != nil {
			return nil, fmt.Errorf("polish: resolve run-local report path: %w", err)
		}
	}
	workDir := filepath.Join(paths.ArtifactsDir, "audita-work")
	stdoutPath := filepath.Join(paths.LogsDir, "audita.stdout.log")
	stderrPath := filepath.Join(paths.LogsDir, "audita.stderr.log")
	generatedConfigPath := filepath.Join(paths.ConfigDir, "audita.generated.yml")
	if runLayout.Enabled {
		workDir = filepath.Join(runLayout.ScratchDir, "audita-work")
		stdoutPath = filepath.Join(runLayout.LogsDir, "audita.stdout.log")
		stderrPath = filepath.Join(runLayout.LogsDir, "audita.stderr.log")
		generatedConfigPath = filepath.Join(runLayout.ConfigDir, "audita.generated.yml")
	}

	reportEnabled := env.Config.Pipeline.Audita.Report != nil && *env.Config.Pipeline.Audita.Report
	req := audita.PolishRequest{
		GeneratedConfigPath:      generatedConfigPath,
		MergedTranscriptPath:     mergedPath,
		OutputProcessedPath:      processedPath,
		GlossaryPath:             glossaryPath,
		ReportPath:               "",
		WorkDir:                  workDir,
		Modules:                  append([]string(nil), env.Config.Pipeline.Audita.Modules...),
		BaseURL:                  env.Config.Pipeline.Audita.BaseURL,
		Model:                    env.Config.Pipeline.Audita.Model,
		TranscriptDescription:    env.Config.Pipeline.Audita.TranscriptDescription,
		ConfigPath:               env.Config.Pipeline.Audita.ConfigPath,
		OutputSchema:             env.Config.Pipeline.Audita.OutputSchema,
		WorkDirRetention:         env.Config.Pipeline.Audita.WorkDirRetention,
		TotalLLMConcurrency:      env.Config.Pipeline.Audita.TotalLLMConcurrency,
		ProposalLLMConcurrency:   env.Config.Pipeline.Audita.ProposalLLMConcurrency,
		ValidationModel:          env.Config.Pipeline.Audita.ValidationModel,
		ValidationLLMConcurrency: env.Config.Pipeline.Audita.ValidationLLMConcurrency,
		StdoutLogPath:            stdoutPath,
		StderrLogPath:            stderrPath,
	}
	if reportEnabled {
		req.ReportPath = reportPath
	}

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

	finalProcessedPath := processedPath
	if strings.TrimSpace(res.ProcessedTranscriptPath) != "" {
		finalProcessedPath = res.ProcessedTranscriptPath
	}
	if err := validateProcessedTranscriptOutput(finalProcessedPath); err != nil {
		return nil, fmt.Errorf("polish: processed transcript %q invalid: %w", finalProcessedPath, err)
	}

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

	materializedProcessed, err := materializeRunLocalOutput(env.ArtifactStore, finalProcessedPath, canonicalProcessedPath, artifacts.Ref{
		Kind:      "transcript_polished",
		Category:  "transcripts",
		SessionID: sessionID,
	})
	if err != nil {
		return nil, fmt.Errorf("polish: materialize canonical polished transcript: %w", err)
	}
	outputs := []artifacts.Ref{materializedProcessed}
	if reportEnabled {
		materializedReport, err := materializeRunLocalOutput(env.ArtifactStore, finalReportPath, canonicalReportPath, artifacts.Ref{
			Kind:      "audita_report",
			Category:  "artifacts",
			SessionID: sessionID,
		})
		if err != nil {
			return nil, fmt.Errorf("polish: materialize canonical report: %w", err)
		}
		outputs = append(outputs, materializedReport)
	}

	var validationConcurrency any
	if env.Config.Pipeline.Audita.ValidationLLMConcurrency != nil {
		validationConcurrency = *env.Config.Pipeline.Audita.ValidationLLMConcurrency
	}
	var totalLLMConcurrency any
	if env.Config.Pipeline.Audita.TotalLLMConcurrency != nil {
		totalLLMConcurrency = *env.Config.Pipeline.Audita.TotalLLMConcurrency
	}
	var proposalLLMConcurrency any
	if env.Config.Pipeline.Audita.ProposalLLMConcurrency != nil {
		proposalLLMConcurrency = *env.Config.Pipeline.Audita.ProposalLLMConcurrency
	}
	reportCanonicalPath := ""
	if reportEnabled {
		reportCanonicalPath = canonicalReportPath
	}

	meta := map[string]any{
		"stage":                         "polish",
		"merged_transcript_path":        mergedPath,
		"merged_transcript_source":      source,
		"glossary_path":                 glossaryPath,
		"run_output_path":               finalProcessedPath,
		"output_path":                   canonicalProcessedPath,
		"run_report_path":               finalReportPath,
		"report_path":                   reportCanonicalPath,
		"audita_work_dir":               workDir,
		"report_enabled":                reportEnabled,
		"modules":                       append([]string(nil), req.Modules...),
		"base_url":                      req.BaseURL,
		"model":                         req.Model,
		"transcript_description":        req.TranscriptDescription,
		"config_path":                   req.ConfigPath,
		"output_schema":                 req.OutputSchema,
		"work_dir_retention":            req.WorkDirRetention,
		"validation_model":              req.ValidationModel,
		"total_llm_concurrency":         totalLLMConcurrency,
		"proposal_llm_concurrency":      proposalLLMConcurrency,
		"validation_llm_concurrency":    validationConcurrency,
		"llm_api_key_env":               env.Config.Pipeline.Audita.LLMAPIKeyEnv,
		"timeout":                       env.Config.Pipeline.Audita.Timeout,
		"binary":                        env.Config.Pipeline.Audita.Binary,
		"generated_config_path":         generatedConfigPath,
		"stdout_log_path":               stdoutPath,
		"stderr_log_path":               stderrPath,
		"adapter_duration_ms":           res.Duration.Milliseconds(),
		"adapter_exit_code":             res.ExitCode,
		"adapter_invoked_binary":        res.InvokedBinary,
		"adapter_processed_output_path": res.ProcessedTranscriptPath,
		"adapter_report_path":           res.ReportPath,
		"adapter_generated_config_path": res.GeneratedConfigPath,
		"adapter_work_dir":              res.WorkDir,
		"adapter_stdout_log_path":       res.StdoutLogPath,
		"adapter_stderr_log_path":       res.StderrLogPath,
		"credential_env_var":            env.Config.Pipeline.Audita.LLMAPIKeyEnv,
		"credential_present":            false,
	}
	if res.Metadata != nil {
		meta["adapter_metadata"] = res.Metadata
		if value, ok := res.Metadata["credential_present"]; ok {
			meta["credential_present"] = value
		}
		if value, ok := res.Metadata["credential_env_var"]; ok {
			meta["credential_env_var"] = value
		}
	}

	return &StageResult{
		Outputs:          outputs,
		Logs:             []string{stdoutPath, stderrPath},
		GeneratedConfigs: []string{generatedConfigPath},
		Metadata:         meta,
	}, nil
}

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

	deduped := dedupeAndSortPaths(candidates)
	for _, p := range deduped {
		if info, err := os.Stat(p); err == nil && !info.IsDir() {
			return p, "manifest.merge.outputs", nil
		}
	}

	fallback := filepath.Join(paths.TranscriptsDir, "base.json")
	if info, err := os.Stat(fallback); err == nil && !info.IsDir() {
		return filepath.Clean(fallback), "fallback.transcripts_dir", nil
	}
	if len(deduped) > 0 {
		return deduped[0], "manifest.merge.outputs", nil
	}
	return "", "", nil
}

func validateProcessedTranscriptOutput(path string) error {
	if err := requireFile(path, "processed transcript"); err != nil {
		return err
	}
	data, err := readExternalResult(path, "audita processed transcript result")
	if err != nil {
		return err
	}

	var payload map[string]any
	if err := json.Unmarshal(data, &payload); err != nil {
		return fmt.Errorf("decode json: %w", err)
	}
	segments, ok := payload["segments"]
	if !ok {
		return fmt.Errorf("top-level segments is required")
	}
	if _, ok := segments.([]any); !ok {
		return fmt.Errorf("top-level segments must be an array")
	}
	return nil
}