package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	clusterrecovery "github.com/pingcap-inc/clinic/pkg/cluster_recovery"
	"github.com/pingcap-inc/clinic/pkg/config"
	"github.com/pingcap-inc/clinic/pkg/datasource"
	"github.com/pingcap-inc/clinic/pkg/db"
)

const (
	exitOK    = 0
	exitUsage = 2
	exitError = 1
)

type commandRuntime struct {
	environment string
	source      clusterrecovery.Source
	readStore   clusterrecovery.ReadStore
	writeStore  clusterrecovery.RecoveryStore
}

type runtimeLoader func(string) (*commandRuntime, error)

type commandError struct {
	code  string
	cause error
}

func (e *commandError) Error() string { return e.code }
func (e *commandError) Unwrap() error { return e.cause }

func fail(code string, cause error) error {
	return &commandError{code: code, cause: cause}
}

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	code := execute(ctx, os.Args[1:], os.Stdout, os.Stderr, loadRuntime)
	os.Exit(code)
}

func execute(ctx context.Context, args []string, stdout, stderr io.Writer, loader runtimeLoader) int {
	if len(args) == 0 {
		_, _ = io.WriteString(stderr, "clinic recovery failed: usage\n")
		return exitUsage
	}
	var err error
	switch args[0] {
	case "plan":
		err = runPlan(ctx, args[1:], stdout, loader)
	case "apply":
		err = runApply(ctx, args[1:], stdout, loader)
	case "help", "-h", "--help":
		_, _ = io.WriteString(stdout, "usage: cluster_recovery <plan|apply> [flags]\n")
		return exitOK
	default:
		err = fail("usage", errors.New("unknown command"))
	}
	if err == nil {
		return exitOK
	}
	code := "internal_error"
	var safe *commandError
	if errors.As(err, &safe) {
		code = safe.code
	}
	_, _ = fmt.Fprintf(stderr, "clinic recovery failed: %s\n", code)
	if code == "usage" {
		return exitUsage
	}
	return exitError
}

type planFlags struct {
	configPath string
	seedPath   string
	outputDir  string
}

func parsePlanFlags(args []string) (planFlags, error) {
	var options planFlags
	flags := flag.NewFlagSet("plan", flag.ContinueOnError)
	flags.SetOutput(io.Discard)
	flags.StringVar(&options.configPath, "config", "", "Clinic configuration path")
	flags.StringVar(&options.seedPath, "seed", "", "incident candidate seed path")
	flags.StringVar(&options.outputDir, "output-dir", "", "new artifact directory")
	if err := flags.Parse(args); err != nil || flags.NArg() != 0 {
		return planFlags{}, fail("usage", err)
	}
	if options.configPath == "" || options.seedPath == "" || options.outputDir == "" {
		return planFlags{}, fail("usage", errors.New("required plan flag is missing"))
	}
	return options, nil
}

func runPlan(ctx context.Context, args []string, stdout io.Writer, loader runtimeLoader) error {
	options, err := parsePlanFlags(args)
	if err != nil {
		return err
	}
	seedFile, err := openPrivateInput(options.seedPath)
	if err != nil {
		return fail("invalid_seed", err)
	}
	seed, decodeErr := clusterrecovery.DecodeSeed(seedFile)
	closeErr := seedFile.Close()
	if decodeErr != nil || closeErr != nil {
		return fail("invalid_seed", errors.Join(decodeErr, closeErr))
	}
	if err := validateOutputDirectory(options.outputDir); err != nil {
		return fail("invalid_output_directory", err)
	}
	runtime, err := loader(options.configPath)
	if err != nil {
		return fail("runtime_initialization_failed", err)
	}
	if err := validateRuntime(runtime, false); err != nil {
		return fail("environment_guard", err)
	}
	output, err := clusterrecovery.NewPlanner(runtime.source, runtime.readStore).Plan(ctx, seed)
	if err != nil {
		return fail("planning_failed", err)
	}
	paths := planArtifactPaths(options.outputDir, output.Report.RunID)
	for _, path := range paths {
		if err := validateNewArtifactPath(path); err != nil {
			return fail("artifact_preflight_failed", err)
		}
	}
	artifacts := []any{&output.Report, &output.CreateManifest, &output.UpdateManifest}
	for index, path := range paths {
		if err := clusterrecovery.WriteJSONArtifact(path, artifacts[index]); err != nil {
			return fail("artifact_write_failed", err)
		}
	}
	summary := planSummary{
		Mode:               "plan",
		RunID:              output.Report.RunID,
		Counts:             output.Report.Counts,
		CreateActions:      len(output.CreateManifest.Actions),
		CreateManifestHash: output.CreateManifest.ManifestHash,
		UpdateActions:      len(output.UpdateManifest.Actions),
		UpdateManifestHash: output.UpdateManifest.ManifestHash,
	}
	if err := json.NewEncoder(stdout).Encode(summary); err != nil {
		return fail("summary_write_failed", err)
	}
	return nil
}

