package cmd

import (
	"fmt"
	"sort"
	"strings"
	"time"
)

const (
	diagnoseIncidentStatusCorrelated      = "correlated"
	diagnoseDefaultCorrelationWindow      = 5 * time.Minute
	diagnoseMinimumDirectIncidentEdge     = 8.0
	diagnoseMinimumBridgeIncidentEdge     = 12.0
	diagnoseSameConnPoints                = 10.0
	diagnoseSameTxnStartTSPoints          = 10.0
	diagnoseSameRegionIDPoints            = 8.0
	diagnoseSameSQLPoints                 = 6.0
	diagnoseSameTablePoints               = 4.0
	diagnoseSameErrorTextPoints           = 4.0
	diagnoseSameInstancePoints            = 2.0
	diagnoseBroadErrorTokenPoints         = 1.0
	diagnoseTimestampWithinFiveSecondsPts = 5.0
	diagnoseTimestampWithinSixtySecondsPt = 2.0
)

type diagnoseScoreEdge struct {
	left    int
	right   int
	points  float64
	reasons []ScoreReason
}

type diagnoseUnionFind struct {
	parent []int
	size   []int
}

func correlateDiagnoseEvidence(scope QueryScope, evidence []Evidence) DiagnoseResult {
	result := DiagnoseResult{
		QueryScope: scope,
		Incidents:  []Incident{},
		Warnings:   []string{},
	}

	if len(evidence) == 0 {
		return result
	}

	edges := make([]diagnoseScoreEdge, 0)
	for left := 0; left < len(evidence); left++ {
		for right := left + 1; right < len(evidence); right++ {
			reasons := diagnoseScoreReasonsForPair(evidence[left], evidence[right])
			if len(reasons) == 0 {
				continue
			}

			points := 0.0
			strong := false
			for _, reason := range reasons {
				points += reason.Points
				if diagnoseScoreReasonIsStrong(reason.Anchor) {
					strong = true
				}
			}
			if !strong || points < diagnoseMinimumDirectIncidentEdge {
				continue
			}

			edges = append(edges, diagnoseScoreEdge{
				left:    left,
				right:   right,
				points:  points,
				reasons: reasons,
			})
		}
	}

	sort.SliceStable(edges, func(i, j int) bool {
		if edges[i].points != edges[j].points {
			return edges[i].points > edges[j].points
		}
		leftDeltaI := diagnoseEvidenceTimeDelta(evidence[edges[i].left], evidence[edges[i].right])
		leftDeltaJ := diagnoseEvidenceTimeDelta(evidence[edges[j].left], evidence[edges[j].right])
		if leftDeltaI != leftDeltaJ {
			return leftDeltaI < leftDeltaJ
		}
		leftIDI := evidence[edges[i].left].ID
		leftIDJ := evidence[edges[j].left].ID
		if leftIDI != leftIDJ {
			return leftIDI < leftIDJ
		}
		rightIDI := evidence[edges[i].right].ID
		rightIDJ := evidence[edges[j].right].ID
		if rightIDI != rightIDJ {
			return rightIDI < rightIDJ
		}
		leftKeyI := stableEvidenceSortKey(evidence[edges[i].left])
		rightKeyI := stableEvidenceSortKey(evidence[edges[i].right])
		leftKeyJ := stableEvidenceSortKey(evidence[edges[j].left])
		rightKeyJ := stableEvidenceSortKey(evidence[edges[j].right])
		if leftKeyI != leftKeyJ {
			return leftKeyI < leftKeyJ
		}
		return rightKeyI < rightKeyJ
	})

	uf := newDiagnoseUnionFind(len(evidence))
	componentStrongAnchors := make([]map[string]struct{}, len(evidence))
	acceptedEdges := make([]diagnoseScoreEdge, 0, len(edges))
	for _, edge := range edges {
		leftRoot := uf.find(edge.left)
		rightRoot := uf.find(edge.right)
		if leftRoot == rightRoot {
			continue
		}

		edgeStrongAnchors := diagnoseStrongAnchorSet(edge.reasons)
		if edge.points < diagnoseMinimumBridgeIncidentEdge &&
			!diagnoseCanUseSubBridgeEdge(edgeStrongAnchors, leftRoot, rightRoot, uf, componentStrongAnchors) {
			continue
		}

		newRoot := uf.union(leftRoot, rightRoot)
		diagnoseMergeComponentStrongAnchors(componentStrongAnchors, newRoot, leftRoot, rightRoot, edgeStrongAnchors)
		acceptedEdges = append(acceptedEdges, edge)
	}

	grouped := make(map[int][]int)
	for idx := range evidence {
		root := uf.find(idx)
		grouped[root] = append(grouped[root], idx)
	}

	incidentByRoot := make(map[int]Incident)
	for root, indexes := range grouped {
		if len(indexes) < 2 {
			continue
		}

		componentEvidence := make([]Evidence, 0, len(indexes))
		for _, idx := range indexes {
			componentEvidence = append(componentEvidence, evidence[idx])
		}
		sort.Slice(componentEvidence, func(i, j int) bool {
			if componentEvidence[i].Time.Equal(componentEvidence[j].Time) {
				leftKey := stableEvidenceSortKey(componentEvidence[i])
				rightKey := stableEvidenceSortKey(componentEvidence[j])
				if leftKey != rightKey {
					return leftKey < rightKey
				}
				return componentEvidence[i].Line < componentEvidence[j].Line
			}
			return componentEvidence[i].Time.Before(componentEvidence[j].Time)
		})

		incidentEdges := diagnoseAcceptedEdgesForComponent(acceptedEdges, evidence, root, uf)
		if len(incidentEdges) == 0 {
			continue
		}

		incident := Incident{
			Seed:          scope.Seed,
			Status:        diagnoseIncidentStatusCorrelated,
			StartTime:     componentEvidence[0].Time,
			EndTime:       componentEvidence[len(componentEvidence)-1].Time,
			Evidence:      append([]Evidence(nil), componentEvidence...),
			EvidenceCount: len(componentEvidence),
		}

		incident.WeakEvidence = diagnoseWeakEvidence(componentEvidence)
		incident.WeakEvidenceCount = len(incident.WeakEvidence)
		incident.ScoreReasons = diagnoseFlattenScoreReasons(incidentEdges)
		incident.Score = diagnoseIncidentScore(incidentEdges)
		incident.Anchors = diagnoseAggregateIncidentAnchors(componentEvidence)
		incident.Summary = diagnoseIncidentSummary(incident.Anchors, len(componentEvidence))
		incident.NextSteps = diagnoseIncidentNextSteps(incident)
		incidentByRoot[root] = incident
	}

	incidents := make([]Incident, 0, len(incidentByRoot))
	for _, incident := range incidentByRoot {
		incidents = append(incidents, incident)
	}
	sortDiagnoseIncidents(incidents)
	assignDiagnoseIncidentIDs(incidents)
	result.Incidents = incidents

	groupedCount := 0
	for _, incident := range incidents {
		groupedCount += incident.EvidenceCount
	}
	if len(evidence) > groupedCount {
		remaining := len(evidence) - groupedCount
		if len(incidents) == 0 {
			result.Warnings = append(result.Warnings, "evidence was found but no incident passed the correlation threshold")
		} else {
			result.Warnings = append(result.Warnings, fmt.Sprintf("%d evidence item(s) remained below the correlation threshold", remaining))
		}
	}

	return result
}

