package server

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"slices"
	"sort"
	"sync"
	"time"

	"github.com/tidbcloud/observability/vm-migration/internal/jsonio"
	"github.com/tidbcloud/observability/vm-migration/internal/migration"
	"github.com/tidbcloud/observability/vm-migration/internal/printer"
	"github.com/tidbcloud/observability/vm-migration/internal/schema"
	"github.com/tidbcloud/observability/vm-migration/internal/status"
)

// errPlanMissingVmselect is returned by StartExportProbe when the plan
// predates the vmselect block.
var errPlanMissingVmselect = errors.New("plan missing vmselect block; recompile required")

// errPlanMissingRuntime is returned when a plan predates typed action runtime
// configuration.
var errPlanMissingRuntime = errors.New("plan missing runtime block; rerun prepare")

// errManualCancel is the cancellation cause set by Cancel; it distinguishes
// an operator stop from root-context shutdown inside run().
var errManualCancel = errors.New("manual cancel")

// Migration is one migration act: an atomic directory of artifacts and a
// single-flight slot for running jobs against it.
//
// Migration is constructed by newMigrationFromDir (called from
// LoadFromDisk). Migration is not safe for concurrent use across goroutines
// without going through the exported methods (which take the mutex).
type Migration struct {
	name string
	dir  string
	plan *schema.MigrationLoopPlan

	rootCtx context.Context

	mu                sync.Mutex
	job               job
	jobRecord         *jobRecord
	jobCancel         context.CancelCauseFunc
	lastStartError    string
	lastStartErrorAt  string
	lastJobRecord     *jobRecord
	logs              *logHub
	cachedLoopStatus  *schema.MigrationLoopStatus
	cachedStatusMtime int64
	cachedJobs        []jobRecord
	cachedJobsMtime   int64
	cachedMarkerMtime int64
	loopOptions       []migration.Option
	provisioner       migration.VmselectProvisioner

	// retryInterval overrides the default delay between retry attempts; zero
	// means the package default. Tests shrink it so retries fire promptly.
	retryInterval time.Duration
}

// Name returns the migration's identifier.
func (m *Migration) Name() string { return m.name }

// Dir returns the migration's private directory path.
func (m *Migration) Dir() string { return m.dir }

// migrationStatus is a snapshot of one Migration's runtime state.
type migrationStatus struct {
	Name          string                      `json:"name"`
	Dir           string                      `json:"dir"`
	Running       bool                        `json:"running"`
	Kind          jobKind                     `json:"kind,omitempty"`
	JobId         string                      `json:"jobId,omitempty"`
	LastJobId     string                      `json:"lastJobId,omitempty"`
	LastKind      jobKind                     `json:"lastKind,omitempty"`
	LastResult    string                      `json:"lastResult,omitempty"`
	ResumeError   string                      `json:"resumeError,omitempty"`
	ResumeErrorAt string                      `json:"resumeErrorAt,omitempty"`
	LoopStatus    *schema.MigrationLoopStatus `json:"loopStatus,omitempty"`
}

// loopFile is the prepared migration-loop.json inside the migration directory.
const loopFile = "migration-loop.json"

// inflightMarkerFile is the on-disk proof that a job is running: written at
// job start, removed after the job's record is appended to jobs.jsonl. A
// marker left behind (crash, SIGKILL, process shutdown) is resolved at the
// next serve start: migrate-run jobs are resumed; read-only (export-probe)
// and vmselect jobs are finalized as not-resumed.
const inflightMarkerFile = "jobs-inflight.json"

// loopPath returns the absolute path to migration-loop.json inside the
// migration directory.
func (m *Migration) loopPath() string { return filepath.Join(m.dir, loopFile) }

// inflightMarkerPath returns the per-migration in-flight marker path.
func (m *Migration) inflightMarkerPath() string {
	return filepath.Join(m.dir, inflightMarkerFile)
}

