package batching

import (
	"fmt"
	"os"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/extractor"
	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/importgraph"
	"github.com/tidbcloud/AgentTeam/kits/agentteam-insight/internal/graphbuilder/scanner"
)

func TestBuildCreatesDeterministicSemanticBatches(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "docs/readme.md", Language: "markdown", Size: 10},
		{Path: "cmd/app/main.go", Language: "go", Size: 12},
		{Path: "internal/util/util.go", Language: "go", Size: 8},
		{Path: "web/app.ts", Language: "typescript", Size: 5},
	}
	imports := importgraph.Graph{ImportMap: map[string][]importgraph.ImportEdge{
		"cmd/app/main.go": {{SourcePath: "cmd/app/main.go", TargetPath: "internal/util/util.go", ImportPath: "example.com/acme/internal/util"}},
	}}

	first := Build(files, imports, Options{MaxBatchBytes: 100, MaxBatchFiles: 10})
	firstPlan := BuildPlan(files, imports, Options{MaxBatchBytes: 100, MaxBatchFiles: 10})
	second := Build([]scanner.FileInfo{files[3], files[1], files[0], files[2]}, imports, Options{MaxBatchBytes: 100, MaxBatchFiles: 10})

	require.Equal(t, first, second)
	require.Equal(t, AlgorithmLouvain, firstPlan.Algorithm)
	require.Len(t, first, 1)
	require.Equal(t, "batch-1.json", first[0].Name)
	require.Equal(t, []string{"cmd/app/main.go", "docs/readme.md", "internal/util/util.go", "web/app.ts"}, batchPaths(first[0]))
	require.Equal(t, []string{"go", "markdown", "typescript"}, first[0].Languages)
	require.Empty(t, first[0].NeighborMap)
	require.Equal(t, map[string][]string{"cmd/app/main.go": {"internal/util/util.go"}, "docs/readme.md": {}, "internal/util/util.go": {}, "web/app.ts": {}}, first[0].BatchImportData)
}

func TestBuildUsesLouvainCommunitiesForDenseImportGraph(t *testing.T) {
	files := make([]scanner.FileInfo, 0, 50)
	importMap := map[string][]importgraph.ImportEdge{}
	for community := 0; community < 2; community++ {
		for i := 0; i < 25; i++ {
			path := communityFileName(community, i)
			files = append(files, scanner.FileInfo{Path: path, Language: "go", Category: scanner.CategoryCode, Size: 1})
		}
	}
	for community := 0; community < 2; community++ {
		paths := []string{}
		for i := 0; i < 25; i++ {
			paths = append(paths, communityFileName(community, i))
		}
		for _, source := range paths {
			for _, target := range paths {
				if source == target {
					continue
				}
				importMap[source] = append(importMap[source], importgraph.ImportEdge{SourcePath: source, TargetPath: target, ImportPath: target})
			}
		}
	}
	bridgeA := communityFileName(0, 0)
	bridgeB := communityFileName(1, 0)
	importMap[bridgeA] = append(importMap[bridgeA], importgraph.ImportEdge{SourcePath: bridgeA, TargetPath: bridgeB, ImportPath: bridgeB})

	plan := BuildPlan(files, importgraph.Graph{ImportMap: importMap}, Options{MaxBatchFiles: 100, MaxBatchBytes: 10_000})

	require.Equal(t, AlgorithmLouvain, plan.Algorithm)
	require.Len(t, plan.Batches, 2)
	require.ElementsMatch(t, []int{25, 25}, []int{len(plan.Batches[0].Files), len(plan.Batches[1].Files)})
	for _, batch := range plan.Batches {
		require.LessOrEqual(t, len(batch.Files), MaxCommunityFiles)
		requireSingleCommunitySuffix(t, batch)
	}
}

func TestCountBasedCommunitiesFallbackChunksAtTwelve(t *testing.T) {
	paths := []string{}
	for i := 14; i >= 0; i-- {
		paths = append(paths, filepath.ToSlash(filepath.Join("pkg", fileName(i))))
	}

	communities := countBasedCommunities(paths, 12)

	require.Equal(t, [][]string{
		{
			"pkg/file00.go", "pkg/file01.go", "pkg/file02.go", "pkg/file03.go",
			"pkg/file04.go", "pkg/file05.go", "pkg/file06.go", "pkg/file07.go",
			"pkg/file08.go", "pkg/file09.go", "pkg/file10.go", "pkg/file11.go",
		},
		{"pkg/file12.go", "pkg/file13.go", "pkg/file14.go"},
	}, communities)
}