func diagnoseScoreReasonsForPair(left, right Evidence) []ScoreReason {
	reasons := make([]ScoreReason, 0, 8)
	timeDelta := diagnoseEvidenceTimeDelta(left, right)
	if timeDelta > diagnoseDefaultCorrelationWindow {
		return nil
	}

	addReason := func(anchor, value, description string, points float64) {
		if value == "" || points <= 0 {
			return
		}
		reasons = append(reasons, ScoreReason{
			EvidenceIDs: []string{left.ID, right.ID},
			Anchor:      anchor,
			Value:       value,
			Points:      points,
			Description: description,
		})
	}

	if left.Anchors.ConnID != "" && left.Anchors.ConnID == right.Anchors.ConnID {
		addReason("conn_id", left.Anchors.ConnID, "shares conn_id anchor", diagnoseSameConnPoints)
	}
	if left.Anchors.TxnStartTS != "" && left.Anchors.TxnStartTS == right.Anchors.TxnStartTS {
		addReason("txn_start_ts", left.Anchors.TxnStartTS, "shares txn_start_ts anchor", diagnoseSameTxnStartTSPoints)
	}
	if left.Anchors.RegionID != "" && left.Anchors.RegionID == right.Anchors.RegionID {
		addReason("region_id", left.Anchors.RegionID, "shares region_id anchor", diagnoseSameRegionIDPoints)
	}
	if left.Anchors.SQL != "" && left.Anchors.SQL == right.Anchors.SQL {
		addReason("sql", left.Anchors.SQL, "shares normalized SQL anchor", diagnoseSameSQLPoints)
	}
	for _, table := range diagnoseSharedTableNames(left.Anchors.TableNames, right.Anchors.TableNames) {
		addReason("table_name", table, "shares table name anchor", diagnoseSameTablePoints)
	}
	if left.Anchors.ErrorText != "" && left.Anchors.ErrorText == right.Anchors.ErrorText {
		addReason("error_text", left.Anchors.ErrorText, "shares exact error text anchor", diagnoseSameErrorTextPoints)
	}

	if left.Instance != "" && left.Instance == right.Instance {
		addReason("instance", left.Instance, "shares instance as weak support", diagnoseSameInstancePoints)
	}

	if len(diagnoseBroadErrorTokens(left.Anchors.ErrorText, right.Anchors.ErrorText)) > 0 && !diagnoseHasReason(reasons, "error_text") {
		tokens := diagnoseBroadErrorTokens(left.Anchors.ErrorText, right.Anchors.ErrorText)
		addReason("error_token", strings.Join(tokens, ","), "shares broad error token as weak support", diagnoseBroadErrorTokenPoints)
	}

	if timeDelta <= 5*time.Second {
		addReason("timestamp", timeDelta.String(), "timestamps within 5 seconds", diagnoseTimestampWithinFiveSecondsPts)
	} else if timeDelta <= 60*time.Second {
		addReason("timestamp", timeDelta.String(), "timestamps within 60 seconds", diagnoseTimestampWithinSixtySecondsPt)
	}

	if len(reasons) == 0 {
		return nil
	}
	sort.SliceStable(reasons, func(i, j int) bool {
		if reasons[i].Points != reasons[j].Points {
			return reasons[i].Points > reasons[j].Points
		}
		if reasons[i].Anchor != reasons[j].Anchor {
			return reasons[i].Anchor < reasons[j].Anchor
		}
		if reasons[i].Value != reasons[j].Value {
			return reasons[i].Value < reasons[j].Value
		}
		return strings.Join(reasons[i].EvidenceIDs, ",") < strings.Join(reasons[j].EvidenceIDs, ",")
	})
	return reasons
}

