package metrics

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"
)

// Stage identifies a request timing segment.
type Stage int

const (
	StageFrameRead Stage = iota
	StageQueueWait
	StageHandler
	StageBatchParse
	StageBatchDecode
	StageBatchMergePlan
	StageBatchExecuteWait
	StageSchemaSync
	StageSchemaCache
	StageCQLMetaCache
	StageMySQLWait
	StageMySQLQuery
	StageEncode
	StageWriteLockWait
	StageWrite
	StageTotal
	stageCount
)

var stageNames = [...]string{
	"frame_read",
	"queue_wait",
	"handler",
	"batch_parse",
	"batch_decode",
	"batch_merge_plan",
	"batch_execute_wait",
	"schema_sync",
	"schema_cache_refresh",
	"cql_meta_cache_refresh",
	"mysql_pool_wait",
	"mysql_query",
	"result_encode",
	"write_lock_wait",
	"write_flush",
	"total",
}

const (
	maxLoggedSQLStatements = 8
	maxLoggedSQLPreviewLen = 1024
)

// RequestMetrics stores per-request timing and target metadata.
type RequestMetrics struct {
	start  time.Time
	opcode byte
	opType string
	stream int16
	// captureDiagnostics enables slow-log-only details such as per-stage breakdown,
	// target metadata, executed SQL previews, and batch summaries.
	captureDiagnostics bool
	mu                 sync.Mutex
	keyspace           string
	table              string
	// requestBytes is the bytes consumed while reading the incoming request frame.
	requestBytes int
	// firstByteWait is the time spent waiting until first readable request byte.
	firstByteWait time.Duration
	// bodyRead is the remaining frame-read time after first byte is observed.
	bodyRead time.Duration
	stages   [stageCount]time.Duration

	executedSQLs         []string
	executedSQLCount     int
	executedSQLBytes     int
	executedSQLDataBytes int
	clientPayload        map[string]string

	batchType             string
	batchStatementCount   int
	batchExecuteItemCount int
	batchMergedGroupCount int
	batchMergedStmtCount  int

	metricUsername  string
	metricLabels    MetricLabels
	metricLabelsSet bool

	inFlightMetricLabels    MetricLabels
	inFlightMetricLabelsSet bool
}

// NewRequestMetrics creates a request tracker.
func NewRequestMetrics(opcode byte, stream int16, start time.Time) *RequestMetrics {
	return NewRequestMetricsWithDiagnostics(opcode, stream, start, true)
}

// NewRequestMetricsWithDiagnostics creates a request tracker with optional slow-log diagnostics.
func NewRequestMetricsWithDiagnostics(opcode byte, stream int16, start time.Time, captureDiagnostics bool) *RequestMetrics {
	return &RequestMetrics{
		start:              start,
		opcode:             opcode,
		stream:             stream,
		captureDiagnostics: captureDiagnostics,
	}
}

