package migration

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/tidbcloud/observability/vm-migration/internal/exitcode"
	"github.com/tidbcloud/observability/vm-migration/internal/printer"
	"github.com/tidbcloud/observability/vm-migration/internal/schema"
	"github.com/tidbcloud/observability/vm-migration/internal/status"
	"github.com/tidbcloud/observability/vm-migration/internal/subprocess"
	"github.com/tidbcloud/observability/vm-migration/internal/vmquery"
)

const vmctlExecutable = "vmctl"

// bytesRe parses vmctl's "total bytes: X.Y UNIT" stats line (case-insensitive).
// vmctl uses 1000-base SI units.
var bytesRe = regexp.MustCompile(`(?i)total bytes:\s*([\d.]+)\s*([kKMGTPE]?B)`)

// unitToBytes maps vmctl's SI unit suffix to multiplier.
var unitToBytes = map[string]int64{
	"B":  1,
	"KB": 1000,
	"MB": 1000 * 1000,
	"GB": 1000 * 1000 * 1000,
	"TB": 1000 * 1000 * 1000 * 1000,
	"PB": 1000 * 1000 * 1000 * 1000 * 1000,
	"EB": 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
}

// VmctlOptions holds one tenant migration's execution options. VmRateLimit
// and VmConcurrency are *int (nil ⇒ user did not pass the flag) to
// distinguish "0" from "unset". TunePath points at the per-migration tune
// document; its vmctl section's non-nil knobs override
// VmRateLimit/VmConcurrency at every shard boundary.
type VmctlOptions struct {
	Manifest          string
	VmRateLimit       *int
	VmConcurrency     *int
	VmctlStepInterval string
	TunePath          string
}

// RunVmctl is the testable entry point. vmctlRunner is injected so tests
// don't need a real vmctl binary.
func RunVmctl(ctx context.Context, a VmctlOptions, runner VmctlRunner, output ActionOutput) error {
	output = output.normalized()
	manifestPath := a.Manifest
	statusPath := status.PathForManifest(manifestPath)

	manifest, err := status.LoadManifest(manifestPath)
	if err != nil {
		return err
	}
	manifestByShard := indexManifest(manifest)

	store, err := openMigrateStore(manifestPath, statusPath)
	if err != nil {
		return err
	}

	var selectedIDs []string
	store.View(func(doc *schema.Document) {
		for i := range doc.Shards {
			if doc.Shards[i].Phase != schema.PhaseImported {
				selectedIDs = append(selectedIDs, doc.Shards[i].ShardID)
			}
		}
	})

	store.View(func(doc *schema.Document) {
		printer.Printf(output.Progress, "migration start: total=%d imported=%d failed=%d pending=%d selected=%d",
			doc.Summary.Total, doc.Summary.Imported, doc.Summary.Failed, doc.Summary.Pending, len(selectedIDs))
	})

	startedAt := time.Now()
	for i, shardID := range selectedIDs {
		entry, ok := manifestByShard[shardID]
		if !ok {
			if err := store.Mutate(ctx, func(doc *schema.Document) error {
				return failShardNotFound(doc, shardID)
			}); err != nil {
				return err
			}
			printer.Printf(output.Progress, "[%d/%d] failed migration shard=%s error=shardID not found in manifest",
				i+1, len(selectedIDs), shardID)
			continue
		}
		if err := runVmctlForShardStore(ctx, a, runner, store, shardID, entry.tenant, entry.shard, output, i+1, len(selectedIDs)); err != nil {
			return err
		}
	}

	totalDuration := time.Since(startedAt)
	selected := make(map[string]struct{}, len(selectedIDs))
	for _, shardID := range selectedIDs {
		selected[shardID] = struct{}{}
	}
	failedSelected := 0
	store.View(func(doc *schema.Document) {
		for i := range doc.Shards {
			if _, ok := selected[doc.Shards[i].ShardID]; ok && doc.Shards[i].Phase != schema.PhaseImported {
				failedSelected++
			}
		}
		printer.Printf(output.Progress, "migration done: total=%d imported=%d verified=%d failed=%d pending=%d bytes=%s duration=%.1fs",
			doc.Summary.Total, doc.Summary.Imported, doc.Summary.Verified,
			doc.Summary.Failed, doc.Summary.Pending,
			printer.FormatBytes(float64(doc.Summary.TotalBytes)),
			totalDuration.Seconds())
	})
	if failedSelected > 0 {
		return exitcode.PartialFailure(fmt.Errorf("%d shard(s) failed during this run", failedSelected))
	}
	return nil
}

// openMigrateStore opens the persisted shard-status store, seeding it from
// the manifest on first run.
func openMigrateStore(manifestPath, statusPath string) (*status.Store, error) {
	return status.Open(statusPath, func() (*schema.Document, error) {
		manifest, err := status.LoadManifest(manifestPath)
		if err != nil {
			return nil, err
		}
		return status.BuildInitialStatus(manifest, manifestPath), nil
	})
}