type planSummary struct {
	Mode               string                                 `json:"mode"`
	RunID              string                                 `json:"run_id"`
	Counts             map[clusterrecovery.Classification]int `json:"counts"`
	CreateActions      int                                    `json:"create_actions"`
	CreateManifestHash string                                 `json:"create_manifest_hash"`
	UpdateActions      int                                    `json:"update_actions"`
	UpdateManifestHash string                                 `json:"update_manifest_hash"`
}

type applyFlags struct {
	configPath               string
	manifestPath             string
	expectedManifestHash     string
	approvalRef              string
	workloadIdentity         string
	mutationClass            string
	phase                    string
	offset                   int
	limit                    int
	rollbackPath             string
	receiptPath              string
	priorReceiptPath         string
	expectedPriorReceiptHash string
	priorVerificationRef     string
	allowWrite               bool
}

func parseApplyFlags(args []string) (applyFlags, error) {
	var options applyFlags
	flags := flag.NewFlagSet("apply", flag.ContinueOnError)
	flags.SetOutput(io.Discard)
	flags.StringVar(&options.configPath, "config", "", "Clinic configuration path")
	flags.StringVar(&options.manifestPath, "manifest", "", "reviewed manifest path")
	flags.StringVar(&options.expectedManifestHash, "expected-manifest-hash", "", "reviewed manifest SHA-256")
	flags.StringVar(&options.approvalRef, "approval-ref", "", "reviewed PR or approval reference")
	flags.StringVar(&options.workloadIdentity, "workload-identity", "", "Kubernetes workload identity")
	flags.StringVar(&options.mutationClass, "mutation-class", "", "create or update")
	flags.StringVar(&options.phase, "phase", "", "canary or batch")
	flags.IntVar(&options.offset, "offset", 0, "first manifest action index")
	flags.IntVar(&options.limit, "limit", 0, "bounded action count")
	flags.StringVar(&options.rollbackPath, "rollback", "", "new rollback evidence path")
	flags.StringVar(&options.receiptPath, "receipt", "", "new execution receipt path")
	flags.StringVar(&options.priorReceiptPath, "prior-receipt", "", "successful prior execution receipt path")
	flags.StringVar(&options.expectedPriorReceiptHash, "expected-prior-receipt-hash", "", "verified prior receipt SHA-256")
	flags.StringVar(&options.priorVerificationRef, "prior-verification-ref", "", "reviewed prior-batch verification reference")
	flags.BoolVar(&options.allowWrite, "allow-write", false, "enable guarded staging writes")
	if err := flags.Parse(args); err != nil || flags.NArg() != 0 {
		return applyFlags{}, fail("usage", err)
	}
	if options.configPath == "" || options.manifestPath == "" || options.expectedManifestHash == "" ||
		options.approvalRef == "" || options.workloadIdentity == "" || options.mutationClass == "" || options.phase == "" || options.limit == 0 ||
		options.rollbackPath == "" || options.receiptPath == "" {
		return applyFlags{}, fail("usage", errors.New("required apply flag is missing"))
	}
	return options, nil
}