func diagnoseAcceptedEdgesForComponent(edges []diagnoseScoreEdge, evidence []Evidence, root int, uf *diagnoseUnionFind) []diagnoseScoreEdge {
	componentEdges := make([]diagnoseScoreEdge, 0)
	for _, edge := range edges {
		if uf.find(edge.left) != root || uf.find(edge.right) != root {
			continue
		}
		componentEdges = append(componentEdges, edge)
	}
	return componentEdges
}

func diagnoseStrongAnchorSet(reasons []ScoreReason) map[string]struct{} {
	anchors := make(map[string]struct{})
	for _, reason := range reasons {
		if !diagnoseScoreReasonIsStrong(reason.Anchor) {
			continue
		}
		anchors[diagnoseStrongAnchorKey(reason)] = struct{}{}
	}
	return anchors
}

func diagnoseCanUseSubBridgeEdge(edgeAnchors map[string]struct{}, leftRoot, rightRoot int, uf *diagnoseUnionFind, componentAnchors []map[string]struct{}) bool {
	if len(edgeAnchors) == 0 {
		return false
	}
	if uf.size[leftRoot] > 1 && !diagnoseAnchorSetsOverlap(componentAnchors[leftRoot], edgeAnchors) {
		return false
	}
	if uf.size[rightRoot] > 1 && !diagnoseAnchorSetsOverlap(componentAnchors[rightRoot], edgeAnchors) {
		return false
	}
	return true
}

