package kb

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"context"
	"os"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
)

const testUASourceSHA = "58cfb20ac8f3f98cd7dede428d147dbe9cdc94b2"

func TestUADashboardArchiveAcceptsFormatsAndRoots(t *testing.T) {
	graph := testUAGraph(t)
	tests := []struct {
		name     string
		filename string
		format   string
		prefix   string
	}{
		{name: "zip ua root", filename: "dashboard.zip", format: "zip", prefix: ".ua/"},
		{name: "zip flat root", filename: "dashboard.zip", format: "zip"},
		{name: "tar gzip", filename: "dashboard.tar.gz", format: "tar.gz", prefix: ".ua/"},
		{name: "tgz", filename: "dashboard.tgz", format: "tar.gz"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			archivePath := writeUATestArchive(t, test.format, testUAArchiveEntries(t, test.prefix, graph))
			staged, err := StageUADashboardArchive(context.Background(), archivePath, test.filename, DefaultUADashboardImportLimits())
			require.NoError(t, err)
			require.Equal(t, test.format, staged.Format)
			require.Equal(t, []string{UADashboardConfigFile, UADashboardDomainGraphFile, UADashboardFingerprintsFile, UADashboardKnowledgeGraphFile, UADashboardMetaFile}, manifestPaths(staged.Manifest))
			require.NotEmpty(t, staged.Files[UADashboardKnowledgeGraphFile])
		})
	}
}

func TestUADashboardArchiveAcceptsArbitraryValidJSONAndRejectsSyntaxViolations(t *testing.T) {
	valid := [][]byte{[]byte(`{}`), []byte(`[]`), []byte(`"value"`), []byte(`42`), []byte(`null`)}
	entries := testUAArchiveEntries(t, "", valid[0])
	for index := range entries {
		entries[index].data = valid[index]
	}
	archivePath := writeUATestArchive(t, "zip", entries)
	_, err := StageUADashboardArchive(context.Background(), archivePath, "dashboard.zip", DefaultUADashboardImportLimits())
	require.NoError(t, err)

	for _, test := range []struct {
		name string
		data []byte
		want error
	}{
		{name: "BOM", data: append([]byte{0xef, 0xbb, 0xbf}, []byte(`{}`)...), want: ErrUADashboardDataInvalid},
		{name: "invalid UTF-8", data: []byte{0xff}, want: ErrUADashboardDataInvalid},
		{name: "malformed", data: []byte(`{`), want: ErrUADashboardDataInvalid},
		{name: "trailing value", data: []byte(`{} []`), want: ErrUADashboardDataInvalid},
	} {
		t.Run(test.name, func(t *testing.T) {
			entries := testUAArchiveEntries(t, "", []byte(`{}`))
			entries[0].data = test.data
			archivePath := writeUATestArchive(t, "zip", entries)
			_, err := StageUADashboardArchive(context.Background(), archivePath, "dashboard.zip", DefaultUADashboardImportLimits())
			require.ErrorIs(t, err, test.want)
		})
	}

	limits := DefaultUADashboardImportLimits()
	limits.JSONDepth = 2
	entries = testUAArchiveEntries(t, "", []byte(`{}`))
	entries[0].data = []byte(`[[[]]]`)
	archivePath = writeUATestArchive(t, "zip", entries)
	_, err = StageUADashboardArchive(context.Background(), archivePath, "dashboard.zip", limits)
	require.ErrorIs(t, err, ErrUADashboardArchiveLimit)
}

