package server

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"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"
	"github.com/tidbcloud/observability/vm-migration/internal/vmquery"
)

// probeJob executes the export probe over every tenant in the plan.
//
// For each tenant, probeJob runs ExportShard over the manifest's shards and
// reports bytes/time. The probe never mutates shard progress. If the manifest
// is missing, the typed manifest action runs in-process and initializes its
// matching status file.
type probeJob struct {
	m        *Migration
	client   *http.Client
	artifact *probeArtifact
	srcURL   string
	progress io.Writer
}

func (j *probeJob) Kind() jobKind { return jobExportProbe }

// Run executes the export probe over every tenant in the plan.
//
// Run validates the plan's vmselect and runtime blocks. Run returns an error
// if those blocks are missing, if a manifest can't be produced for a tenant,
// or if ctx is cancelled. Per-shard export errors are recorded in the artifact
// but don't fail the job.
func (j *probeJob) Run(ctx context.Context) error {
	if err := j.m.validateRunnable(); err != nil {
		return err
	}
	if j.srcURL == "" {
		j.srcURL = j.m.plan.Vmselect.Migration
	}

	j.artifact = &probeArtifact{
		Kind:    string(jobExportProbe),
		Tenants: []probeTenantResult{},
	}

	for _, t := range j.m.plan.Tenants {
		if err := ctx.Err(); err != nil {
			return err
		}
		printer.Printf(j.progressWriter(), "tenant=%s probe status=running", t.Name)
		tenant, err := j.runTenant(ctx, t)
		if err != nil {
			return err
		}
		j.artifact.Tenants = append(j.artifact.Tenants, tenant)
		printer.Printf(j.progressWriter(), "tenant=%s probe status=completed shards=%d", t.Name, len(tenant.Shards))
	}
	j.artifact.computeTotals()
	return nil
}

// runTenant ensures the manifest exists, then exports each shard in the manifest.
func (j *probeJob) runTenant(ctx context.Context, t schema.MigrationLoopTenant) (probeTenantResult, error) {
	action, err := manifestAction(t)
	if err != nil {
		return probeTenantResult{}, err
	}
	manifestPath := migration.Resolve(j.m.dir, action.Manifest)
	if _, err := os.Stat(manifestPath); err != nil {
		if !os.IsNotExist(err) {
			return probeTenantResult{}, fmt.Errorf("stat manifest for %s: %w", t.Name, err)
		}
		if err := j.generateManifest(ctx, action); err != nil {
			return probeTenantResult{}, fmt.Errorf("generate manifest for %s: %w", t.Name, err)
		}
	}

	manifest, err := status.LoadManifest(manifestPath)
	if err != nil {
		return probeTenantResult{}, fmt.Errorf("load manifest for %s: %w", t.Name, err)
	}
	shards := status.IterShards(manifest)

	result := probeTenantResult{Tenant: t.Name}
	result.Shards = make([]probeShardRow, 0, len(shards))
	for _, s := range shards {
		if err := ctx.Err(); err != nil {
			return result, err
		}
		res := vmquery.ExportShard(ctx, j.client, s, vmquery.ExportOptions{
			SourceURLTemplate: j.srcURL,
			Match:             vmquery.MatchAll,
			Timeout:           60 * time.Second,
		})
		row := probeShardRow{
			ShardID: res.ShardID,
			Window:  s.TimeStart + " - " + s.TimeEnd,
			Bytes:   res.Bytes,
			DurMs:   res.DurMs,
			Error:   res.Error,
		}
		if res.DurMs > 0 {
			row.ThroughputBps = res.Bytes * 1000 / res.DurMs
		}
		result.Shards = append(result.Shards, row)
		if row.Error != "" {
			printer.Printf(j.progressWriter(), "tenant=%s shard=%s probe status=failed error=%s", t.Name, row.ShardID, row.Error)
		} else {
			printer.Printf(j.progressWriter(), "tenant=%s shard=%s probe status=completed bytes=%s dur=%s",
				t.Name, row.ShardID, printer.FormatBytes(float64(row.Bytes)), formatShardDur(row.DurMs))
		}
	}
	result.computeTotals()
	return result, nil
}

// manifestAction returns the tenant's canonical typed manifest step.
func manifestAction(t schema.MigrationLoopTenant) (*schema.MigrationLoopAction, error) {
	for _, c := range t.Cmds {
		if c.Step == "manifest" && c.Action != nil && c.Action.Kind == schema.ActionManifest {
			return c.Action, nil
		}
	}
	return nil, fmt.Errorf("plan has no manifest action for tenant %s; cannot generate", t.Name)
}

// generateManifest executes the manifest action in-process.
func (j *probeJob) generateManifest(ctx context.Context, action *schema.MigrationLoopAction) error {
	return migration.RunAction(ctx, j.m.plan, j.m.dir, action, nil, migration.ActionOutput{
		Progress: j.progressWriter(),
		Log:      j.progressWriter(),
	})
}

func (j *probeJob) progressWriter() io.Writer {
	if j.progress == nil {
		return io.Discard
	}
	return j.progress
}

// formatShardDur renders a shard duration: sub-second in ms, then
// fractional seconds.
func formatShardDur(ms int64) string {
	if ms < 1000 {
		return fmt.Sprintf("%dms", ms)
	}
	return fmt.Sprintf("%.1fs", float64(ms)/1000)
}

func (j *probeJob) WriteArtifact(path string, rec *jobRecord) error {
	j.artifact.ID = rec.ID
	j.artifact.StartedAt = rec.StartedAt
	j.artifact.FinishedAt = rec.FinishedAt
	return jsonio.WriteFile(path, j.artifact)
}