func diagnoseMergeComponentStrongAnchors(componentAnchors []map[string]struct{}, newRoot, leftRoot, rightRoot int, edgeAnchors map[string]struct{}) {
	merged := make(map[string]struct{})
	for _, root := range []int{leftRoot, rightRoot, newRoot} {
		for anchor := range componentAnchors[root] {
			merged[anchor] = struct{}{}
		}
	}
	for anchor := range edgeAnchors {
		merged[anchor] = struct{}{}
	}

	componentAnchors[newRoot] = merged
	if leftRoot != newRoot {
		componentAnchors[leftRoot] = nil
	}
	if rightRoot != newRoot {
		componentAnchors[rightRoot] = nil
	}
}

func diagnoseAnchorSetsOverlap(left, right map[string]struct{}) bool {
	if len(left) == 0 || len(right) == 0 {
		return false
	}
	if len(left) > len(right) {
		left, right = right, left
	}
	for anchor := range left {
		if _, ok := right[anchor]; ok {
			return true
		}
	}
	return false
}

func diagnoseIncidentScore(edges []diagnoseScoreEdge) float64 {
	score := 0.0
	for _, edge := range edges {
		score += edge.points
	}
	return score
}

func diagnoseFlattenScoreReasons(edges []diagnoseScoreEdge) []ScoreReason {
	reasons := make([]ScoreReason, 0)
	for _, edge := range edges {
		reasons = append(reasons, edge.reasons...)
	}
	sort.SliceStable(reasons, func(i, j int) bool {
		if reasons[i].Points != reasons[j].Points {
			return reasons[i].Points > reasons[j].Points
		}
		if reasons[i].Anchor != reasons[j].Anchor {
			return reasons[i].Anchor < reasons[j].Anchor
		}
		if reasons[i].Value != reasons[j].Value {
			return reasons[i].Value < reasons[j].Value
		}
		return strings.Join(reasons[i].EvidenceIDs, ",") < strings.Join(reasons[j].EvidenceIDs, ",")
	})
	return reasons
}

func diagnoseAggregateIncidentAnchors(evidence []Evidence) Anchors {
	return Anchors{
		ConnID:      diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.ConnID }),
		TxnStartTS:  diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.TxnStartTS }),
		RegionID:    diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.RegionID }),
		StoreAddr:   diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.StoreAddr }),
		Keyspace:    diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.Keyspace }),
		SQL:         diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.SQL }),
		TableNames:  diagnoseMostFrequentTableNames(evidence),
		ErrorText:   diagnoseMostFrequentValue(evidence, func(item Evidence) string { return item.Anchors.ErrorText }),
		StackFrames: diagnoseMostFrequentStringsFromEvidence(evidence, func(item Evidence) []string { return item.Anchors.StackFrames }),
	}
}

func diagnoseIncidentSummary(anchors Anchors, evidenceCount int) string {
	parts := make([]string, 0, 4)
	if anchors.ConnID != "" {
		parts = append(parts, "conn_id")
	}
	if anchors.RegionID != "" {
		parts = append(parts, "region_id")
	}
	if anchors.TxnStartTS != "" {
		parts = append(parts, "txn_start_ts")
	}
	if anchors.SQL != "" {
		parts = append(parts, "normalized SQL")
	} else if len(anchors.TableNames) > 0 {
		parts = append(parts, "table name")
	}
	if anchors.ErrorText != "" {
		parts = append(parts, "error text")
	}

	if len(parts) == 0 {
		return fmt.Sprintf("Correlated %d nearby logs with shared anchors.", evidenceCount)
	}
	if len(parts) == 1 {
		return fmt.Sprintf("Shares %s anchor with nearby logs.", parts[0])
	}
	if len(parts) == 2 {
		return fmt.Sprintf("Shares %s and %s anchors with nearby logs.", parts[0], parts[1])
	}
	return fmt.Sprintf("Shares %s, %s, and %s anchors with nearby logs.", parts[0], parts[1], parts[2])
}