func runApply(ctx context.Context, args []string, stdout io.Writer, loader runtimeLoader) error {
	flags, err := parseApplyFlags(args)
	if err != nil {
		return err
	}
	manifestFile, err := openPrivateInput(flags.manifestPath)
	if err != nil {
		return fail("invalid_manifest", err)
	}
	manifest, decodeErr := clusterrecovery.DecodeManifest(manifestFile)
	closeErr := manifestFile.Close()
	if decodeErr != nil || closeErr != nil {
		return fail("invalid_manifest", errors.Join(decodeErr, closeErr))
	}

	var priorReceipt *clusterrecovery.Receipt
	if flags.priorReceiptPath != "" {
		priorFile, err := openPrivateInput(flags.priorReceiptPath)
		if err != nil {
			return fail("invalid_prior_receipt", err)
		}
		priorReceipt, decodeErr = clusterrecovery.DecodeReceipt(priorFile)
		closeErr = priorFile.Close()
		if decodeErr != nil || closeErr != nil {
			return fail("invalid_prior_receipt", errors.Join(decodeErr, closeErr))
		}
	}
	applyOptions := clusterrecovery.ApplyOptions{
		AllowWrite:               flags.allowWrite,
		ExpectedManifestHash:     flags.expectedManifestHash,
		ApprovalRef:              flags.approvalRef,
		WorkloadIdentity:         flags.workloadIdentity,
		MutationClass:            clusterrecovery.MutationClass(flags.mutationClass),
		Phase:                    clusterrecovery.ApplyPhase(flags.phase),
		Offset:                   flags.offset,
		Limit:                    flags.limit,
		PriorReceipt:             priorReceipt,
		ExpectedPriorReceiptHash: flags.expectedPriorReceiptHash,
		PriorVerificationRef:     flags.priorVerificationRef,
	}
	if err := clusterrecovery.ValidateApplyRequest(manifest, applyOptions, time.Now()); err != nil {
		return fail("apply_guard", err)
	}
	if flags.rollbackPath == flags.receiptPath {
		return fail("artifact_preflight_failed", errors.New("receipt and rollback paths must differ"))
	}
	for _, path := range []string{flags.rollbackPath, flags.receiptPath} {
		if err := validateNewArtifactPath(path); err != nil {
			return fail("artifact_preflight_failed", err)
		}
	}
	receiptArtifact, err := reserveJSONArtifact(flags.receiptPath)
	if err != nil {
		return fail("artifact_preflight_failed", err)
	}
	runtime, err := loader(flags.configPath)
	if err != nil {
		receiptArtifact.abort()
		return fail("runtime_initialization_failed", err)
	}
	if err := validateRuntime(runtime, true); err != nil {
		receiptArtifact.abort()
		return fail("environment_guard", err)
	}
	receipt, executeErr := clusterrecovery.NewExecutor(runtime.source, runtime.writeStore).Execute(
		ctx,
		manifest,
		applyOptions,
		clusterrecovery.FileRollbackSink{Path: flags.rollbackPath},
	)
	if receipt != nil {
		if err := receiptArtifact.commit(receipt); err != nil {
			return fail("receipt_write_failed", errors.Join(executeErr, err))
		}
		if err := json.NewEncoder(stdout).Encode(summarizeReceipt(receipt)); err != nil {
			return fail("summary_write_failed", errors.Join(executeErr, err))
		}
	}
	if executeErr != nil {
		if receipt == nil {
			receiptArtifact.abort()
		}
		return fail("execution_stopped", executeErr)
	}
	if receipt == nil {
		receiptArtifact.abort()
		return fail("execution_receipt_missing", errors.New("executor returned no receipt"))
	}
	return nil
}

type applySummary struct {
	Mode        string `json:"mode"`
	RunID       string `json:"run_id"`
	ReceiptHash string `json:"receipt_hash"`
	Stopped     bool   `json:"stopped"`
	Applied     int    `json:"applied"`
	Noop        int    `json:"noop"`
	Failed      int    `json:"failed"`
}

func summarizeReceipt(receipt *clusterrecovery.Receipt) applySummary {
	summary := applySummary{Mode: "apply", RunID: receipt.RunID, ReceiptHash: receipt.ReceiptHash, Stopped: receipt.Stopped}
	for _, result := range receipt.Results {
		switch result.Outcome {
		case clusterrecovery.MutationOutcomeApplied:
			summary.Applied++
		case clusterrecovery.MutationOutcomeNoop:
			summary.Noop++
		case clusterrecovery.MutationOutcomeFailed:
			summary.Failed++
		}
	}
	return summary
}

func loadRuntime(configPath string) (*commandRuntime, error) {
	cfg, err := config.InitConfig(configPath)
	if err != nil {
		return nil, err
	}
	if cfg == nil || cfg.Env != clusterrecovery.StagingEnvironment {
		return nil, errors.New("configuration is not staging")
	}
	if cfg.DBConfig == nil {
		return nil, errors.New("database configuration is missing")
	}
	if err := db.Initialize(cfg.DBConfig, true); err != nil {
		return nil, err
	}
	if err := datasource.Initialize(cfg.DBaasConfig, nil); err != nil {
		return nil, err
	}
	store := clusterrecovery.NewGormStore(db.GetDB())
	return &commandRuntime{
		environment: cfg.Env,
		source:      clusterrecovery.NewCloudSource(),
		readStore:   store,
		writeStore:  store,
	}, nil
}