func TestBuildSplitsBySizeAndFileCount(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "a.go", Language: "go", Size: 20},
		{Path: "b.go", Language: "go", Size: 20},
		{Path: "c.go", Language: "go", Size: 20},
	}
	batches := Build(files, importgraph.Graph{}, Options{MaxBatchBytes: 35, MaxBatchFiles: 2})
	require.Len(t, batches, 3)
	require.Equal(t, []string{"a.go"}, batchPaths(batches[0]))
	require.Equal(t, []string{"b.go"}, batchPaths(batches[1]))
	require.Equal(t, []string{"c.go"}, batchPaths(batches[2]))
}

func TestBuildSplitsByEstimatedSemanticWork(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "a.go", Language: "go", Category: scanner.CategoryCode, Size: 10, Lines: 800},
		{Path: "b.go", Language: "go", Category: scanner.CategoryCode, Size: 10, Lines: 800},
		{Path: "c.go", Language: "go", Category: scanner.CategoryCode, Size: 10, Lines: 20},
	}
	imports := importgraph.Graph{ImportMap: map[string][]importgraph.ImportEdge{
		"a.go": {{SourcePath: "a.go", TargetPath: "b.go"}},
		"b.go": {{SourcePath: "b.go", TargetPath: "c.go"}},
	}}

	batches := Build(files, imports, Options{MaxBatchBytes: 1_000, MaxBatchFiles: 10, MaxBatchWork: 45})

	require.Len(t, batches, 2)
	require.Equal(t, []string{"a.go"}, batchPaths(batches[0]))
	require.Equal(t, []string{"b.go", "c.go"}, batchPaths(batches[1]))
	for _, batch := range batches {
		require.LessOrEqual(t, batch.EstimatedWork, 45)
	}
}

func TestBuildEstimatesWorkFromSignificantEvidence(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "a.go", Language: "go", Category: scanner.CategoryCode, Size: 10, Lines: 2_000},
		{Path: "b.go", Language: "go", Category: scanner.CategoryCode, Size: 10, Lines: 20},
	}
	imports := importgraph.Graph{
		Analyses: map[string]extractor.StructuralAnalysis{
			"a.go": {FilePath: "a.go", HasStructuralAnalysis: true, Functions: []extractor.Function{
				{Name: "small", LineRange: [2]int{1, 2}},
				{Name: "Exported", Exported: true, LineRange: [2]int{3, 4}},
			}},
			"b.go": {FilePath: "b.go", HasStructuralAnalysis: true, Functions: []extractor.Function{
				{Name: "large", LineRange: [2]int{1, 20}},
			}},
		},
		ImportMap: map[string][]importgraph.ImportEdge{
			"a.go": {{SourcePath: "a.go", TargetPath: "b.go"}},
		},
	}

	batches := Build(files, imports, Options{MaxBatchBytes: 1_000, MaxBatchFiles: 10, MaxBatchWork: 10})

	require.Len(t, batches, 1)
	// Each file contributes one file node and one significant function. The
	// private two-line helper and the raw 2,000-line file size do not inflate it.
	require.Equal(t, 4, batches[0].EstimatedWork)
}

func TestBuildAddsCrossBatchNeighborMapWithStableBatchIndex(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "a.go", Language: "go", Size: 40},
		{Path: "web/b.ts", Language: "typescript", Size: 40},
	}
	imports := importgraph.Graph{
		Analyses: map[string]extractor.StructuralAnalysis{
			"web/b.ts": {FilePath: "web/b.ts", Exports: []string{"B"}},
		},
		ImportMap: map[string][]importgraph.ImportEdge{
			"a.go": {{SourcePath: "a.go", TargetPath: "web/b.ts", ImportPath: "example.com/b"}},
		},
	}
	batches := Build(files, imports, Options{MaxBatchBytes: 50, MaxBatchFiles: 1})
	require.Len(t, batches, 2)
	require.Equal(t, []Neighbor{{Path: "web/b.ts", BatchIndex: 2, Symbols: []string{"B"}}}, batches[0].NeighborMap["a.go"])
	require.Equal(t, []Neighbor{{Path: "a.go", BatchIndex: 1, Symbols: []string{}}}, batches[1].NeighborMap["web/b.ts"])
}

func TestBatchFileNameContract(t *testing.T) {
	require.Equal(t, "batch-3.json", BatchFileName(3, 0))
	require.Equal(t, "batch-3-part-2.json", BatchFileName(3, 2))
}