func diagnoseIncidentNextSteps(incident Incident) []ExpansionHint {
	return nil
}

func diagnoseWeakEvidence(evidence []Evidence) []Evidence {
	weak := make([]Evidence, 0)
	for _, item := range evidence {
		if diagnoseEvidenceIsWeakOnly(item) {
			weak = append(weak, item)
		}
	}
	return weak
}

func diagnoseEvidenceIsWeakOnly(evidence Evidence) bool {
	return evidence.Anchors.ConnID == "" &&
		evidence.Anchors.TxnStartTS == "" &&
		evidence.Anchors.RegionID == "" &&
		evidence.Anchors.SQL == "" &&
		len(evidence.Anchors.TableNames) == 0 &&
		evidence.Anchors.ErrorText == ""
}

func diagnoseSharedTableNames(left, right []string) []string {
	if len(left) == 0 || len(right) == 0 {
		return nil
	}

	seen := make(map[string]struct{}, len(left))
	for _, table := range left {
		if table != "" {
			seen[table] = struct{}{}
		}
	}
	shared := make([]string, 0)
	for _, table := range right {
		if table == "" {
			continue
		}
		if _, ok := seen[table]; ok {
			shared = append(shared, table)
		}
	}
	sort.Strings(shared)
	return shared
}

func diagnoseBroadErrorTokens(left, right string) []string {
	leftTokens := diagnoseTokenizeBroadErrorText(left)
	rightTokens := diagnoseTokenizeBroadErrorText(right)
	if len(leftTokens) == 0 || len(rightTokens) == 0 {
		return nil
	}

	common := make([]string, 0)
	for token := range leftTokens {
		if _, ok := rightTokens[token]; ok {
			common = append(common, token)
		}
	}
	sort.Strings(common)
	return common
}

func diagnoseTokenizeBroadErrorText(text string) map[string]struct{} {
	tokens := make(map[string]struct{})
	for _, token := range strings.FieldsFunc(strings.ToLower(text), func(r rune) bool {
		return !(r >= 'a' && r <= 'z') && !(r >= '0' && r <= '9')
	}) {
		token = strings.TrimSpace(token)
		if token == "" || len(token) < 4 || diagnoseBroadErrorTokenStopWords[token] {
			continue
		}
		tokens[token] = struct{}{}
	}
	return tokens
}

var diagnoseBroadErrorTokenStopWords = map[string]bool{
	"error":    true,
	"failed":   true,
	"failure":  true,
	"fatal":    true,
	"panic":    true,
	"request":  true,
	"range":    true,
	"bound":    true,
	"timed":    true,
	"timeout":  true,
	"exceeds":  true,
	"exceeded": true,
}

func diagnoseScoreReasonIsStrong(anchor string) bool {
	switch anchor {
	case "conn_id", "txn_start_ts", "region_id", "sql", "table_name", "error_text":
		return true
	default:
		return false
	}
}

func diagnoseStrongAnchorKey(reason ScoreReason) string {
	return reason.Anchor + "\x00" + reason.Value
}

func diagnoseHasReason(reasons []ScoreReason, anchor string) bool {
	for _, reason := range reasons {
		if reason.Anchor == anchor {
			return true
		}
	}
	return false
}

func diagnoseEvidenceTimeDelta(left, right Evidence) time.Duration {
	delta := left.Time.Sub(right.Time)
	if delta < 0 {
		delta = -delta
	}
	return delta
}

func sortDiagnoseIncidents(incidents []Incident) {
	sort.SliceStable(incidents, func(i, j int) bool {
		if incidents[i].Score != incidents[j].Score {
			return incidents[i].Score > incidents[j].Score
		}
		if incidents[i].EvidenceCount != incidents[j].EvidenceCount {
			return incidents[i].EvidenceCount > incidents[j].EvidenceCount
		}
		if !incidents[i].StartTime.Equal(incidents[j].StartTime) {
			return incidents[i].StartTime.Before(incidents[j].StartTime)
		}
		return stableDiagnoseIncidentAnchorKey(incidents[i]) < stableDiagnoseIncidentAnchorKey(incidents[j])
	})
}

