package stage

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

	"gitea.maximumdirect.net/eric/narratio/internal/adapters/notarius"
	"gitea.maximumdirect.net/eric/narratio/internal/artifactmodel"
	"gitea.maximumdirect.net/eric/narratio/internal/artifacts"
	"gitea.maximumdirect.net/eric/narratio/internal/config"
	"gitea.maximumdirect.net/eric/narratio/internal/fileops"
	"gitea.maximumdirect.net/eric/narratio/internal/manifest"
	"gitea.maximumdirect.net/eric/narratio/internal/pathsafe"
)

const (
	extractSkipReason      = "notarius_disabled"
	extractLaneOutputKind  = "notarius_lane"
	extractIndexOutputKind = "notarius_index"
	maxDiagnosticSummaries = 100
)

type extractStage struct{}

func (extractStage) Name() string { return "extract" }

func (extractStage) Declares() IODecl {
	return IODecl{
		Inputs: []artifacts.Ref{{
			Kind:         artifactmodel.TranscriptOutputKindFinalTrimmed,
			SourceID:     artifactmodel.SourceTranscriptFinalTrimmed,
			Category:     "transcripts",
			RelativePath: artifactmodel.TranscriptPathFinalTrimmed,
		}},
		Outputs: []artifacts.Ref{
			{Kind: extractLaneOutputKind, Category: "artifacts", RelativePath: "artifacts/notarius/<run-id>/lanes/*.json"},
			{Kind: extractIndexOutputKind, Category: "artifacts", RelativePath: "artifacts/notarius/<run-id>/index.json"},
		},
	}
}

