package agent

import (
	"context"
	"encoding/json"
	"testing"

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

func TestEnhanceBatchDefaultsAgentEdgesBeforePersistence(t *testing.T) {
	input := sampleGraph()
	runner := RunnerFunc(func(context.Context, schema.KnowledgeGraph, Request) (schema.KnowledgeGraph, error) {
		data, err := json.Marshal(input)
		require.NoError(t, err)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(data, &decoded))
		decoded["edges"] = []any{map[string]any{"source": input.Nodes[0].ID, "target": input.Nodes[0].ID, "type": "related"}}
		data, err = json.Marshal(decoded)
		require.NoError(t, err)
		var enhanced schema.KnowledgeGraph
		require.NoError(t, json.Unmarshal(data, &enhanced))
		return enhanced, nil
	})

	enhanced, ran, err := EnhanceBatch(context.Background(), input, Options{Enabled: true, Runner: runner}, Request{})

	require.NoError(t, err)
	require.True(t, ran)
	require.Equal(t, "forward", enhanced.Edges[len(enhanced.Edges)-1].Direction)
	require.Equal(t, 0.5, enhanced.Edges[len(enhanced.Edges)-1].Weight)
	_, err = json.Marshal(enhanced)
	require.NoError(t, err)
}

func TestEnhanceDisabledByDefaultIsNoop(t *testing.T) {
	graph := sampleGraph()
	enhanced, ran, err := Enhance(context.Background(), graph, Options{})
	require.NoError(t, err)
	require.False(t, ran)
	require.Equal(t, graph, enhanced)
}

func TestEnhanceWithFakeRunnerCanEnrichWithoutBreakingSchema(t *testing.T) {
	graph := sampleGraph()
	enhanced, ran, err := Enhance(context.Background(), graph, Options{
		Enabled:        true,
		OutputLanguage: "Simplified Chinese",
		Runner: RunnerFunc(func(ctx context.Context, graph schema.KnowledgeGraph, request Request) (schema.KnowledgeGraph, error) {
			require.Equal(t, "Simplified Chinese", request.OutputLanguage)
			graph.Nodes[0].Summary = "Agent-enhanced summary"
			graph.Nodes[0].Tags = append(graph.Nodes[0].Tags, "agent")
			return graph, nil
		}),
	})
	require.NoError(t, err)
	require.True(t, ran)
	require.Equal(t, "Agent-enhanced summary", enhanced.Nodes[0].Summary)
	require.Contains(t, enhanced.Nodes[0].Tags, "agent")
	result := schema.ValidateGraph(enhanced)
	require.True(t, result.Success, result.Fatal)
}

func TestEnhanceBatchPassesBatchContext(t *testing.T) {
	graph := sampleGraph()
	enhanced, ran, err := EnhanceBatch(context.Background(), graph, Options{
		Enabled:        true,
		OutputLanguage: "zh",
		Runner: RunnerFunc(func(ctx context.Context, graph schema.KnowledgeGraph, request Request) (schema.KnowledgeGraph, error) {
			require.Equal(t, "zh", request.OutputLanguage)
			require.Equal(t, 7, request.BatchIndex)
			require.Equal(t, 9, request.TotalBatches)
			require.Len(t, request.Files, 1)
			require.Equal(t, "main.go", request.Files[0].Path)
			require.Equal(t, map[string][]string{"main.go": {"dep.go"}}, request.BatchImportData)
			require.Contains(t, request.NeighborMap, "main.go")
			graph.Nodes[0].Summary = "Batch-aware summary"
			return graph, nil
		}),
	}, Request{
		BatchIndex:      7,
		TotalBatches:    9,
		Files:           []scanner.FileInfo{{Path: "main.go"}},
		BatchImportData: map[string][]string{"main.go": {"dep.go"}},
		NeighborMap:     map[string][]batching.Neighbor{"main.go": {{Path: "dep.go", BatchIndex: 2}}},
	})
	require.NoError(t, err)
	require.True(t, ran)
	require.Equal(t, "Batch-aware summary", enhanced.Nodes[0].Summary)
}

func TestEnhanceRejectsInvalidRunnerGraph(t *testing.T) {
	graph := sampleGraph()
	_, _, err := Enhance(context.Background(), graph, Options{
		Enabled: true,
		Runner: RunnerFunc(func(context.Context, schema.KnowledgeGraph, Request) (schema.KnowledgeGraph, error) {
			graph.Nodes = nil
			return graph, nil
		}),
	})
	require.Error(t, err)
	require.Contains(t, err.Error(), "invalid graph")
}

func sampleGraph() schema.KnowledgeGraph {
	return schema.KnowledgeGraph{
		Version: "1.0.0",
		Kind:    schema.KindCodebase,
		Project: schema.ProjectMeta{Name: "tiny", Languages: []string{"go"}, Frameworks: []string{}, Description: "tiny", AnalyzedAt: "2026-07-02T00:00:00Z", GitCommitHash: "abc"},
		Nodes: []schema.GraphNode{{
			ID:         "file:main.go",
			Type:       "file",
			Name:       "main.go",
			FilePath:   "main.go",
			Summary:    "main",
			Tags:       []string{"go"},
			Complexity: "simple",
		}},
		Edges:  []schema.GraphEdge{},
		Layers: []schema.Layer{},
		Tour:   []schema.TourStep{},
	}
}