func assignDiagnoseIncidentIDs(incidents []Incident) {
	for idx := range incidents {
		incidents[idx].ID = fmt.Sprintf("I%d", idx+1)
	}
}

func stableDiagnoseIncidentAnchorKey(incident Incident) string {
	tableNames := append([]string(nil), incident.Anchors.TableNames...)
	stackFrames := append([]string(nil), incident.Anchors.StackFrames...)
	sort.Strings(tableNames)
	sort.Strings(stackFrames)

	return strings.Join([]string{
		incident.Anchors.ConnID,
		incident.Anchors.TxnStartTS,
		incident.Anchors.RegionID,
		incident.Anchors.StoreAddr,
		incident.Anchors.Keyspace,
		incident.Anchors.SQL,
		strings.Join(tableNames, ","),
		incident.Anchors.ErrorText,
		strings.Join(stackFrames, ","),
	}, "\x1f")
}

func newDiagnoseUnionFind(size int) *diagnoseUnionFind {
	uf := &diagnoseUnionFind{
		parent: make([]int, size),
		size:   make([]int, size),
	}
	for idx := range uf.parent {
		uf.parent[idx] = idx
		uf.size[idx] = 1
	}
	return uf
}

func (uf *diagnoseUnionFind) find(index int) int {
	if uf.parent[index] != index {
		uf.parent[index] = uf.find(uf.parent[index])
	}
	return uf.parent[index]
}

func (uf *diagnoseUnionFind) union(leftRoot, rightRoot int) int {
	leftRoot = uf.find(leftRoot)
	rightRoot = uf.find(rightRoot)
	if leftRoot == rightRoot {
		return leftRoot
	}
	if uf.size[leftRoot] < uf.size[rightRoot] {
		leftRoot, rightRoot = rightRoot, leftRoot
	}
	uf.parent[rightRoot] = leftRoot
	uf.size[leftRoot] += uf.size[rightRoot]
	return leftRoot
}

func diagnoseMostFrequentValue(evidence []Evidence, getter func(Evidence) string) string {
	counts := make(map[string]int)
	for _, item := range evidence {
		if value := strings.TrimSpace(getter(item)); value != "" {
			counts[value]++
		}
	}
	return diagnoseMostFrequentString(counts)
}

func diagnoseMostFrequentTableNames(evidence []Evidence) []string {
	counts := make(map[string]int)
	for _, item := range evidence {
		seen := make(map[string]struct{})
		for _, table := range item.Anchors.TableNames {
			table = strings.TrimSpace(table)
			if table == "" {
				continue
			}
			if _, ok := seen[table]; ok {
				continue
			}
			seen[table] = struct{}{}
			counts[table]++
		}
	}
	return diagnoseMostFrequentStrings(counts)
}

func diagnoseMostFrequentStringsFromEvidence(evidence []Evidence, getter func(Evidence) []string) []string {
	counts := make(map[string]int)
	for _, item := range evidence {
		seen := make(map[string]struct{})
		for _, value := range getter(item) {
			value = strings.TrimSpace(value)
			if value == "" {
				continue
			}
			if _, ok := seen[value]; ok {
				continue
			}
			seen[value] = struct{}{}
			counts[value]++
		}
	}
	return diagnoseMostFrequentStrings(counts)
}

func diagnoseMostFrequentString(counts map[string]int) string {
	values := diagnoseMostFrequentValues(counts)
	if len(values) == 0 {
		return ""
	}
	return values[0]
}

func diagnoseMostFrequentStrings(counts map[string]int) []string {
	values := diagnoseMostFrequentValues(counts)
	if len(values) == 0 {
		return nil
	}
	return values
}

func diagnoseMostFrequentValues(counts map[string]int) []string {
	if len(counts) == 0 {
		return nil
	}
	maxCount := 0
	for _, count := range counts {
		if count > maxCount {
			maxCount = count
		}
	}

	values := make([]string, 0)
	for value, count := range counts {
		if count == maxCount {
			values = append(values, value)
		}
	}
	sort.Strings(values)
	return values
}