// Status snapshots the migration's runtime state plus the on-disk loop-status
// document.
//
// Status caches the parsed loop-status by mtime to avoid re-reading and
// re-parsing on every poll. Disk IO runs outside m.mu so a slow stat/read
// doesn't block job lifecycle (start, cancel, run-done).
func (m *Migration) Status() migrationStatus {
	m.mu.Lock()
	s := migrationStatus{
		Name:    m.name,
		Dir:     m.dir,
		Running: m.job != nil,
	}
	if m.job != nil {
		s.Kind = m.job.Kind()
	}
	if m.jobRecord != nil {
		s.JobId = m.jobRecord.ID
	}
	if m.lastJobRecord != nil {
		s.LastJobId = m.lastJobRecord.ID
		s.LastKind = m.lastJobRecord.Kind
		s.LastResult = m.lastJobRecord.Result
	}
	s.ResumeError = m.lastStartError
	s.ResumeErrorAt = m.lastStartErrorAt
	statusPath := migration.DeriveLoopStatusPath(m.loopPath())
	cachedMtime := m.cachedStatusMtime
	cached := m.cachedLoopStatus
	m.mu.Unlock()

	cached = m.refreshLoopStatus(statusPath, cachedMtime, cached)
	s.LoopStatus = cached
	return s
}

// refreshLoopStatus reads and parses the status file when its mtime has
// advanced since the last cache fill.
//
// refreshLoopStatus updates the cache under m.mu. Stat, read, or parse errors
// return the previous cache (or nil).
func (m *Migration) refreshLoopStatus(statusPath string, cachedMtime int64, cached *schema.MigrationLoopStatus) *schema.MigrationLoopStatus {
	fi, err := os.Stat(statusPath)
	if err != nil {
		return cached
	}
	mtime := fi.ModTime().UnixNano()
	if cached != nil && mtime == cachedMtime {
		return cached
	}
	doc, _ := loadLoopStatus(m.loopPath())
	if doc == nil {
		return cached
	}
	m.mu.Lock()
	m.cachedLoopStatus = doc
	m.cachedStatusMtime = mtime
	m.mu.Unlock()
	return doc
}

// validateRunnable enforces the plan preconditions every typed job shares.
//
// The vmselect and runtime blocks must be present. validateRunnable is
// centralized so the probe start and run paths stay in sync. The loop path
// enforces the same contract via migration.ValidateStep at load time.
func (m *Migration) validateRunnable() error {
	if m.plan == nil || m.plan.Vmselect == nil {
		return errPlanMissingVmselect
	}
	if m.plan.Runtime == nil {
		return errPlanMissingRuntime
	}
	return nil
}

// start is the shared single-flight entry point for all StartXxx methods.
//
// start writes the in-flight marker before touching any m.* field, allocates
// the jobRecord, spawns the run goroutine, and returns the record's ID.
func (m *Migration) start(j job) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.job != nil {
		return "", errAlreadyRunning
	}
	record := newJobRecord(j.Kind())
	if err := status.DumpJSONAtomic(m.inflightMarkerPath(), record); err != nil {
		return "", fmt.Errorf("write %s: %w", m.inflightMarkerPath(), err)
	}
	ctx, cancel := context.WithCancelCause(m.rootCtx)
	m.job = j
	m.jobRecord = record
	m.jobCancel = cancel
	m.lastStartError, m.lastStartErrorAt = "", ""
	m.logs.GetOrCreate() // ensure an open stream before the job produces output
	go m.run(ctx, j, record)
	return record.ID, nil
}