func validateRuntime(runtime *commandRuntime, requireWrite bool) error {
	if runtime == nil || runtime.environment != clusterrecovery.StagingEnvironment || runtime.source == nil || runtime.readStore == nil {
		return errors.New("runtime is not an initialized staging runtime")
	}
	if requireWrite && runtime.writeStore == nil {
		return errors.New("staging write store is unavailable")
	}
	return nil
}

func openPrivateInput(path string) (*os.File, error) {
	if !canonicalAbsolutePath(path) {
		return nil, errors.New("input path must be absolute")
	}
	resolved, err := filepath.EvalSymlinks(path)
	if err != nil || resolved != path {
		return nil, errors.New("input path must not contain symlinks")
	}
	info, err := os.Lstat(path)
	if err != nil {
		return nil, err
	}
	if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
		return nil, errors.New("input must be a regular non-symlink file")
	}
	if info.Mode().Perm()&0o077 != 0 {
		return nil, errors.New("input file permissions must not allow group or other access")
	}
	return os.Open(path)
}

func validateOutputDirectory(path string) error {
	if !canonicalAbsolutePath(path) {
		return errors.New("output directory must be absolute")
	}
	resolved, err := filepath.EvalSymlinks(path)
	if err != nil || resolved != path {
		return errors.New("output directory path must not contain symlinks")
	}
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}
	if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
		return errors.New("output directory must be a non-symlink directory")
	}
	if info.Mode().Perm()&0o022 != 0 {
		return errors.New("output directory must not be group or world writable")
	}
	return nil
}

func validateNewArtifactPath(path string) error {
	if !canonicalAbsolutePath(path) || filepath.Base(path) == "." || filepath.Base(path) == string(filepath.Separator) {
		return errors.New("artifact path must be an absolute file path")
	}
	if err := validateOutputDirectory(filepath.Dir(path)); err != nil {
		return err
	}
	if _, err := os.Lstat(path); err == nil {
		return errors.New("artifact path already exists")
	} else if !errors.Is(err, os.ErrNotExist) {
		return err
	}
	return nil
}

func canonicalAbsolutePath(path string) bool {
	return filepath.IsAbs(path) && filepath.Clean(path) == path
}

type reservedJSONArtifact struct {
	path string
	file *os.File
}

func reserveJSONArtifact(path string) (*reservedJSONArtifact, error) {
	if err := validateNewArtifactPath(path); err != nil {
		return nil, err
	}
	file, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
	if err != nil {
		return nil, err
	}
	if err := file.Chmod(0o600); err != nil {
		_ = file.Close()
		_ = os.Remove(path)
		return nil, err
	}
	if err := file.Sync(); err != nil {
		_ = file.Close()
		_ = os.Remove(path)
		return nil, err
	}
	if err := syncDirectory(filepath.Dir(path)); err != nil {
		_ = file.Close()
		_ = os.Remove(path)
		return nil, err
	}
	return &reservedJSONArtifact{path: path, file: file}, nil
}

func syncDirectory(path string) error {
	directory, err := os.Open(path)
	if err != nil {
		return err
	}
	if err := directory.Sync(); err != nil {
		_ = directory.Close()
		return err
	}
	return directory.Close()
}

func (a *reservedJSONArtifact) commit(value any) error {
	if a == nil || a.file == nil {
		return errors.New("receipt artifact is not reserved")
	}
	payload, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		_ = a.file.Close()
		a.file = nil
		return err
	}
	payload = append(payload, '\n')
	written, err := a.file.Write(payload)
	if err == nil && written != len(payload) {
		err = io.ErrShortWrite
	}
	if err != nil {
		_ = a.file.Close()
		a.file = nil
		return err
	}
	if err := a.file.Sync(); err != nil {
		_ = a.file.Close()
		a.file = nil
		return err
	}
	if err := a.file.Close(); err != nil {
		a.file = nil
		return err
	}
	a.file = nil
	return nil
}

func (a *reservedJSONArtifact) abort() {
	if a == nil || a.file == nil {
		return
	}
	_ = a.file.Close()
	a.file = nil
	_ = os.Remove(a.path)
}

func planArtifactPaths(outputDir, runID string) []string {
	return []string{
		filepath.Join(outputDir, runID+"-report.json"),
		filepath.Join(outputDir, runID+"-create-manifest.json"),
		filepath.Join(outputDir, runID+"-update-manifest.json"),
	}
}