func (extractStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (*StageResult, error) {
	if env == nil || env.Config == nil || env.Config.Pipeline == nil || env.Config.Session == nil {
		return nil, fmt.Errorf("extract: resolved stage environment config is required")
	}
	notariusConfig := env.Config.Pipeline.Notarius
	if notariusConfig == nil || !notariusConfig.Enabled {
		return &StageResult{
			Disposition: StageDispositionSkipped,
			SkipReason:  extractSkipReason,
			Metadata: map[string]any{
				"stage":            "extract",
				"notarius_enabled": false,
				"reason":           extractSkipReason,
			},
		}, nil
	}
	if env.ArtifactStore == nil {
		return nil, fmt.Errorf("extract: artifact store is required")
	}
	if env.Notarius == nil {
		return nil, fmt.Errorf("extract: notarius adapter is required")
	}
	if m == nil {
		return nil, fmt.Errorf("extract: session manifest is required")
	}

	sessionID := strings.TrimSpace(m.SessionID)
	if sessionID == "" {
		sessionID = strings.TrimSpace(env.Config.Session.SessionID)
	}
	campaign := strings.TrimSpace(m.Campaign)
	if campaign == "" {
		campaign = strings.TrimSpace(env.Config.Session.Campaign)
	}
	runID := strings.TrimSpace(m.RunID)
	if sessionID == "" || campaign == "" || runID == "" {
		return nil, fmt.Errorf("extract: session id, campaign, and run id are required")
	}
	if !safePathSegment(runID) {
		return nil, fmt.Errorf("extract: run id %q is not a safe path segment", runID)
	}

	paths := sessionPathsForEnv(env, sessionID)
	input, err := artifacts.ResolveSessionArtifact(paths, m, artifacts.ArtifactTranscriptFinalTrimmed)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve final-trimmed transcript: %w", err)
	}

	timeout, err := time.ParseDuration(strings.TrimSpace(notariusConfig.Timeout))
	if err != nil || timeout <= 0 {
		return nil, fmt.Errorf("extract: invalid notarius timeout %q", notariusConfig.Timeout)
	}
	resolvedBinary, err := resolveExecutable(notariusConfig.Binary)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve notarius binary: %w", err)
	}
	configPath, err := absolutePath(notariusConfig.ConfigPath)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve notarius config path: %w", err)
	}
	inputPath, err := absolutePath(input.Path)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve transcript input path: %w", err)
	}
	workingDirectory, err := absolutePath(notariusConfig.WorkingDirectory)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve notarius working directory: %w", err)
	}

	workspaceRoot := strings.TrimSpace(paths.WorkspaceRoot)
	if workspaceRoot == "" {
		workspaceRoot = env.Config.Pipeline.Workspace.Root
	}
	receiptPath, err := absolutePath(artifacts.SessionRunNotariusReceiptPathForCampaign(workspaceRoot, campaign, sessionID, runID))
	if err != nil {
		return nil, fmt.Errorf("extract: resolve receipt path: %w", err)
	}
	logPath, err := absolutePath(artifacts.SessionRunNotariusLogPathForCampaign(workspaceRoot, campaign, sessionID, runID))
	if err != nil {
		return nil, fmt.Errorf("extract: resolve stderr path: %w", err)
	}
	outputRoot, err := absolutePath(artifacts.SessionRunNotariusOutputRootForCampaign(workspaceRoot, campaign, sessionID, runID))
	if err != nil {
		return nil, fmt.Errorf("extract: resolve output root: %w", err)
	}
	durableBundle, err := absolutePath(artifacts.SessionNotariusBundleDirForCampaign(workspaceRoot, campaign, sessionID, runID))
	if err != nil {
		return nil, fmt.Errorf("extract: resolve durable bundle path: %w", err)
	}
	for _, directory := range []string{filepath.Dir(receiptPath), outputRoot, filepath.Dir(durableBundle)} {
		if err := os.MkdirAll(directory, 0o755); err != nil {
			return nil, fmt.Errorf("extract: create directory %q: %w", directory, err)
		}
	}

	fingerprint, err := extractionFingerprint(resolvedBinary, configPath, notariusConfig, timeout, workingDirectory)
	if err != nil {
		return nil, fmt.Errorf("extract: build configuration fingerprint: %w", err)
	}
	request := notarius.RunRequest{
		Binary: resolvedBinary, ConfigPath: configPath, PipelineID: notariusConfig.PipelineID,
		InputPath: inputPath, OutputRoot: outputRoot, WorkingDirectory: workingDirectory,
		ReceiptPath: receiptPath, LogPath: logPath, Timeout: timeout,
	}
	adapterResult, err := env.Notarius.Run(ctx, request)
	if err != nil {
		return nil, fmt.Errorf("extract: run notarius: %w", err)
	}
	if strings.TrimSpace(adapterResult.BundleRoot) == "" || strings.TrimSpace(adapterResult.Index.Path) == "" {
		return nil, fmt.Errorf("extract: notarius result is missing bundle or index path")
	}
	if adapterResult.Receipt.RunID == "" || adapterResult.Receipt.PipelineID != notariusConfig.PipelineID {
		return nil, fmt.Errorf("extract: notarius receipt identity is missing or incompatible")
	}

	selected, err := selectRequiredNotariusLanes(env.ArtifactStore, notariusConfig.Outputs, adapterResult)
	if err != nil {
		return nil, err
	}
	indexRelative, err := pathsafe.SlashRelativeFromRoot(adapterResult.BundleRoot, adapterResult.Index.Path)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve staging index relative path: %w", err)
	}
	rejectionsRelative, err := pathsafe.SlashRelativeFromRoot(adapterResult.BundleRoot, adapterResult.Index.RejectedPath)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve staging rejections relative path: %w", err)
	}
	warningsRelative, err := pathsafe.SlashRelativeFromRoot(adapterResult.BundleRoot, adapterResult.Index.WarningsPath)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve staging warnings relative path: %w", err)
	}
	stagingIndexChecksum, err := checksumRegularFile(adapterResult.Index.Path, false)
	if err != nil {
		return nil, fmt.Errorf("extract: validate staging index: %w", err)
	}

	if err := fileops.PromoteDirectory(adapterResult.BundleRoot, durableBundle); err != nil {
		return nil, fmt.Errorf("extract: promote notarius bundle: %w", err)
	}
	promotedIndexPath, err := pathsafe.JoinSlashRelativeUnderRoot(durableBundle, indexRelative)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve promoted index: %w", err)
	}
	promotedIndexChecksum, err := checksumRegularFile(promotedIndexPath, false)
	if err != nil {
		return nil, fmt.Errorf("extract: checksum promoted index: %w", err)
	}
	if promotedIndexChecksum != stagingIndexChecksum {
		return nil, fmt.Errorf("extract: promoted index checksum differs from staging index")
	}
	promotedRejectionsPath, err := pathsafe.JoinSlashRelativeUnderRoot(durableBundle, rejectionsRelative)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve promoted rejections: %w", err)
	}
	promotedWarningsPath, err := pathsafe.JoinSlashRelativeUnderRoot(durableBundle, warningsRelative)
	if err != nil {
		return nil, fmt.Errorf("extract: resolve promoted warnings: %w", err)
	}

	outputs := make([]artifacts.Ref, 0, len(selected)+1)
	for _, lane := range selected {
		promotedPath, err := pathsafe.JoinSlashRelativeUnderRoot(durableBundle, lane.RelativePath)
		if err != nil {
			return nil, fmt.Errorf("extract: resolve promoted lane %q: %w", lane.Descriptor.LaneID, err)
		}
		checksum, err := checksumRegularFile(promotedPath, true)
		if err != nil {
			return nil, fmt.Errorf("extract: validate promoted lane %q: %w", lane.Descriptor.LaneID, err)
		}
		if checksum != lane.StagingChecksum {
			return nil, fmt.Errorf("extract: promoted lane %q checksum differs from staging payload", lane.Descriptor.LaneID)
		}
		relativePath, err := pathsafe.SlashRelativeFromRoot(paths.Root, promotedPath)
		if err != nil {
			return nil, fmt.Errorf("extract: derive lane %q session-relative path: %w", lane.Descriptor.LaneID, err)
		}
		outputs = append(outputs, artifacts.Ref{
			Kind: extractLaneOutputKind, SourceID: artifacts.ExtractionArtifactSourceID(lane.Key),
			Category: "artifacts", SessionID: sessionID, RelativePath: relativePath,
			AbsolutePath: promotedPath, Checksum: checksum,
			Contract: &artifactmodel.ContractMetadata{
				MediaType: lane.Descriptor.MediaType, SchemaID: lane.Descriptor.SchemaID,
				SchemaVersion: lane.Descriptor.SchemaVersion, ModuleKey: lane.Descriptor.ModuleKey,
			},
			ExternalProvenance: &artifactmodel.ExternalProvenance{
				System: "notarius", RunID: adapterResult.Receipt.RunID,
				PipelineID: adapterResult.Receipt.PipelineID, ArtifactID: lane.Descriptor.LaneID,
			},
		})
	}
	indexSessionRelative, err := pathsafe.SlashRelativeFromRoot(paths.Root, promotedIndexPath)
	if err != nil {
		return nil, fmt.Errorf("extract: derive index session-relative path: %w", err)
	}
	outputs = append(outputs, artifacts.Ref{
		Kind: extractIndexOutputKind, Category: "artifacts", SessionID: sessionID,
		RelativePath: indexSessionRelative, AbsolutePath: promotedIndexPath, Checksum: promotedIndexChecksum,
	})

	metadata := map[string]any{
		"stage":                     "extract",
		"notarius_enabled":          true,
		"bundle_root":               durableBundle,
		"receipt_path":              receiptPath,
		"diagnostic_path":           logPath,
		"rejections_path":           promotedRejectionsPath,
		"warnings_path":             promotedWarningsPath,
		"narratio_run_id":           runID,
		"configuration_fingerprint": fingerprint,
		"receipt": map[string]any{
			"run_id": adapterResult.Receipt.RunID, "pipeline_id": adapterResult.Receipt.PipelineID,
			"normalized_output_count": adapterResult.Receipt.NormalizedOutputCount,
			"rejected_output_count":   adapterResult.Receipt.RejectedOutputCount,
			"warning_count":           adapterResult.Receipt.WarningCount,
			"validation_status":       adapterResult.Receipt.ValidationStatus,
		},
		"rejections": boundedRejectionMetadata(adapterResult.Rejections),
		"warnings":   boundedWarningMetadata(adapterResult.Warnings),
	}
	return &StageResult{
		Outputs:  outputs,
		Logs:     []string{receiptPath, logPath},
		Metadata: metadata,
	}, nil
}