// run is the goroutine entry point.
//
// job.Run runs to completion (or ctx cancellation); the record is finalized,
// persisted to jobs.jsonl, and cached as lastJobRecord for Status and GetJob.
// Record finalization happens under m.mu so Status cannot observe partial
// fields. Artifact and history persistence finish before the immutable record
// is published as the last job and the single-flight slot is released. The job
// log stream closes last so /logs subscribers drain buffered output before the
// channel shuts.
func (m *Migration) run(ctx context.Context, j job, record *jobRecord) {
	err := j.Run(ctx)

	// Process shutdown (SIGTERM/SIGINT cancels the root context) killed the
	// job mid-run: leave the in-flight marker and jobs.jsonl untouched so
	// the next serve start resolves this job (migrate-run resumes; a raced-to-
	// completion job re-runs an idempotent pass; export-probe finalizes as
	// not-resumed). A job cancelled via Cancel() finalizes as a manual stop
	// instead even when shutdown lands during the unwind — an operator's
	// cancel must never be auto-resumed.
	if m.rootCtx.Err() != nil && context.Cause(ctx) != errManualCancel {
		m.mu.Lock()
		m.job = nil
		m.jobRecord = nil
		m.jobCancel = nil
		logs := m.logs
		m.mu.Unlock()
		logs.Close()
		slog.Info("job interrupted by shutdown; resumable on next start",
			"id", record.ID, "kind", record.Kind)
		return
	}

	m.mu.Lock()
	if source, ok := j.(warningSource); ok {
		record.Warnings = source.Warnings()
	}
	record.FinishedAt = status.FormatTime(status.UTCNow())
	if err != nil {
		record.Result = err.Error()
	}
	if ctx.Err() == context.Canceled {
		record.Cancelled = true
	}
	logs := m.logs
	m.mu.Unlock()

	if w, ok := j.(artifactWriter); ok && err == nil {
		path := filepath.Join(m.dir, "artifacts", record.ID+".json")
		if mkErr := os.MkdirAll(filepath.Dir(path), 0o755); mkErr != nil {
			m.addJobWarning(record, logs, fmt.Sprintf("job=%s artifact directory unavailable", record.ID))
			slog.Error("artifact mkdir", "id", record.ID, "err", mkErr)
		} else if writeErr := w.WriteArtifact(path, record); writeErr != nil {
			m.addJobWarning(record, logs, fmt.Sprintf("job=%s artifact unavailable", record.ID))
			slog.Error("artifact write", "id", record.ID, "err", writeErr)
		}
	}
	if appendErr := appendJobRecord(m.jobsJSONLPath(), record); appendErr != nil {
		m.addJobWarning(record, logs, fmt.Sprintf("job=%s history unavailable", record.ID))
		slog.Error("jobs.jsonl append", "id", record.ID, "err", appendErr)
	} else if rmErr := os.Remove(m.inflightMarkerPath()); rmErr != nil && !os.IsNotExist(rmErr) {
		// Keep-on-failure is the safe direction: a leftover marker with a
		// finished record is cleaned as stale on the next start.
		slog.Warn("remove in-flight marker", "id", record.ID, "err", rmErr)
	}

	m.mu.Lock()
	m.job = nil
	m.jobRecord = nil
	m.jobCancel = nil
	m.lastJobRecord = record
	m.mu.Unlock()

	logs.Close()
}

func (m *Migration) addJobWarning(record *jobRecord, logs *logHub, message string) {
	m.mu.Lock()
	record.Warnings = append(record.Warnings, message)
	m.mu.Unlock()
	printer.Printf(logs, "warning: %s", message)
}

// jobsJSONLPath returns the per-migration jobs.jsonl path.
func (m *Migration) jobsJSONLPath() string { return filepath.Join(m.dir, "jobs.jsonl") }

// readInflightMarker reads jobs-inflight.json.
//
// Returns (nil, nil) when no marker exists. Read and parse failures carry
// context; a corrupt marker is kept on disk for inspection by the caller's
// policy.
func (m *Migration) readInflightMarker() (*jobRecord, error) {
	var r jobRecord
	if err := jsonio.ReadFile(m.inflightMarkerPath(), &r); err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("read in-flight marker: %w", err)
	}
	return &r, nil
}

