package ingest

import (
	"encoding/json"
	"os"
	"testing"

	"github.com/agentic-research/mache/api"
	"github.com/agentic-research/mache/internal/graph"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func loadSchema(t *testing.T, path string) *api.Topology {
	t.Helper()
	data, err := os.ReadFile(path)
	require.NoError(t, err)
	var topo api.Topology
	require.NoError(t, json.Unmarshal(data, &topo))
	return &topo
}

func TestEngine_IngestSQLite_KEV(t *testing.T) {
	dbPath := createTestDB(t, []string{
		`{"schema":"kev","identifier":"CVE-2024-0001","item":{"cveID":"CVE-2024-0001","vendorProject":"Acme Corp","product":"Widget","shortDescription":"Remote code execution in Widget","dateAdded":"2024-01-15","dueDate":"2024-02-15"}}`,
		`{"schema":"kev","identifier":"CVE-2024-0002","item":{"cveID":"CVE-2024-0002","vendorProject":"Globex","product":"Gizmo","shortDescription":"SQL injection in Gizmo","dateAdded":"2024-02-01","dueDate":"2024-03-01"}}`,
		`{"schema":"kev","identifier":"CVE-2024-0003","item":{"cveID":"CVE-2024-0003","vendorProject":"Initech","product":"TPS Report","shortDescription":"XSS in TPS Report viewer","dateAdded":"2024-03-01","dueDate":"2024-04-01"}}`,
	})

	schema := loadSchema(t, "../../examples/kev-schema.json")
	store := graph.NewMemoryStore()
	engine := NewEngine(schema, store)

	require.NoError(t, engine.Ingest(dbPath))

	// vulns root dir exists with 3 children
	vulns, err := store.GetNode("vulns")
	require.NoError(t, err)
	assert.True(t, vulns.Mode.IsDir())
	assert.Len(t, vulns.Children, 3)

	// Spot-check first CVE files
	vendor, err := store.GetNode("vulns/CVE-2024-0001/vendor")
	require.NoError(t, err)
	assert.Equal(t, "Acme Corp", string(vendor.Data))

	product, err := store.GetNode("vulns/CVE-2024-0001/product")
	require.NoError(t, err)
	assert.Equal(t, "Widget", string(product.Data))

	desc, err := store.GetNode("vulns/CVE-2024-0001/description")
	require.NoError(t, err)
	assert.Equal(t, "Remote code execution in Widget", string(desc.Data))

	// raw.json is valid JSON containing the full record
	raw, err := store.GetNode("vulns/CVE-2024-0001/raw.json")
	require.NoError(t, err)
	var rawData map[string]any
	require.NoError(t, json.Unmarshal(raw.Data, &rawData))
	assert.Equal(t, "CVE-2024-0001", rawData["identifier"])
}

func TestEngine_IngestSQLite_NVD(t *testing.T) {
	dbPath := createTestDB(t, []string{
		`{"schema":"nvd","identifier":"CVE-2024-0001","item":{"cve":{"id":"CVE-2024-0001","descriptions":[{"lang":"en","value":"A buffer overflow in FooBar allows remote code execution."}],"published":"2024-01-15T00:00:00Z","vulnStatus":"Analyzed"}}}`,
		`{"schema":"nvd","identifier":"CVE-2024-0002","item":{"cve":{"id":"CVE-2024-0002","descriptions":[{"lang":"en","value":"An authentication bypass in BazQux."}],"published":"2024-02-01T00:00:00Z","vulnStatus":"Modified"}}}`,
		`{"schema":"nvd","identifier":"CVE-2023-0001","item":{"cve":{"id":"CVE-2023-0001","descriptions":[{"lang":"en","value":"A null pointer dereference in Quux."}],"published":"2023-06-01T00:00:00Z","vulnStatus":"Analyzed"}}}`,
	})

	schema := loadSchema(t, "../../examples/nvd-schema.json")
	store := graph.NewMemoryStore()
	engine := NewEngine(schema, store)

	require.NoError(t, engine.Ingest(dbPath))

	// by-cve root dir exists with year-level children
	byCve, err := store.GetNode("by-cve")
	require.NoError(t, err)
	assert.True(t, byCve.Mode.IsDir())
	assert.Len(t, byCve.Children, 2) // "2024" and "2023"

	// Year directory has month children
	year2024, err := store.GetNode("by-cve/2024")
	require.NoError(t, err)
	assert.True(t, year2024.Mode.IsDir())
	assert.Len(t, year2024.Children, 2) // "01" and "02"

	// Month directory has CVE children
	month01, err := store.GetNode("by-cve/2024/01")
	require.NoError(t, err)
	assert.True(t, month01.Mode.IsDir())
	assert.Len(t, month01.Children, 1) // CVE-2024-0001

	month02, err := store.GetNode("by-cve/2024/02")
	require.NoError(t, err)
	assert.Len(t, month02.Children, 1) // CVE-2024-0002

	year2023, err := store.GetNode("by-cve/2023")
	require.NoError(t, err)
	assert.Len(t, year2023.Children, 1) // "06"

	month06, err := store.GetNode("by-cve/2023/06")
	require.NoError(t, err)
	assert.Len(t, month06.Children, 1) // CVE-2023-0001

	// description uses nested template with index (now under year/month dir)
	desc, err := store.GetNode("by-cve/2024/01/CVE-2024-0001/description")
	require.NoError(t, err)
	assert.Equal(t, "A buffer overflow in FooBar allows remote code execution.", string(desc.Data))

	// status from nested path
	status, err := store.GetNode("by-cve/2024/01/CVE-2024-0001/status")
	require.NoError(t, err)
	assert.Equal(t, "Analyzed", string(status.Data))

	// raw.json is valid JSON
	raw, err := store.GetNode("by-cve/2024/01/CVE-2024-0001/raw.json")
	require.NoError(t, err)
	var rawData map[string]any
	require.NoError(t, json.Unmarshal(raw.Data, &rawData))
	item := rawData["item"].(map[string]any)
	cve := item["cve"].(map[string]any)
	assert.Equal(t, "CVE-2024-0001", cve["id"])

	// Cross-year: 2023 CVE also works
	desc2023, err := store.GetNode("by-cve/2023/06/CVE-2023-0001/description")
	require.NoError(t, err)
	assert.Equal(t, "A null pointer dereference in Quux.", string(desc2023.Data))
}

func TestEngine_SchemaRefs(t *testing.T) {
	// Schema with refs: each tool node emits a ref token for its tool name.
	// This tests the cross-reference feature for JSON-projected data.
	schema := &api.Topology{
		Version: "v1",
		Nodes: []api.Node{
			{
				Name:     "servers",
				Selector: "$",
				Children: []api.Node{
					{
						Name:     "{{.item.name}}",
						Selector: "$[*]",
						Children: []api.Node{
							{
								Name:     "tools",
								Selector: "$",
								Children: []api.Node{
									{
										Name:     "{{.name}}",
										Selector: "$.item.tools[*]",
										Refs:     []string{"{{.name}}"},
										Files: []api.Leaf{
											{Name: "description", ContentTemplate: "{{.description}}"},
										},
									},
								},
							},
						},
					},
				},
			},
		},
	}

	dbPath := createTestDB(t, []string{
		`{"schema":"mcp","identifier":"server-a","item":{"name":"server-a","tools":[{"name":"search","description":"Search things"},{"name":"create","description":"Create things"}]}}`,
		`{"schema":"mcp","identifier":"server-b","item":{"name":"server-b","tools":[{"name":"search","description":"Search other things"},{"name":"delete","description":"Delete things"}]}}`,
	})

	store := graph.NewMemoryStore()
	require.NoError(t, store.InitRefsDB())
	defer func() { _ = store.Close() }()
	engine := NewEngine(schema, store)

	require.NoError(t, engine.Ingest(dbPath))
	require.NoError(t, store.FlushRefs())

	// Both servers ingested
	_, err := store.GetNode("servers/server-a")
	require.NoError(t, err)
	_, err = store.GetNode("servers/server-b")
	require.NoError(t, err)

	// Tool nodes exist under each server
	_, err = store.GetNode("servers/server-a/tools/search")
	require.NoError(t, err)
	_, err = store.GetNode("servers/server-b/tools/search")
	require.NoError(t, err)

	// Refs: "search" token should have callers from both servers
	callers, err := store.GetCallers("search")
	require.NoError(t, err)
	assert.Len(t, callers, 2, "search should have callers from 2 servers")

	// Refs: "create" token should have 1 caller
	callers, err = store.GetCallers("create")
	require.NoError(t, err)
	assert.Len(t, callers, 1, "create should have 1 caller")

	// Refs: "delete" token should have 1 caller
	callers, err = store.GetCallers("delete")
	require.NoError(t, err)
	assert.Len(t, callers, 1, "delete should have 1 caller")
}
