package cmd

import (
	"fmt"
	"strconv"
	"strings"
	"time"
)

const (
	diagnoseDefaultMaxSeedExpand  = 20
	diagnoseDefaultMaxRelatedLogs = 1000
	diagnoseRelatedQueryLimit     = 100
)

type diagnoseExpansionOutput struct {
	QueryPlan []DiagnoseQueryRecord
	Seeds     []DiagnoseLogEntry
	Related   []DiagnoseLogEntry
	Warnings  []DiagnoseWarning
}

type diagnoseRelatedQuerySpec struct {
	Purpose  string
	Anchor   string
	Value    string
	Window   time.Duration
	Strength string
}

func expandDiagnoseRelatedLogs(
	ctx *RuntimeContext,
	clusterID string,
	selector string,
	seedMatches []DiagnoseLogEntry,
	opts diagnoseOptions,
	plan *diagnoseQueryPlanBuilder,
	location *time.Location,
) diagnoseExpansionOutput {
	output := diagnoseExpansionOutput{
		Seeds:    seedMatches,
		Related:  []DiagnoseLogEntry{},
		Warnings: []DiagnoseWarning{},
	}
	if location == nil {
		location = time.Local
	}
	if len(seedMatches) == 0 {
		return output
	}
	if opts.NoExpand {
		output.Warnings = append(output.Warnings, diagnoseExpansionWarning("expansion_skipped", "related_logs", "", "related-log expansion skipped by --no-expand"))
		return output
	}
	if len(seedMatches) > opts.MaxSeedExpand {
		output.Warnings = append(output.Warnings, diagnoseExpansionWarning(
			"expansion_skipped",
			"seed_matches",
			"",
			fmt.Sprintf("seed match count %d exceeds --max-seed-expand %d", len(seedMatches), opts.MaxSeedExpand),
		))
		return output
	}
	if opts.MaxRelatedLogs == 0 {
		output.Warnings = append(output.Warnings, diagnoseExpansionWarning("limit", "related_logs", "", "related-log expansion would run but --max-related-logs is 0"))
		return output
	}

	relatedLimit := diagnoseRelatedQueryLimit
	if relatedLimit > lokiMaxLimit {
		relatedLimit = lokiMaxLimit
	}
	seedIndex := make(map[string]int, len(output.Seeds))
	for i, seed := range output.Seeds {
		seedIndex[diagnoseEntryIdentity(seed)] = i
	}
	relatedIndex := make(map[string]int)

	for _, seed := range seedMatches {
		seedTime, err := time.Parse(time.RFC3339Nano, seed.TimeUTC)
		if err != nil {
			continue
		}
		for _, spec := range diagnoseRelatedQuerySpecs(seed) {
			record := plan.newLokiRelatedRecord(
				spec.Purpose,
				diagnoseRelatedLogQL(selector, spec),
				seedTime.Add(-spec.Window).UTC(),
				seedTime.Add(spec.Window).UTC(),
				relatedLimit,
			)
			started := time.Now()
			body, err := executeLokiQueryRangeRawWithRetry(ctx, diagnoseLokiQueryRangeRequest(clusterID, record))
			resultCount := 0
			if err == nil {
				resultCount, err = countDiagnoseLokiResults(body)
			}
			completeDiagnoseQueryRecord(&record, started, resultCount, err)
			output.QueryPlan = append(output.QueryPlan, record)
			if err != nil {
				output.Warnings = append(output.Warnings, diagnoseExpansionWarning("query_error", spec.Purpose, record.ID, err.Error()))
				continue
			}
			if record.HitLimit {
				output.Warnings = append(output.Warnings, diagnoseExpansionWarning("limit", spec.Purpose, record.ID, "related query reached its per-query limit"))
			}

			entries, parseWarnings, err := normalizeDiagnoseLogEntries(body, opts.Seed, []string{record.ID}, location)
			if err != nil {
				output.Warnings = append(output.Warnings, diagnoseExpansionWarning("query_error", spec.Purpose, record.ID, err.Error()))
				continue
			}
			output.Warnings = append(output.Warnings, parseWarnings...)
			for _, entry := range entries {
				key := diagnoseEntryIdentity(entry)
				if seedPos, ok := seedIndex[key]; ok {
					output.Seeds[seedPos].QueryIDs = mergeDiagnoseQueryIDs(output.Seeds[seedPos].QueryIDs, []string{record.ID})
					continue
				}
				relation := diagnoseRelationForEntry(seed, entry, spec)
				if relatedPos, ok := relatedIndex[key]; ok {
					output.Related[relatedPos].QueryIDs = mergeDiagnoseQueryIDs(output.Related[relatedPos].QueryIDs, []string{record.ID})
					output.Related[relatedPos].Relations = appendDiagnoseRelation(output.Related[relatedPos].Relations, relation)
					continue
				}
				if len(output.Related) >= opts.MaxRelatedLogs {
					output.Warnings = append(output.Warnings, diagnoseExpansionWarning("limit", "related_logs", record.ID, "global related-log cap reached"))
					return output
				}
				entry.Relations = []DiagnoseRelation{relation}
				entry.ID = fmt.Sprintf("log-%03d", len(output.Seeds)+len(output.Related)+1)
				relatedIndex[key] = len(output.Related)
				output.Related = append(output.Related, entry)
			}
		}
	}
	return output
}

