package vmquery

import (
	"context"
	"errors"
	"fmt"
	"io"
	"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"
)

// MetricNames is the list of canary metrics verified per shard. Order matters:
// short labels and table cells render in this order.
var MetricNames = []string{
	"node_os_info",
	"tidb_executor_statement_total",
}

// MetricShort is the compact table label per metric name.
var MetricShort = map[string]string{
	"node_os_info":                  "k8s",
	"tidb_executor_statement_total": "tidb",
}

// FormatCounts renders per-metric counts as "short=count" pairs joined by
// "/", in MetricNames order.
func FormatCounts(counts map[string]int64) string {
	parts := make([]string, 0, len(MetricNames))
	for _, n := range MetricNames {
		short := MetricShort[n]
		if short == "" {
			short = n
		}
		parts = append(parts, fmt.Sprintf("%s=%d", short, counts[n]))
	}
	return strings.Join(parts, "/")
}

// QueryWindowSeconds is the canary sampling window each shard is verified
// over, evaluated at the shard's TimeEnd. Sampling is the design: a
// full-window count per shard would be prohibitively expensive on both
// clusters, so verification trades interior coverage for bounded query
// cost. Data missing earlier in the shard window is not detected.
const QueryWindowSeconds = 1800

const DefaultTolerance = 0.01

// VerifyOptions carries the verify run parameters.
type VerifyOptions struct {
	Manifest             string
	OriginVmselectURL    string
	CandidateVmselectURL string
	TimeoutSeconds       int
	Tolerance            float64
}

// Validate checks cross-flag invariants.
func (o VerifyOptions) Validate() error {
	var errs []error
	if o.Tolerance < 0 {
		errs = append(errs, errors.New("--tolerance must be greater than or equal to 0"))
	}
	if o.CandidateVmselectURL == "" {
		errs = append(errs, errors.New("--candidate-vmselect-url is required"))
	}
	return errors.Join(errs...)
}

// VerifyResult captures one shard's verify outcome for presentation.
type VerifyResult struct {
	Index        int
	Shard        schema.Shard
	Verification *schema.Verification
	OK           bool
	Duration     time.Duration
	ErrorLabel   string
}

// VerifyOutput separates concise progress from detailed query failures.
type VerifyOutput struct {
	Progress io.Writer
	Log      io.Writer
}

func (o VerifyOutput) normalized() VerifyOutput {
	if o.Progress == nil {
		o.Progress = io.Discard
	}
	if o.Log == nil {
		o.Log = io.Discard
	}
	return o
}

// VerifyRunner executes one verify pass over a manifest. Construct via NewVerifyRunner;
// caller invokes Run, which streams operator progress and returns
// per-shard results for table rendering.
type VerifyRunner struct {
	opts   VerifyOptions
	client Client
	output VerifyOutput
}

// NewVerifyRunner validates opts and returns a VerifyRunner.
func NewVerifyRunner(opts VerifyOptions, client Client, output VerifyOutput) (*VerifyRunner, error) {
	if err := opts.Validate(); err != nil {
		return nil, err
	}
	return &VerifyRunner{opts: opts, client: client, output: output.normalized()}, nil
}