// StartExportProbe starts an export probe job for the migration.
//
// StartExportProbe returns the dispatched job's ID and errAlreadyRunning if
// a job is already in flight. The request body is empty; defaults come from
// the plan (vmselect.migration URL, match-all query, 60s timeout, no
// rate-limit).
func (m *Migration) StartExportProbe() (string, error) {
	if err := m.validateRunnable(); err != nil {
		return "", err
	}
	job := &probeJob{
		m:        m,
		client:   vmquery.NewSharedClient(),
		progress: m.logs,
	}
	return m.start(job)
}

// probeArtifact is the JSON artifact written to artifacts/<id>.json after
// an export probe job completes.
type probeArtifact struct {
	ID         string              `json:"id"`
	Kind       string              `json:"kind"`
	StartedAt  string              `json:"startedAt"`
	FinishedAt string              `json:"finishedAt"`
	Tenants    []probeTenantResult `json:"tenants"`
	Totals     probeTotals         `json:"totals"`
}

type probeTenantResult struct {
	Tenant string          `json:"tenant"`
	Shards []probeShardRow `json:"shards"`
	Totals probeTotals     `json:"totals"`
}

type probeShardRow struct {
	ShardID       string `json:"shardID"`
	Window        string `json:"window"`
	Bytes         int64  `json:"bytes"`
	DurMs         int64  `json:"durMs"`
	ThroughputBps int64  `json:"throughputBps"`
	Error         string `json:"error,omitempty"`
}

type probeTotals struct {
	Bytes  int64 `json:"bytes"`
	DurMs  int64 `json:"durMs"`
	AvgBps int64 `json:"avgBps"`
}

func (a *probeArtifact) computeTotals() {
	var t probeTotals
	for i := range a.Tenants {
		a.Tenants[i].computeTotals()
		t.Bytes += a.Tenants[i].Totals.Bytes
		t.DurMs += a.Tenants[i].Totals.DurMs
	}
	if t.DurMs > 0 {
		t.AvgBps = t.Bytes * 1000 / t.DurMs
	}
	a.Totals = t
}

func (r *probeTenantResult) computeTotals() {
	var t probeTotals
	for _, s := range r.Shards {
		t.Bytes += s.Bytes
		t.DurMs += s.DurMs
	}
	if t.DurMs > 0 {
		t.AvgBps = t.Bytes * 1000 / t.DurMs
	}
	r.Totals = t
}

// probeErrWidth caps the error column; longer messages are truncated.
const probeErrWidth = 48

// renderProbeText writes a as a fixed-width table: one row per shard, a
// totals row at the end. The error column appears only when some shard
// failed.
func renderProbeText(w io.Writer, a *probeArtifact) {
	shards := 0
	hasErr := false
	for _, t := range a.Tenants {
		shards += len(t.Shards)
		for _, s := range t.Shards {
			hasErr = hasErr || s.Error != ""
		}
	}
	if shards == 0 {
		printer.Printf(w, "probe job=%s: no shards probed", a.ID)
		return
	}

	cols := probeTableColumns(hasErr)
	rows := [][]string{}
	for _, t := range a.Tenants {
		for _, s := range t.Shards {
			row := []string{
				t.Tenant,
				s.ShardID,
				s.Window,
				printer.FormatBytes(float64(s.Bytes)),
				formatShardDur(s.DurMs),
				formatBps(s.ThroughputBps),
			}
			if hasErr {
				row = append(row, truncateErr(s.Error))
			}
			rows = append(rows, row)
		}
	}
	totals := []string{
		"total",
		strconv.Itoa(shards) + " shards",
		"-",
		printer.FormatBytes(float64(a.Totals.Bytes)),
		formatShardDur(a.Totals.DurMs),
		formatBps(a.Totals.AvgBps),
	}
	if hasErr {
		totals = append(totals, "")
	}

	printer.Printf(w, "job=%s finished=%s", a.ID, a.FinishedAt)
	widths := make([]int, len(cols))
	for i, c := range cols {
		widths[i] = len(c.Name)
	}
	for _, row := range append(rows, totals) {
		for i, cell := range row {
			widths[i] = max(widths[i], len(cell))
		}
	}
	printer.PrintRow(w, cols, widths)
	printer.PrintSeparator(w, widths)
	for _, row := range rows {
		printer.PrintValues(w, anySlice(row), cols, widths)
	}
	printer.PrintSeparator(w, widths)
	printer.PrintValues(w, anySlice(totals), cols, widths)
}

// probeTableColumns returns the table spec; numeric columns right-justify.
func probeTableColumns(hasErr bool) []printer.ColumnSpec {
	cols := []printer.ColumnSpec{
		{Name: "tenant"},
		{Name: "shardID"},
		{Name: "window"},
		{Name: "bytes", Align: ">"},
		{Name: "dur", Align: ">"},
		{Name: "avg", Align: ">"},
	}
	if hasErr {
		cols = append(cols, printer.ColumnSpec{Name: "error"})
	}
	return cols
}

// formatBps renders a throughput with a per-second suffix.
func formatBps(bps int64) string {
	return printer.FormatBytes(float64(bps)) + "/s"
}

// truncateErr caps an error cell at probeErrWidth with a "..." tail.
func truncateErr(s string) string {
	if len(s) <= probeErrWidth {
		return s
	}
	return s[:probeErrWidth-3] + "..."
}

// anySlice converts row cells for printer.PrintValues.
func anySlice(row []string) []any {
	out := make([]any, len(row))
	for i, cell := range row {
		out[i] = cell
	}
	return out
}