func diagnoseRelatedQuerySpecs(seed DiagnoseLogEntry) []diagnoseRelatedQuerySpec {
	specs := make([]diagnoseRelatedQuerySpec, 0, 8)
	if value := diagnoseAnchorString(seed, "conn", "conn_id"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_connection", Anchor: "conn", Value: value, Window: 10 * time.Minute, Strength: "strong"})
	}
	if value := diagnoseAnchorString(seed, "txn_start_ts", "timestamp_or_start_ts"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_transaction_ts", Anchor: "txn_start_ts", Value: value, Window: 10 * time.Minute, Strength: "strong"})
	}
	if value := diagnoseFirstAnchorString(seed, "task_id", "job_id", "subtask_id"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_task_or_job", Anchor: "task_or_job", Value: value, Window: 30 * time.Minute, Strength: "strong"})
	}
	if value := diagnoseAnchorString(seed, "instance"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_instance_nearby", Anchor: "instance", Value: value, Window: 5 * time.Minute, Strength: "weak"})
	}
	if value := diagnoseAnchorString(seed, "region_id"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_region", Anchor: "region_id", Value: value, Window: 30 * time.Minute, Strength: "medium"})
	}
	if value := diagnoseFirstAnchorString(seed, "store_addr", "peer_addr"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_store_address", Anchor: "store_or_peer_addr", Value: value, Window: 10 * time.Minute, Strength: "medium"})
	}
	if value := diagnoseFirstAnchorString(seed, "sql", "sql_digest", "table"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "same_sql_or_table", Anchor: "sql_or_table", Value: value, Window: 30 * time.Minute, Strength: "medium"})
	}
	if value := diagnoseAnchorString(seed, "namespace"); value != "" {
		specs = append(specs, diagnoseRelatedQuerySpec{Purpose: "nearby_warn_error_same_namespace", Anchor: "namespace", Value: value, Window: 5 * time.Minute, Strength: "weak"})
	}
	if len(specs) > 8 {
		return specs[:8]
	}
	return specs
}

func diagnoseRelatedLogQL(selector string, spec diagnoseRelatedQuerySpec) string {
	query := selector + " |= " + strconv.Quote(spec.Value)
	if spec.Purpose == "nearby_warn_error_same_namespace" {
		query += " |= " + strconv.Quote("WARN") + " |= " + strconv.Quote("ERROR")
	}
	return query
}

func diagnoseRelationForEntry(seed DiagnoseLogEntry, entry DiagnoseLogEntry, spec diagnoseRelatedQuerySpec) DiagnoseRelation {
	strength := spec.Strength
	if strength == "" {
		strength = diagnoseRelationStrength(seed, entry)
	}
	return DiagnoseRelation{
		SeedID:   seed.ID,
		Type:     spec.Purpose,
		Strength: strength,
		Anchor:   spec.Anchor,
		Value:    spec.Value,
	}
}

func diagnoseRelationStrength(seed DiagnoseLogEntry, entry DiagnoseLogEntry) string {
	for _, key := range []string{"conn", "conn_id", "txn_start_ts", "request_id", "trace_id", "task_id", "job_id", "subtask_id"} {
		if diagnoseSharedAnchor(seed, entry, key) {
			return "strong"
		}
	}
	for _, key := range []string{"region_id", "sql", "sql_digest", "store_addr", "peer_addr", "table"} {
		if diagnoseSharedAnchor(seed, entry, key) {
			return "medium"
		}
	}
	return "weak"
}

func diagnoseSharedAnchor(left, right DiagnoseLogEntry, key string) bool {
	leftValue := diagnoseAnchorString(left, key)
	rightValue := diagnoseAnchorString(right, key)
	return leftValue != "" && leftValue == rightValue
}

func diagnoseAnchorString(entry DiagnoseLogEntry, keys ...string) string {
	for _, key := range keys {
		raw, ok := entry.Anchors[key]
		if !ok {
			continue
		}
		switch value := raw.(type) {
		case string:
			if strings.TrimSpace(value) != "" {
				return strings.TrimSpace(value)
			}
		case []string:
			if len(value) > 0 && strings.TrimSpace(value[0]) != "" {
				return strings.TrimSpace(value[0])
			}
		case []any:
			if len(value) > 0 {
				if text, ok := value[0].(string); ok && strings.TrimSpace(text) != "" {
					return strings.TrimSpace(text)
				}
			}
		}
	}
	return ""
}

func diagnoseFirstAnchorString(entry DiagnoseLogEntry, keys ...string) string {
	return diagnoseAnchorString(entry, keys...)
}

func diagnoseEntryIdentity(entry DiagnoseLogEntry) string {
	return strings.Join([]string{entry.TimeUTC, entry.Stream, entry.Raw}, "\x00")
}

func appendDiagnoseRelation(existing []DiagnoseRelation, addition DiagnoseRelation) []DiagnoseRelation {
	for _, relation := range existing {
		if relation.SeedID == addition.SeedID && relation.Type == addition.Type && relation.Anchor == addition.Anchor && relation.Value == addition.Value {
			return existing
		}
	}
	return append(existing, addition)
}

func diagnoseExpansionWarning(code, source, queryID, message string) DiagnoseWarning {
	return DiagnoseWarning{
		Code:    code,
		Source:  source,
		QueryID: queryID,
		Message: message,
	}
}