// failShardNotFound flips a shard to failed with a generic error.
func failShardNotFound(doc *schema.Document, shardID string) error {
	s := status.FindShard(doc, shardID)
	if s == nil {
		return fmt.Errorf("shard %s: not present in status document", shardID)
	}
	if err := s.Phase.Transition(schema.PhaseFailed); err != nil {
		return fmt.Errorf("shard %s: %w", shardID, err)
	}
	s.Phase = schema.PhaseFailed
	s.LastError = new("shardID not found in manifest")
	s.FinishedAt = new(status.UTCNow())
	return nil
}

// VmctlRunner abstracts "run vmctl with these args, stream output to w".
// Production uses defaultVmctlRunner (exec vmctl); tests inject a fake.
type VmctlRunner interface {
	Run(ctx context.Context, args []string, out io.Writer) (exitCode int, bytesMigrated int64)
}

type defaultVmctlRunner struct{}

// NewDefaultVmctlRunner returns the production VmctlRunner (exec vmctl).
func NewDefaultVmctlRunner() VmctlRunner { return &defaultVmctlRunner{} }

func (r *defaultVmctlRunner) Run(ctx context.Context, args []string, out io.Writer) (int, int64) {
	sniffer := &byteSnifferWriter{inner: out}
	res := subprocess.Exec(subprocess.ExecOptions{
		Ctx:    ctx,
		Bin:    args[0],
		Args:   args[1:],
		Stdout: sniffer,
		Stderr: sniffer,
	})
	if res.Err == nil {
		return 0, sniffer.bytes()
	}
	var se *subprocess.Error
	if errors.As(res.Err, &se) {
		return se.Code, sniffer.bytes()
	}
	slog.Error("vmctl exec", "cmd", strings.Join(args, " "), "err", res.Err)
	return 127, sniffer.bytes()
}

// byteSnifferWriter forwards writes to inner while scanning for vmctl's
// "total bytes: X.Y UNIT" stats line.
type byteSnifferWriter struct {
	inner    io.Writer
	leftover []byte
	found    int64
}

func (w *byteSnifferWriter) Write(p []byte) (int, error) {
	if w.inner != nil {
		if _, err := w.inner.Write(p); err != nil {
			return 0, err
		}
	}
	w.leftover = append(w.leftover, p...)
	for {
		i := bytes.IndexByte(w.leftover, '\n')
		if i < 0 {
			break
		}
		line := w.leftover[:i]
		w.leftover = w.leftover[i+1:]
		if v := parseVmctlBytesLine(string(line)); v > 0 {
			w.found = v
		}
	}
	return len(p), nil
}

func (w *byteSnifferWriter) bytes() int64 { return w.found }

// parseVmctlBytesLine extracts the byte count from a vmctl "total bytes:" line.
// Returns 0 if the line doesn't match.
func parseVmctlBytesLine(line string) int64 {
	m := bytesRe.FindStringSubmatch(line)
	if m == nil {
		return 0
	}
	// Normalize the unit to canonical case (KB not kb/kB).
	unit := strings.ToUpper(m[2])
	mult, ok := unitToBytes[unit]
	if !ok {
		return 0
	}
	f, err := strconv.ParseFloat(m[1], 64)
	if err != nil {
		return 0
	}
	return int64(f * float64(mult))
}

// vmctlCommand builds the vmctl vm-native command for one shard.
func vmctlCommand(tenant schema.ManifestTenant, shard schema.ManifestShard, stepInterval string, vmRateLimit, vmConcurrency *int) []string {
	cmd := []string{
		vmctlExecutable,
		"vm-native",
		"--vm-native-src-addr", tenant.SourceAddr,
		"--vm-native-dst-addr", tenant.DestinationAddr,
		"--vm-native-filter-match", vmquery.MatchAll,
		"--vm-native-filter-time-start", shard.TimeStart,
		"--vm-native-filter-time-end", shard.TimeEnd,
		"--vm-native-step-interval", stepInterval,
		"--disable-progress-bar",
		"-s",
	}
	if vmRateLimit != nil {
		cmd = append(cmd, "--vm-rate-limit", strconv.Itoa(*vmRateLimit))
	}
	if vmConcurrency != nil {
		cmd = append(cmd, "--vm-concurrency", strconv.Itoa(*vmConcurrency))
	}
	return cmd
}