// Jobs returns the job history from jobs.jsonl plus the in-flight marker's
// record when one exists (job running, or interrupted before the next
// start), sorted by submission time (StartedAt ascending). A marker whose
// ID is already in the history is a stale leftover from the
// append-then-remove window and is not duplicated.
//
// Jobs is cached by both files' mtimes so repeated calls don't re-read.
func (m *Migration) Jobs() []jobRecord {
	m.mu.Lock()
	cachedHistMtime := m.cachedJobsMtime
	cachedMarkMtime := m.cachedMarkerMtime
	cached := m.cachedJobs
	histPath := m.jobsJSONLPath()
	markPath := m.inflightMarkerPath()
	m.mu.Unlock()

	var histMtime, markMtime int64
	if fi, err := os.Stat(histPath); err == nil {
		histMtime = fi.ModTime().UnixNano()
	}
	if fi, err := os.Stat(markPath); err == nil {
		markMtime = fi.ModTime().UnixNano()
	}
	if cached != nil && histMtime == cachedHistMtime && markMtime == cachedMarkMtime {
		return append([]jobRecord(nil), cached...)
	}

	got, err := readJobRecords(histPath)
	if err != nil {
		return append([]jobRecord(nil), cached...)
	}
	merged := got
	marker, err := m.readInflightMarker()
	switch {
	case err != nil:
		slog.Warn("jobs-inflight.json unreadable", "path", markPath, "err", err)
	case marker != nil && !slices.ContainsFunc(got, func(r jobRecord) bool { return r.ID == marker.ID }):
		merged = append(got, *marker)
	}
	// Submission order. StartedAt is parsed, not string-compared:
	// RFC3339Nano's variable fraction width breaks lexical ordering inside
	// one second ("…:00Z" sorts after "…:00.1Z"). IDs break exact ties;
	// unparseable StartedAt values fall back to ID order.
	starts := make([]time.Time, len(merged))
	parsed := make([]bool, len(merged))
	for i := range merged {
		t, err := time.Parse(time.RFC3339Nano, merged[i].StartedAt)
		starts[i], parsed[i] = t, err == nil
	}
	sort.SliceStable(merged, func(i, j int) bool {
		if parsed[i] && parsed[j] && !starts[i].Equal(starts[j]) {
			return starts[i].Before(starts[j])
		}
		return merged[i].ID < merged[j].ID
	})
	m.mu.Lock()
	m.cachedJobs = merged
	m.cachedJobsMtime = histMtime
	m.cachedMarkerMtime = markMtime
	m.mu.Unlock()
	return append([]jobRecord(nil), merged...)
}

// GetJob looks up a single jobRecord by id.
//
// GetJob returns ok=false if not found.
func (m *Migration) GetJob(id string) (jobRecord, bool) {
	for _, r := range m.Jobs() {
		if r.ID == id {
			return r, true
		}
	}
	return jobRecord{}, false
}

// JobArtifact reads artifacts/<id>.json.
//
// JobArtifact returns an os.IsNotExist-wrapped error when the artifact is
// absent (loop jobs, unknown IDs).
func (m *Migration) JobArtifact(id string) ([]byte, error) {
	path := filepath.Join(m.dir, "artifacts", id+".json")
	return os.ReadFile(path)
}

// compiledPlanFile is the compiled-plan snapshot Prepare commits; the deploy
// job and tune views derive the vmselect specs from it.
const compiledPlanFile = "compiled-plan.json"

// compiledPlan loads the compiled-plan.json snapshot taken at prepare.
//
// A snapshot from before the vmselects restructure decodes with a zero
// vmselects block; that is rejected here so /tune and deploy surface a
// recompile instruction instead of empty specs.
func (m *Migration) compiledPlan() (*schema.CompiledDoc, error) {
	var doc schema.CompiledDoc
	if err := jsonio.ReadFile(filepath.Join(m.dir, compiledPlanFile), &doc); err != nil {
		return nil, fmt.Errorf("load %s: %w", compiledPlanFile, err)
	}
	if doc.Vmselects.Image == "" || doc.Vmselects.Migration.Name == "" || doc.Vmselects.Candidate.Name == "" {
		return nil, fmt.Errorf("%s predates the vmselects restructure; recompile the plan and re-prepare the migration", compiledPlanFile)
	}
	// Re-check the replica-span invariant plan validation enforces: the
	// snapshot is operator-editable on disk, and a target >= source would
	// otherwise build a negative ordinal range downstream.
	if doc.TargetReplicaCount <= 0 || doc.TargetReplicaCount >= doc.SourceReplicaCount {
		return nil, fmt.Errorf("%s has an invalid replica span: targetReplicaCount=%d must be positive and below sourceReplicaCount=%d; recompile the plan",
			compiledPlanFile, doc.TargetReplicaCount, doc.SourceReplicaCount)
	}
	return &doc, nil
}

// tunePath returns the per-migration tune.json path.
func (m *Migration) tunePath() string {
	return filepath.Join(m.dir, migration.TuneFile)
}