func TestUADashboardArchiveRejectsUnsafeAndInvalidLayouts(t *testing.T) {
	graph := testUAGraph(t)
	tests := []struct {
		name       string
		format     string
		entries    []uaTestArchiveEntry
		want       error
		postMutate func(*testing.T, string)
	}{
		{name: "wrapper", format: "zip", entries: []uaTestArchiveEntry{{name: "project/.ua/knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveLayout},
		{name: "legacy root", format: "zip", entries: []uaTestArchiveEntry{{name: ".understand-anything/knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveLayout},
		{name: "mixed roots", format: "tar.gz", entries: []uaTestArchiveEntry{{name: ".ua/knowledge-graph.json", data: graph}, {name: "domain-graph.json", data: testUADomainGraph(t)}}, want: ErrUADashboardArchiveLayout},
		{name: "empty ua plus flat exact five", format: "zip", entries: append([]uaTestArchiveEntry{{name: ".ua", typeflag: tar.TypeDir}}, testUAArchiveEntries(t, "", graph)...), want: ErrUADashboardArchiveLayout},
		{name: "missing primary graph", format: "zip", entries: []uaTestArchiveEntry{{name: "domain-graph.json", data: testUADomainGraph(t)}}, want: ErrUADashboardArchiveLayout},
		{name: "missing domain graph", format: "zip", entries: []uaTestArchiveEntry{{name: "knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveLayout},
		{name: "former diff", format: "zip", entries: append(testUAArchiveEntries(t, "", graph), uaTestArchiveEntry{name: "diff-overlay.json", data: []byte(`{}`)}), want: ErrUADashboardArchiveLayout},
		{name: "unknown file", format: "zip", entries: append(testUAArchiveEntries(t, "", graph), uaTestArchiveEntry{name: "intermediate/cache.json", data: []byte(`{}`)}), want: ErrUADashboardArchiveLayout},
		{name: "traversal", format: "zip", entries: []uaTestArchiveEntry{{name: "../knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveUnsafe},
		{name: "backslash", format: "zip", entries: []uaTestArchiveEntry{{name: `.ua\knowledge-graph.json`, data: graph}}, want: ErrUADashboardArchiveUnsafe},
		{name: "percent ambiguity", format: "tar.gz", entries: []uaTestArchiveEntry{{name: "%2e%2e/knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveUnsafe},
		{name: "duplicate", format: "tar.gz", entries: []uaTestArchiveEntry{{name: "knowledge-graph.json", data: graph}, {name: "knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveUnsafe},
		{name: "symlink", format: "tar.gz", entries: []uaTestArchiveEntry{{name: "knowledge-graph.json", data: graph}, {name: "domain-graph.json", typeflag: tar.TypeSymlink, linkname: "knowledge-graph.json"}}, want: ErrUADashboardArchiveUnsafe},
		{name: "trailing zip", format: "zip", entries: []uaTestArchiveEntry{{name: "knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveUnsafe, postMutate: func(t *testing.T, path string) {
			file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0)
			require.NoError(t, err)
			_, err = file.WriteString("trailing")
			require.NoError(t, err)
			require.NoError(t, file.Close())
		}},
		{name: "trailing gzip", format: "tar.gz", entries: []uaTestArchiveEntry{{name: "knowledge-graph.json", data: graph}}, want: ErrUADashboardArchiveUnsafe, postMutate: func(t *testing.T, path string) {
			file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0)
			require.NoError(t, err)
			_, err = file.WriteString("trailing")
			require.NoError(t, err)
			require.NoError(t, file.Close())
		}},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			archivePath := writeUATestArchive(t, test.format, test.entries)
			if test.postMutate != nil {
				test.postMutate(t, archivePath)
			}
			_, err := StageUADashboardArchive(context.Background(), archivePath, "dashboard."+map[string]string{"zip": "zip", "tar.gz": "tgz"}[test.format], DefaultUADashboardImportLimits())
			require.ErrorIs(t, err, test.want)
		})
	}
}

func TestUADashboardArchiveEnforcesResourceBoundsAndCancellation(t *testing.T) {
	graph := testUAGraph(t)
	tests := []struct {
		name   string
		limits func(UADashboardImportLimits) UADashboardImportLimits
	}{
		{name: "compressed", limits: func(l UADashboardImportLimits) UADashboardImportLimits { l.HTTPBytes = 4; return l }},
		{name: "entries", limits: func(l UADashboardImportLimits) UADashboardImportLimits { l.Entries = 1; return l }},
		{name: "one file", limits: func(l UADashboardImportLimits) UADashboardImportLimits { l.FileBytes = 32; return l }},
		{name: "extracted", limits: func(l UADashboardImportLimits) UADashboardImportLimits { l.ExtractedBytes = 32; return l }},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			archivePath := writeUATestArchive(t, "zip", testUAArchiveEntries(t, "", graph))
			_, err := StageUADashboardArchive(context.Background(), archivePath, "dashboard.zip", test.limits(DefaultUADashboardImportLimits()))
			require.ErrorIs(t, err, ErrUADashboardArchiveLimit)
		})
	}
	archivePath := writeUATestArchive(t, "zip", []uaTestArchiveEntry{{name: "knowledge-graph.json", data: graph}})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := StageUADashboardArchive(ctx, archivePath, "dashboard.zip", DefaultUADashboardImportLimits())
	require.ErrorIs(t, err, context.Canceled)
}

func TestDefaultUADashboardImportLimitsAllow128MiBFile(t *testing.T) {
	require.EqualValues(t, 128<<20, DefaultUADashboardImportLimits().FileBytes)
}

func TestUADashboardArchiveUsesMagicAndCompatibleFilename(t *testing.T) {
	archivePath := writeUATestArchive(t, "zip", []uaTestArchiveEntry{{name: "knowledge-graph.json", data: testUAGraph(t)}})
	_, err := StageUADashboardArchive(context.Background(), archivePath, "dashboard.tgz", DefaultUADashboardImportLimits())
	require.ErrorIs(t, err, ErrUADashboardArchiveUnsupported)
	require.NoError(t, os.WriteFile(archivePath, []byte("not an archive"), 0o600))
	_, err = StageUADashboardArchive(context.Background(), archivePath, "dashboard.zip", DefaultUADashboardImportLimits())
	require.ErrorIs(t, err, ErrUADashboardArchiveUnsupported)
}

func TestUADashboardArchiveRejectsEncryptedZIP(t *testing.T) {
	archivePath := writeUATestArchive(t, "zip", []uaTestArchiveEntry{{name: "knowledge-graph.json", data: testUAGraph(t)}})
	corruptZipEncryptedFlag(t, archivePath)
	_, err := StageUADashboardArchive(context.Background(), archivePath, "dashboard.zip", DefaultUADashboardImportLimits())
	require.ErrorIs(t, err, ErrUADashboardArchiveUnsafe)
}

type uaTestArchiveEntry struct {
	name     string
	data     []byte
	typeflag byte
	linkname string
}

func writeUATestArchive(t *testing.T, format string, entries []uaTestArchiveEntry) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "dashboard.archive")
	file, err := os.Create(path)
	require.NoError(t, err)
	switch format {
	case "zip":
		writer := zip.NewWriter(file)
		for _, entry := range entries {
			header := &zip.FileHeader{Name: entry.name, Method: zip.Deflate}
			if entry.typeflag == tar.TypeDir {
				header.Name += "/"
			}
			part, err := writer.CreateHeader(header)
			require.NoError(t, err)
			_, err = part.Write(entry.data)
			require.NoError(t, err)
		}
		require.NoError(t, writer.Close())
	case "tar.gz":
		gzipWriter := gzip.NewWriter(file)
		writer := tar.NewWriter(gzipWriter)
		for _, entry := range entries {
			typeflag := entry.typeflag
			if typeflag == 0 {
				typeflag = tar.TypeReg
			}
			header := &tar.Header{Name: entry.name, Mode: 0o600, Typeflag: typeflag, Linkname: entry.linkname}
			if typeflag == tar.TypeReg || typeflag == tar.TypeRegA {
				header.Size = int64(len(entry.data))
			}
			require.NoError(t, writer.WriteHeader(header))
			if header.Size > 0 {
				_, err := writer.Write(entry.data)
				require.NoError(t, err)
			}
		}
		require.NoError(t, writer.Close())
		require.NoError(t, gzipWriter.Close())
	default:
		t.Fatalf("unknown archive format %q", format)
	}
	require.NoError(t, file.Close())
	return path
}

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

func testUAGraph(t *testing.T) []byte {
	t.Helper()
	data, err := os.ReadFile("testdata/ua/pinned-codebase/knowledge-graph.json")
	require.NoError(t, err)
	return data
}

func testUADomainGraph(t *testing.T) []byte {
	t.Helper()
	data, err := os.ReadFile("testdata/ua/pinned-codebase/domain-graph.json")
	require.NoError(t, err)
	return data
}

func testUAArchiveEntries(t *testing.T, prefix string, graph []byte) []uaTestArchiveEntry {
	t.Helper()
	return []uaTestArchiveEntry{
		{name: prefix + UADashboardKnowledgeGraphFile, data: graph},
		{name: prefix + UADashboardDomainGraphFile, data: testUADomainGraph(t)},
		{name: prefix + UADashboardMetaFile, data: []byte(`{}`)},
		{name: prefix + UADashboardConfigFile, data: []byte(`{}`)},
		{name: prefix + UADashboardFingerprintsFile, data: []byte(`{}`)},
	}
}

func corruptZipEncryptedFlag(t *testing.T, archivePath string) {
	t.Helper()
	data, err := os.ReadFile(archivePath)
	require.NoError(t, err)
	for _, record := range []struct {
		signature  []byte
		flagOffset int
	}{{signature: []byte{'P', 'K', 3, 4}, flagOffset: 6}, {signature: []byte{'P', 'K', 1, 2}, flagOffset: 8}} {
		for offset := 0; ; {
			index := bytes.Index(data[offset:], record.signature)
			if index < 0 {
				break
			}
			index += offset
			flagOffset := index + record.flagOffset
			data[flagOffset] |= 1
			offset = index + 4
		}
	}
	require.NoError(t, os.WriteFile(archivePath, data, 0o600))
}