// Run executes the verify pass. Writes operator progress, persists
// per-shard outcomes to the status store, and returns one VerifyResult per
// verified shard in iteration order.
func (r *VerifyRunner) Run(ctx context.Context) ([]VerifyResult, error) {
	statusPath := status.PathForManifest(r.opts.Manifest)
	store, err := status.Open(statusPath, nil)
	if err != nil {
		return nil, err
	}

	var selectedIDs []string
	store.View(func(d *schema.Document) {
		for i := range d.Shards {
			s := &d.Shards[i]
			if s.Phase != schema.PhaseImported {
				continue
			}
			if s.Verified {
				continue
			}
			selectedIDs = append(selectedIDs, s.ShardID)
		}
	})
	failed := []string{}
	store.View(func(d *schema.Document) {
		printer.Printf(r.output.Progress, "verification start: total=%d imported=%d verified=%d selected=%d",
			d.Summary.Total, d.Summary.Imported, d.Summary.Verified, len(selectedIDs))
	})

	timeout := time.Duration(r.opts.TimeoutSeconds) * time.Second
	runStarted := time.Now()
	results := make([]VerifyResult, 0, len(selectedIDs))
	for i, shardID := range selectedIDs {
		printer.Printf(r.output.Progress, "[%d/%d] shardID=%s", i+1, len(selectedIDs), shardID)
		started := time.Now()

		var shardSnap schema.Shard
		store.View(func(d *schema.Document) {
			if s := status.FindShard(d, shardID); s != nil {
				shardSnap = *s
			}
		})

		verification, ok, qerr := verifyShardCounts(ctx, r.client, shardSnap, r.opts.OriginVmselectURL, r.opts.CandidateVmselectURL, timeout, r.opts.Tolerance)
		result := VerifyResult{Index: i + 1, Shard: shardSnap}
		if qerr != nil {
			result.ErrorLabel = classifyError(qerr.Error())
			failed = append(failed, shardID)
			printer.Printf(r.output.Log, "[shard %s] %s: %s", shardID, result.ErrorLabel, truncateMsg(qerr.Error(), 200))
			if err := store.Mutate(ctx, func(d *schema.Document) error {
				return mutateVerifyFailure(d, shardID, qerr)
			}); err != nil {
				return results, persistFailure(shardID, err)
			}
		} else {
			result.OK = ok
			if !ok {
				failed = append(failed, shardID)
				printer.Printf(r.output.Log, "[shard %s] verification mismatch", shardID)
			}
			// Counts go to progress so a hollow 0-vs-0 pass is visible
			// where the table is not rendered (service-mode /logs).
			statusWord := "verified"
			if !ok {
				statusWord = "mismatch"
			}
			printer.Printf(r.output.Progress, "[%d/%d] shardID=%s status=%s origin=%s candidate=%s diff=%s",
				i+1, len(selectedIDs), shardID, statusWord,
				FormatCounts(verification.OriginCounts), FormatCounts(verification.CandidateCounts),
				FormatCounts(verification.Diffs))
			if err := store.Mutate(ctx, func(d *schema.Document) error {
				return mutateVerifyResult(d, shardID, verification, ok, r.opts.Tolerance)
			}); err != nil {
				return results, persistFailure(shardID, err)
			}
		}
		result.Duration = time.Since(started)
		result.Verification = verification
		store.View(func(d *schema.Document) {
			if s := status.FindShard(d, shardID); s != nil {
				result.Shard = *s
			}
		})
		results = append(results, result)
	}

	totalDuration := time.Since(runStarted)
	var total, verified int
	store.View(func(d *schema.Document) {
		total = d.Summary.Total
		verified = d.Summary.Verified
		printer.Printf(r.output.Progress, "verification done: total=%d imported=%d verified=%d failed=%d pending=%d duration=%.1fs",
			d.Summary.Total, d.Summary.Imported, d.Summary.Verified,
			d.Summary.Failed, d.Summary.Pending, totalDuration.Seconds())
	})

	if len(failed) > 0 {
		printer.Printf(r.output.Progress, "verification failed shards: %s", strings.Join(failed, ", "))
		return results, exitcode.PartialFailure(fmt.Errorf("%d shard(s) failed verification: %s", len(failed), strings.Join(failed, ", ")))
	}
	if verified != total {
		return results, exitcode.PartialFailure(fmt.Errorf("verification incomplete: verified=%d total=%d", verified, total))
	}
	return results, nil
}

// persistFailure converts a status persistence failure into a run error.
// ctx cancellation is returned as-is so callers treat it as shutdown, not a
// verification failure.
func persistFailure(shardID string, err error) error {
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return err
	}
	return exitcode.PartialFailure(fmt.Errorf("persist status for shard %s: %w", shardID, err))
}