type selectedNotariusLane struct {
	Key             string
	Descriptor      notarius.LaneDescriptor
	RelativePath    string
	StagingChecksum string
}

func selectRequiredNotariusLanes(
	store artifacts.Store,
	required map[string]config.NotariusOutputConfig,
	result notarius.RunResult,
) ([]selectedNotariusLane, error) {
	keys := make([]string, 0, len(required))
	for key := range required {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	selected := make([]selectedNotariusLane, 0, len(keys))
	for _, key := range keys {
		expected := required[key]
		for _, rejection := range result.Rejections {
			if rejection.LaneID == expected.LaneID {
				return nil, fmt.Errorf("extract: required lane %q was rejected (reason_code=%q)", expected.LaneID, rejection.ReasonCode)
			}
		}
		matches := make([]notarius.LaneDescriptor, 0, 1)
		for _, descriptor := range result.Index.Lanes {
			if descriptor.LaneID == expected.LaneID {
				matches = append(matches, descriptor)
			}
		}
		if len(matches) != 1 {
			return nil, fmt.Errorf("extract: required lane %q has %d descriptors, want exactly one", expected.LaneID, len(matches))
		}
		descriptor := matches[0]
		if descriptor.MediaType != expected.MediaType || descriptor.SchemaID != expected.SchemaID ||
			descriptor.SchemaVersion != expected.SchemaVersion ||
			(expected.ModuleKey != "" && descriptor.ModuleKey != expected.ModuleKey) {
			return nil, fmt.Errorf("extract: required lane %q descriptor contract is incompatible", expected.LaneID)
		}
		checksum, err := checksumRegularFile(descriptor.Path, true)
		if err != nil {
			return nil, fmt.Errorf("extract: validate required lane %q: %w", expected.LaneID, err)
		}
		if store != nil {
			storeChecksum, err := store.Checksum(descriptor.Path)
			if err != nil {
				return nil, fmt.Errorf("extract: checksum required lane %q: %w", expected.LaneID, err)
			}
			if storeChecksum != checksum {
				return nil, fmt.Errorf("extract: inconsistent staging checksum for required lane %q", expected.LaneID)
			}
		}
		relative, err := pathsafe.SlashRelativeFromRoot(result.BundleRoot, descriptor.Path)
		if err != nil {
			return nil, fmt.Errorf("extract: resolve required lane %q relative path: %w", expected.LaneID, err)
		}
		selected = append(selected, selectedNotariusLane{
			Key: key, Descriptor: descriptor, RelativePath: relative, StagingChecksum: checksum,
		})
	}
	return selected, nil
}

func checksumRegularFile(path string, requireJSON bool) (string, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return "", err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return "", fmt.Errorf("path %q must be a regular file without symlinks", path)
	}
	if info.Size() == 0 {
		return "", fmt.Errorf("path %q must be non-empty", path)
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	if requireJSON && !json.Valid(data) {
		return "", fmt.Errorf("path %q is not valid JSON", path)
	}
	digest := sha256.Sum256(data)
	return hex.EncodeToString(digest[:]), nil
}

type fingerprintOutput struct {
	Key           string `json:"key"`
	LaneID        string `json:"lane_id"`
	MediaType     string `json:"media_type"`
	SchemaID      string `json:"schema_id"`
	SchemaVersion string `json:"schema_version"`
	ModuleKey     string `json:"module_key,omitempty"`
}

type fingerprintDocument struct {
	Binary           string              `json:"binary"`
	ConfigPath       string              `json:"config_path"`
	PipelineID       string              `json:"pipeline_id"`
	Timeout          string              `json:"timeout"`
	WorkingDirectory string              `json:"working_directory"`
	Outputs          []fingerprintOutput `json:"outputs"`
}

func extractionFingerprint(
	binary, configPath string,
	cfg *config.NotariusConfig,
	timeout time.Duration,
	workingDirectory string,
) (string, error) {
	keys := make([]string, 0, len(cfg.Outputs))
	for key := range cfg.Outputs {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	outputs := make([]fingerprintOutput, 0, len(keys))
	for _, key := range keys {
		output := cfg.Outputs[key]
		outputs = append(outputs, fingerprintOutput{
			Key: key, LaneID: output.LaneID, MediaType: output.MediaType,
			SchemaID: output.SchemaID, SchemaVersion: output.SchemaVersion, ModuleKey: output.ModuleKey,
		})
	}
	payload, err := json.Marshal(fingerprintDocument{
		Binary: binary, ConfigPath: configPath, PipelineID: cfg.PipelineID,
		Timeout: timeout.String(), WorkingDirectory: workingDirectory, Outputs: outputs,
	})
	if err != nil {
		return "", err
	}
	digest := sha256.Sum256(payload)
	return hex.EncodeToString(digest[:]), nil
}

func resolveExecutable(value string) (string, error) {
	resolved, err := exec.LookPath(strings.TrimSpace(value))
	if err != nil {
		return "", err
	}
	return absolutePath(resolved)
}

func absolutePath(value string) (string, error) {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return "", fmt.Errorf("path is required")
	}
	resolved, err := filepath.Abs(trimmed)
	if err != nil {
		return "", err
	}
	return filepath.Clean(resolved), nil
}

func safePathSegment(value string) bool {
	return value != "" && value != "." && value != ".." && filepath.Base(value) == value &&
		!strings.ContainsAny(value, `/\\`)
}

func boundedRejectionMetadata(values []notarius.RejectionSummary) []map[string]any {
	limit := len(values)
	if limit > maxDiagnosticSummaries {
		limit = maxDiagnosticSummaries
	}
	result := make([]map[string]any, 0, limit)
	for _, value := range values[:limit] {
		result = append(result, map[string]any{
			"stage": value.Stage, "step_id": value.StepID, "lane_id": value.LaneID,
			"module_key": value.ModuleKey, "chunk_id": value.ChunkID,
			"validator_name": value.ValidatorName, "reason_code": value.ReasonCode,
		})
	}
	return result
}

func boundedWarningMetadata(values []notarius.WarningSummary) []map[string]any {
	limit := len(values)
	if limit > maxDiagnosticSummaries {
		limit = maxDiagnosticSummaries
	}
	result := make([]map[string]any, 0, limit)
	for _, value := range values[:limit] {
		result = append(result, map[string]any{"scope": value.Scope, "reason_code": value.ReasonCode})
	}
	return result
}