func TestBuildGroupsNonCodeFilesAndPoolsSmallMergeableBatches(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "deploy/app/Dockerfile", Category: scanner.CategoryInfra, Language: "docker", Size: 1},
		{Path: "deploy/app/.dockerignore", Category: scanner.CategoryConfig, Language: "dockerignore", Size: 1},
		{Path: "deploy/app/docker-compose.yml", Category: scanner.CategoryConfig, Language: "yaml", Size: 1},
		{Path: ".github/workflows/ci.yml", Category: scanner.CategoryInfra, Language: "github-actions", Size: 1},
		{Path: ".github/workflows/release.yaml", Category: scanner.CategoryInfra, Language: "github-actions", Size: 1},
		{Path: ".gitlab-ci.yml", Category: scanner.CategoryInfra, Language: "yaml", Size: 1},
		{Path: ".circleci/config.yml", Category: scanner.CategoryInfra, Language: "yaml", Size: 1},
		{Path: "db/migrations/001_init.sql", Category: scanner.CategoryData, Language: "sql", Size: 1},
		{Path: "db/migrations/002_add.sql", Category: scanner.CategoryData, Language: "sql", Size: 1},
		{Path: "docs/a.md", Category: scanner.CategoryDocs, Language: "markdown", Size: 1},
		{Path: "docs/b.md", Category: scanner.CategoryDocs, Language: "markdown", Size: 1},
		{Path: "configs/app.yaml", Category: scanner.CategoryConfig, Language: "yaml", Size: 1},
	}

	batches := Build(files, importgraph.Graph{}, Options{MaxBatchFiles: 25})
	require.Len(t, batches, 5)
	require.Equal(t, []string{"deploy/app/.dockerignore", "deploy/app/Dockerfile", "deploy/app/docker-compose.yml"}, batchPaths(batches[0]))
	require.Equal(t, []string{".github/workflows/ci.yml", ".github/workflows/release.yaml"}, batchPaths(batches[1]))
	require.Equal(t, []string{".circleci/config.yml", ".gitlab-ci.yml"}, batchPaths(batches[2]))
	require.Equal(t, []string{"db/migrations/001_init.sql", "db/migrations/002_add.sql"}, batchPaths(batches[3]))
	require.Equal(t, []string{"configs/app.yaml", "docs/a.md", "docs/b.md"}, batchPaths(batches[4]))
}

func TestFilterChangedFilesRetainsFullGraphBatchIndexAndNeighborMap(t *testing.T) {
	files := []scanner.FileInfo{
		{Path: "a.go", Language: "go", Category: scanner.CategoryCode, Size: 40},
		{Path: "web/b.ts", Language: "typescript", Category: scanner.CategoryCode, Size: 40},
	}
	imports := importgraph.Graph{
		Analyses: map[string]extractor.StructuralAnalysis{
			"a.go":     {FilePath: "a.go", Exports: []string{"A"}},
			"web/b.ts": {FilePath: "web/b.ts", Exports: []string{"B"}},
		},
		ImportMap: map[string][]importgraph.ImportEdge{
			"a.go": {{SourcePath: "a.go", TargetPath: "web/b.ts", ImportPath: "example.com/b"}},
		},
	}
	batches := Build(files, imports, Options{MaxBatchBytes: 50, MaxBatchFiles: 1})
	require.Len(t, batches, 2)

	filtered := FilterChangedFiles(batches, []string{"web/b.ts"})
	require.Len(t, filtered, 1)
	require.Equal(t, 2, filtered[0].BatchIndex)
	require.Equal(t, []string{"web/b.ts"}, batchPaths(filtered[0]))
	require.Equal(t, []Neighbor{{Path: "a.go", BatchIndex: 1, Symbols: []string{"A"}}}, filtered[0].NeighborMap["web/b.ts"])
}

func TestVerifyOutputsRequiresExpectedBatchFileOrPart(t *testing.T) {
	dir := t.TempDir()
	batches := []Batch{{BatchIndex: 1}, {BatchIndex: 2}}
	require.NoError(t, os.WriteFile(filepath.Join(dir, "batch-1.json"), []byte(`{"nodes":[]}`), 0o644))
	require.NoError(t, os.WriteFile(filepath.Join(dir, "batch-2-part-1.json"), []byte(`{"nodes":[]}`), 0o644))
	require.NoError(t, VerifyOutputs(dir, batches))

	require.NoError(t, os.Remove(filepath.Join(dir, "batch-2-part-1.json")))
	err := VerifyOutputs(dir, batches)
	require.Error(t, err)
	require.Contains(t, err.Error(), "incomplete-batches")
}

func batchPaths(batch Batch) []string {
	paths := make([]string, 0, len(batch.Files))
	for _, file := range batch.Files {
		paths = append(paths, file.Path)
	}
	return paths
}

func fileName(index int) string {
	return fmt.Sprintf("file%02d.go", index)
}

func communityFileName(community int, index int) string {
	return filepath.ToSlash(filepath.Join("pkg", fmt.Sprintf("service-%02d", index), fmt.Sprintf("file_%c.go", rune('a'+community))))
}

func requireSingleCommunitySuffix(t *testing.T, batch Batch) {
	t.Helper()
	suffix := ""
	for _, file := range batch.Files {
		name := filepath.Base(filepath.ToSlash(file.Path))
		if suffix == "" {
			suffix = name
			continue
		}
		require.Equal(t, suffix, name)
	}
}