// mutateVerifyFailure records a failed query attempt on a shard.
func mutateVerifyFailure(doc *schema.Document, shardID string, qerr error) error {
	s := status.FindShard(doc, shardID)
	if s == nil {
		return nil
	}
	s.Verified = false
	s.VerifiedAt = nil
	s.Verification = nil
	s.LastError = new(qerr.Error())
	return nil
}

// mutateVerifyResult records a verify query's outcome on a shard. ok=true marks
// the shard verified; ok=false records the appropriate LastError message.
func mutateVerifyResult(doc *schema.Document, shardID string, verification *schema.Verification, ok bool, tolerance float64) error {
	s := status.FindShard(doc, shardID)
	if s == nil {
		return nil
	}
	s.Verification = verification
	s.Verified = ok
	if ok {
		s.VerifiedAt = new(status.UTCNow())
		s.LastError = nil
		return nil
	}
	s.VerifiedAt = nil
	msg := fmt.Sprintf("verification count diff exceeds tolerance %g", tolerance)
	s.LastError = new(msg)
	return nil
}

// verifyShardCounts issues per-metric queries against origin and candidate
// vmselect, returning a Verification record with the counts and per-metric
// diffs.
func verifyShardCounts(ctx context.Context, client Client, shard schema.Shard, originTemplate, candidateTemplate string, timeout time.Duration, tolerance float64) (*schema.Verification, bool, error) {
	timeUnix, err := shardQueryTimeUnix(shard)
	if err != nil {
		return nil, false, err
	}
	originCounts, err := fetchCountsByMetric(ctx, client, originTemplate, shard.Tenant, timeUnix, timeout)
	if err != nil {
		return nil, false, fmt.Errorf("origin: %w", err)
	}
	candidateCounts, err := fetchCountsByMetric(ctx, client, candidateTemplate, shard.Tenant, timeUnix, timeout)
	if err != nil {
		return nil, false, fmt.Errorf("candidate: %w", err)
	}

	diffs := map[string]int64{}
	ok := true
	for _, metric := range MetricNames {
		origin := originCounts[metric]
		candidate := candidateCounts[metric]
		diff := origin - candidate
		if diff < 0 {
			diff = -diff
		}
		diffs[metric] = diff
		allowed := float64(candidate) * tolerance
		if candidate > 0 {
			if float64(diff) > allowed {
				ok = false
			}
		} else {
			if origin != 0 {
				ok = false
			}
		}
	}
	return &schema.Verification{
		OriginCounts:    originCounts,
		CandidateCounts: candidateCounts,
		Diffs:           diffs,
	}, ok, nil
}

func fetchCountsByMetric(ctx context.Context, client Client, urlTemplate, tenant string, timeUnix int64, timeout time.Duration) (map[string]int64, error) {
	out := map[string]int64{}
	tenantURL := RenderURL(urlTemplate, tenant)
	for _, metric := range MetricNames {
		params := map[string]string{
			"query": MetricQuery(metric),
			"time":  strconv.FormatInt(timeUnix, 10),
		}
		resp, err := client.Query(ctx, tenantURL, "/api/v1/query", params, timeout)
		if err != nil {
			return nil, fmt.Errorf("metric %s: %w", metric, err)
		}
		count, err := resp.Count()
		if err != nil {
			return nil, fmt.Errorf("metric %s decode: %w", metric, err)
		}
		out[metric] = int64(count)
	}
	return out, nil
}

// MetricQuery returns the PromQL count_over_time expression for a canary metric.
func MetricQuery(metric string) string {
	return fmt.Sprintf("sum(count_over_time(%s[%ds]))", metric, QueryWindowSeconds)
}

func shardQueryTimeUnix(shard schema.Shard) (int64, error) {
	t, err := status.ParseTime(shard.TimeEnd)
	if err != nil {
		return 0, fmt.Errorf("parse timeEnd %q: %w", shard.TimeEnd, err)
	}
	return t.UTC().Unix(), nil
}

func classifyError(msg string) string {
	lower := strings.ToLower(msg)
	if strings.Contains(lower, "http ") {
		return "http"
	}
	return "error"
}

func truncateMsg(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n]
}