// Tune returns the effective knobs per section: the plan values with
// tune.json overrides applied. The vmctl pair takes effect at the next shard
// boundary; the vmselect specs at the next deploy.
func (m *Migration) Tune() (schema.TuneView, error) {
	doc, err := migration.LoadTuneDoc(m.tunePath())
	if err != nil {
		return schema.TuneView{}, err
	}
	planRate, planConcurrency := migration.PlanKnobs(m.plan)
	vmctl := doc.Vmctl.MergedOver(planRate, planConcurrency)
	compiled, err := m.compiledPlan()
	if err != nil {
		return schema.TuneView{}, err
	}
	specs, err := migration.EffectiveVmselectSpecs(compiled, doc.Vmselect)
	if err != nil {
		return schema.TuneView{}, err
	}
	return schema.TuneView{
		Vmctl:    vmctl,
		Vmselect: vmselectTuneView(compiled, specs, doc.Vmselect.UpdatedAt),
	}, nil
}

func vmselectTuneView(compiled *schema.CompiledDoc, pair migration.VmselectSpecPair, updatedAt string) schema.VmselectTuneView {
	view := func(spec migration.VmselectSpec, url string) schema.TuneTargetView {
		return schema.TuneTargetView{
			Name:      spec.Name,
			URL:       url,
			Image:     spec.Image,
			Replicas:  spec.Replicas,
			Resources: spec.TuneQuantities(),
		}
	}
	return schema.VmselectTuneView{
		Migration: view(pair.Migration, compiled.Vmselects.Migration.URLPattern),
		Candidate: view(pair.Candidate, compiled.Vmselects.Candidate.URLPattern),
		UpdatedAt: updatedAt,
	}
}

// errInvalidTune marks a tune rejected by section validation; the HTTP layer
// maps it to 422.
var errInvalidTune = errors.New("invalid tune")

// SetTune applies a sectioned replacement.
//
// A present section replaces its override wholesale (validated and
// stamped); an empty section removes its override; an absent section is
// untouched. An all-empty result removes tune.json. vmctl applies at the
// next shard boundary, never to the in-flight shard; vmselect applies on
// the next deploy.
func (m *Migration) SetTune(req schema.TuneSetRequest) (schema.TuneView, error) {
	if req.Vmctl != nil {
		section := req.Vmctl.Section()
		if err := section.Validate(); err != nil {
			return schema.TuneView{}, fmt.Errorf("%w: vmctl: %w", errInvalidTune, err)
		}
	}
	if req.Vmselect != nil {
		section := req.Vmselect.Section()
		if err := section.Validate(); err != nil {
			return schema.TuneView{}, fmt.Errorf("%w: %w", errInvalidTune, err)
		}
	}
	err := migration.UpdateTuneDoc(m.tunePath(), func(doc *schema.TuneDoc) {
		now := status.FormatTime(status.UTCNow())
		if req.Vmctl != nil {
			doc.Vmctl = req.Vmctl.Section()
			if !doc.Vmctl.IsEmpty() {
				doc.Vmctl.UpdatedAt = now
			}
		}
		if req.Vmselect != nil {
			doc.Vmselect = req.Vmselect.Section()
			if !doc.Vmselect.IsEmpty() {
				doc.Vmselect.UpdatedAt = now
			}
		}
	})
	if err != nil {
		return schema.TuneView{}, err
	}
	return m.Tune()
}

// Cancel aborts the in-flight job's context.
//
// Cancel returns errNotRunning if no job is running. id must match the
// in-flight job's ID; anything else returns errJobIdMismatch, which protects
// against racing a cancel against a finished job whose ID has been reused.
func (m *Migration) Cancel(id string) error {
	m.mu.Lock()
	cancel := m.jobCancel
	currentID := ""
	if m.jobRecord != nil {
		currentID = m.jobRecord.ID
	}
	m.mu.Unlock()
	if cancel == nil {
		return errNotRunning
	}
	if id != currentID {
		return errJobIdMismatch
	}
	cancel(errManualCancel)
	return nil
}

// NoteStartFailure records why the most recent automated start attempt
// failed, surfaced as migrationStatus.ResumeError for polling monitors.
//
// Failures of manual StartXxx calls are visible in the HTTP response and
// need no recording; this is for the startup resume path, which has no
// caller to return to. The record is cleared by the next successful start.
func (m *Migration) NoteStartFailure(err error) {
	m.mu.Lock()
	m.lastStartError = err.Error()
	m.lastStartErrorAt = status.FormatTime(status.UTCNow())
	m.mu.Unlock()
}