// AddStage records duration for a stage.
func (r *RequestMetrics) AddStage(stage Stage, d time.Duration) {
	if r == nil || stage < 0 || stage >= stageCount || !r.captureDiagnostics {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.stages[stage] += d
}

// StageDuration returns the recorded duration for a stage.
func (r *RequestMetrics) StageDuration(stage Stage) time.Duration {
	if r == nil || stage < 0 || stage >= stageCount {
		return 0
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.stages[stage]
}

// Start returns the request start time.
func (r *RequestMetrics) Start() time.Time {
	if r == nil {
		return time.Time{}
	}
	return r.start
}

// Stream returns the request stream ID.
func (r *RequestMetrics) Stream() int16 {
	if r == nil {
		return 0
	}
	return r.stream
}

// Opcode returns the request opcode.
func (r *RequestMetrics) Opcode() byte {
	if r == nil {
		return 0
	}
	return r.opcode
}

// SetOperationType stores a string operation type for request-completion metrics.
func (r *RequestMetrics) SetOperationType(opType string) {
	if r == nil {
		return
	}
	opType = strings.TrimSpace(opType)
	if opType == "" {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.opType = opType
}

// OperationType returns the stored string operation type for request-completion metrics.
func (r *RequestMetrics) OperationType() string {
	if r == nil {
		return ""
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.opType
}

// SetTarget stores keyspace/table identifiers when known.
func (r *RequestMetrics) SetTarget(keyspace, table string) {
	if r == nil || !r.captureDiagnostics {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if keyspace != "" {
		r.keyspace = keyspace
	}
	if table != "" {
		r.table = table
	}
}

// Target returns the stored keyspace/table identifiers.
func (r *RequestMetrics) Target() (string, string) {
	if r == nil {
		return "", ""
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.keyspace, r.table
}

// SetFrameReadBreakdown stores frame-read breakdown for slow-log diagnostics.
func (r *RequestMetrics) SetFrameReadBreakdown(requestBytes int, firstByteWait, bodyRead time.Duration) {
	if r == nil {
		return
	}
	if requestBytes < 0 {
		requestBytes = 0
	}
	if firstByteWait < 0 {
		firstByteWait = 0
	}
	if bodyRead < 0 {
		bodyRead = 0
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.requestBytes = requestBytes
	r.firstByteWait = firstByteWait
	r.bodyRead = bodyRead
}

// FrameReadBreakdown returns request bytes and read timing split.
func (r *RequestMetrics) FrameReadBreakdown() (requestBytes int, firstByteWait, bodyRead time.Duration) {
	if r == nil {
		return 0, 0, 0
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.requestBytes, r.firstByteWait, r.bodyRead
}

// RecordSQL appends a SQL preview and accumulates approximate SQL/data sizes.
func (r *RequestMetrics) RecordSQL(sql string, args ...interface{}) {
	if r == nil || !r.captureDiagnostics {
		return
	}
	trimmed := strings.TrimSpace(sql)
	dataBytes := estimateSQLArgsSize(args)

	r.mu.Lock()
	defer r.mu.Unlock()

	r.executedSQLCount++
	r.executedSQLBytes += len(trimmed)
	r.executedSQLDataBytes += dataBytes

	if trimmed == "" || len(r.executedSQLs) >= maxLoggedSQLStatements {
		return
	}
	r.executedSQLs = append(r.executedSQLs, trimmed)
}

// SQLSummary returns recorded SQL diagnostics for slow-log reporting.
func (r *RequestMetrics) SQLSummary() (count, sqlBytes, dataBytes, omitted int, sqls []string) {
	if r == nil {
		return 0, 0, 0, 0, nil
	}
	r.mu.Lock()
	rawSQLs := append([]string(nil), r.executedSQLs...)
	count, sqlBytes, dataBytes, omitted = r.executedSQLCount, r.executedSQLBytes, r.executedSQLDataBytes, r.executedSQLCount-len(r.executedSQLs)
	r.mu.Unlock()

	if len(rawSQLs) > 0 {
		sqls = make([]string, 0, len(rawSQLs))
		for _, sql := range rawSQLs {
			sqls = append(sqls, previewSQL(sql))
		}
	}
	return count, sqlBytes, dataBytes, omitted, sqls
}

// SetClientPayload stores decoded client payload for diagnostics.
func (r *RequestMetrics) SetClientPayload(payload map[string]string) {
	if r == nil || !r.captureDiagnostics {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(payload) == 0 {
		r.clientPayload = nil
		return
	}
	cloned := make(map[string]string, len(payload))
	for k, v := range payload {
		cloned[k] = v
	}
	r.clientPayload = cloned
}

// SetMetricUsername stores the username used to derive metric labels.
func (r *RequestMetrics) SetMetricUsername(username string) {
	if r == nil {
		return
	}
	username = strings.TrimSpace(username)
	r.mu.Lock()
	defer r.mu.Unlock()
	r.metricUsername = username
	r.metricLabels = MetricLabels{}
	r.metricLabelsSet = false
}

// MetricUsername returns the raw username used for metric labels.
func (r *RequestMetrics) MetricUsername() string {
	if r == nil {
		return ""
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.metricUsername
}

// MetricLabels returns the resolved metric labels for the request.
func (r *RequestMetrics) MetricLabels(ctx context.Context) MetricLabels {
	if r == nil {
		return defaultMetricLabels()
	}
	r.mu.Lock()
	if r.metricLabelsSet {
		labels := r.metricLabels.withDefaults()
		r.mu.Unlock()
		return labels
	}
	username := r.metricUsername
	r.mu.Unlock()

	labels := ResolveMetricLabels(ctx, username)

	r.mu.Lock()
	r.metricLabels = labels
	r.metricLabelsSet = true
	r.mu.Unlock()
	return labels
}

// CaptureInFlightMetricLabels stores the labels used by the in-flight increment.
func (r *RequestMetrics) CaptureInFlightMetricLabels(ctx context.Context) MetricLabels {
	if r == nil {
		return defaultMetricLabels()
	}
	labels := r.MetricLabels(ctx)
	r.mu.Lock()
	defer r.mu.Unlock()
	r.inFlightMetricLabels = labels.withDefaults()
	r.inFlightMetricLabelsSet = true
	return r.inFlightMetricLabels
}

// InFlightMetricLabels returns the labels captured at the in-flight increment.
func (r *RequestMetrics) InFlightMetricLabels(ctx context.Context) MetricLabels {
	if r == nil {
		return defaultMetricLabels()
	}
	r.mu.Lock()
	if r.inFlightMetricLabelsSet {
		labels := r.inFlightMetricLabels.withDefaults()
		r.mu.Unlock()
		return labels
	}
	r.mu.Unlock()
	return r.MetricLabels(ctx)
}

// ClientPayload returns stored client payload.
func (r *RequestMetrics) ClientPayload() map[string]string {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.clientPayload) == 0 {
		return nil
	}
	cloned := make(map[string]string, len(r.clientPayload))
	for k, v := range r.clientPayload {
		cloned[k] = v
	}
	return cloned
}

// SetBatchInfo stores batch execution diagnostics for slow-log reporting.
func (r *RequestMetrics) SetBatchInfo(batchType string, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount int) {
	if r == nil || !r.captureDiagnostics {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()

	if batchType != "" {
		r.batchType = batchType
	}
	if statementCount >= 0 {
		r.batchStatementCount = statementCount
	}
	if executeItemCount >= 0 {
		r.batchExecuteItemCount = executeItemCount
	}
	if mergedGroupCount >= 0 {
		r.batchMergedGroupCount = mergedGroupCount
	}
	if mergedStmtCount >= 0 {
		r.batchMergedStmtCount = mergedStmtCount
	}
}

// BatchInfo returns recorded batch diagnostics.
func (r *RequestMetrics) BatchInfo() (batchType string, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount int) {
	if r == nil {
		return "", 0, 0, 0, 0
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.batchType, r.batchStatementCount, r.batchExecuteItemCount, r.batchMergedGroupCount, r.batchMergedStmtCount
}

type requestContextKey struct{}

// WithRequestMetrics attaches the request tracker to the context.
func WithRequestMetrics(ctx context.Context, r *RequestMetrics) context.Context {
	if r == nil {
		return ctx
	}
	if username := usernameFromContext(ctx); username != "" {
		r.SetMetricUsername(username)
	}
	return context.WithValue(ctx, requestContextKey{}, r)
}

// FromContext extracts the request tracker from the context.
func FromContext(ctx context.Context) *RequestMetrics {
	if ctx == nil {
		return nil
	}
	if v := ctx.Value(requestContextKey{}); v != nil {
		if r, ok := v.(*RequestMetrics); ok {
			return r
		}
	}
	return nil
}

// ObserveStageFromContext records a stage duration using the tracker on the context.
func ObserveStageFromContext(ctx context.Context, stage Stage, d time.Duration) {
	r := FromContext(ctx)
	if r == nil {
		state := activeState()
		if state != nil && state.enabled && state.stageHist != nil && stage >= 0 && stage < stageCount {
			labels := labelsFromContext(ctx)
			state.stageHist.WithLabelValues(stageNames[stage], labels.KeyspaceName, labels.ClusterID).Observe(d.Seconds())
		}
		return
	}
	ObserveStage(stage, d, r)
}

// SetRequestTarget updates the request tracker stored on the context.
func SetRequestTarget(ctx context.Context, keyspace, table string) {
	r := FromContext(ctx)
	if r == nil {
		return
	}
	r.SetTarget(keyspace, table)
}

// SetRequestOperationType updates the request tracker stored on the context.
func SetRequestOperationType(ctx context.Context, opType string) {
	r := FromContext(ctx)
	if r == nil {
		return
	}
	r.SetOperationType(opType)
}

// RecordRequestSQL records executed SQL and approximate bind-data size on the request tracker.
func RecordRequestSQL(ctx context.Context, sql string, args ...interface{}) {
	r := FromContext(ctx)
	if r == nil {
		return
	}
	r.RecordSQL(sql, args...)
}

// SetRequestBatchInfo records batch execution diagnostics on the request tracker.
func SetRequestBatchInfo(ctx context.Context, batchType string, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount int) {
	r := FromContext(ctx)
	if r == nil {
		return
	}
	r.SetBatchInfo(batchType, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount)
}

// SetRequestClientPayload records decoded client payload on the request tracker.
func SetRequestClientPayload(ctx context.Context, payload map[string]string) {
	r := FromContext(ctx)
	if r == nil {
		return
	}
	r.SetClientPayload(payload)
}

func previewSQL(sql string) string {
	normalized := strings.Join(strings.Fields(strings.TrimSpace(sql)), " ")
	if normalized == "" {
		return ""
	}
	if len(normalized) <= maxLoggedSQLPreviewLen {
		return normalized
	}
	return normalized[:maxLoggedSQLPreviewLen] + fmt.Sprintf("...<truncated %d bytes>", len(normalized)-maxLoggedSQLPreviewLen)
}

func estimateSQLArgsSize(args []interface{}) int {
	total := 0
	for _, arg := range args {
		total += estimateSQLArgSize(arg)
	}
	return total
}

func estimateSQLArgSize(arg interface{}) int {
	switch v := arg.(type) {
	case nil:
		return 0
	case bool:
		return 1
	case int8, uint8:
		return 1
	case int16, uint16:
		return 2
	case int32, uint32, float32:
		return 4
	case int, int64, uint, uint64, float64:
		return 8
	case string:
		return len(v)
	case []byte:
		return len(v)
	case time.Time:
		return len(v.Format(time.RFC3339Nano))
	case []interface{}:
		return estimateSQLArgsSize(v)
	case fmt.Stringer:
		return len(v.String())
	default:
		return len(fmt.Sprint(v))
	}
}