// runVmctlForShardStore drives one shard through running → (imported | failed)
// using Store.Mutate for atomic persistence. Returns nil on success or a
// wrapped error from any Mutate call; a non-zero vmctl exit is recorded on the
// shard (migration continues with the next shard).
//
// State changes (each Mutate is atomic):
//  1. flip → running, reset per-attempt fields, persist
//  2. run vmctl (NO lock — subprocess execution can take hours)
//  3. flip → imported | failed, set Bytes / LastError, persist
//
// The shard is "checked out" by ID; concurrent callers operating on different
// shard IDs do not conflict. Re-entering with the same shard ID while running
// correctly fails the second Mutate's transition check.
func runVmctlForShardStore(
	ctx context.Context,
	a VmctlOptions,
	runner VmctlRunner,
	store *status.Store,
	shardID string,
	tenant schema.ManifestTenant,
	shard schema.ManifestShard,
	output ActionOutput,
	current, total int,
) error {
	position := ""
	if current > 0 && total > 0 {
		position = fmt.Sprintf("[%d/%d] ", current, total)
	}
	tune, err := EffectiveTune(a.TunePath, a.VmRateLimit, a.VmConcurrency)
	if err != nil {
		return fmt.Errorf("%sload tune: %w", position, err)
	}
	cmdArgs := vmctlCommand(tenant, shard, a.VmctlStepInterval, tune.RateLimitBytes, tune.Concurrency)

	var label string
	var nextAttempt int
	store.View(func(doc *schema.Document) {
		if s := status.FindShard(doc, shardID); s != nil {
			label = shardLabel(*s)
			nextAttempt = s.Attempts + 1
		}
	})
	printer.Printf(output.Progress, "%sstart migration %s attempt=%d", position, label, nextAttempt)
	printer.Printf(output.Log, "%svmctl command: %s", position, strings.Join(cmdArgs, " "))

	startedAt := time.Now()
	// Mutate #1: flip pending | failed → running.
	if err := store.Mutate(ctx, func(doc *schema.Document) error {
		return mutateShardRunning(doc, shardID, nextAttempt)
	}); err != nil {
		return fmt.Errorf("%s%w", position, err)
	}

	exitCode, bytesMigrated := runner.Run(ctx, cmdArgs, output.Log)

	duration := time.Since(startedAt)
	// Mutate #2: flip running → imported | failed.
	var resultLabel string
	if err := store.Mutate(ctx, func(doc *schema.Document) error {
		var err error
		resultLabel, err = mutateShardFinal(doc, shardID, exitCode, bytesMigrated)
		return err
	}); err != nil {
		return fmt.Errorf("%s%w", position, err)
	}
	if exitCode == 0 {
		printer.Printf(output.Progress, "%sdone migration %s duration=%.1fs bytes=%s",
			position, resultLabel, duration.Seconds(),
			printer.FormatBytes(float64(bytesMigrated)))
	} else {
		printer.Printf(output.Progress, "%sfailed migration %s duration=%.1fs error=vmctl exited with code %d",
			position, resultLabel, duration.Seconds(), exitCode)
	}
	return nil
}

// mutateShardRunning flips a shard to PhaseRunning and resets per-attempt fields.
// Returns a wrapped error if the shard is missing or the transition is illegal.
func mutateShardRunning(doc *schema.Document, shardID string, nextAttempt int) error {
	s := status.FindShard(doc, shardID)
	if s == nil {
		return fmt.Errorf("shard %s: not present in status document", shardID)
	}
	if err := s.Phase.Transition(schema.PhaseRunning); err != nil {
		return fmt.Errorf("shard %s: %w", shardID, err)
	}
	s.Phase = schema.PhaseRunning
	s.Verified = false
	s.Attempts = nextAttempt
	s.Bytes = nil
	s.StartedAt = new(status.UTCNow())
	s.FinishedAt = nil
	s.LastError = nil
	s.VerifiedAt = nil
	s.Verification = nil
	return nil
}

// mutateShardFinal flips a running shard to imported (exitCode == 0) or failed
// (otherwise), recording Bytes / LastError / FinishedAt. Returns the shard's
// post-mutation label string for stdout.
func mutateShardFinal(doc *schema.Document, shardID string, exitCode int, bytesMigrated int64) (string, error) {
	s := status.FindShard(doc, shardID)
	if s == nil {
		return "", fmt.Errorf("shard %s: not present in status document", shardID)
	}
	target := schema.PhaseImported
	if exitCode != 0 {
		target = schema.PhaseFailed
	}
	if err := s.Phase.Transition(target); err != nil {
		return "", fmt.Errorf("shard %s: %w", shardID, err)
	}
	// Mutations land only after the transition validates, so an illegal
	// transition leaves no stray fields in the in-memory document.
	s.Phase = target
	s.FinishedAt = new(status.UTCNow())
	if exitCode == 0 {
		s.Bytes = new(bytesMigrated)
		s.LastError = nil
	} else {
		s.LastError = new(fmt.Sprintf("vmctl exited with code %d", exitCode))
	}
	return shardLabel(*s), nil
}

func shardLabel(s schema.Shard) string {
	return strings.Join([]string{
		"shard=" + s.ShardID,
		"tenant=" + s.Tenant,
		fmt.Sprintf("time=%s..%s", s.TimeStart, s.TimeEnd),
	}, " ")
}

func indexManifest(m *schema.Manifest) map[string]vmctlShardLookup {
	out := map[string]vmctlShardLookup{}
	for _, tenant := range m.Tenants {
		for _, shard := range tenant.Shards {
			out[shard.ShardID] = vmctlShardLookup{tenant: tenant, shard: shard}
		}
	}
	return out
}

type vmctlShardLookup struct {
	tenant schema.ManifestTenant
	shard  schema.ManifestShard
}