package main

import (
	"bytes"
	"fmt"
	"io"
	"path/filepath"
	"reflect"
	"sort"
	"strings"
	"testing"

	"android/soong/third_party/zip"
)

// TestProcessAndSplit_BasicSplit tests a simple scenario.
func TestProcessAndSplit_BasicSplit(t *testing.T) {
	zip1Entries := []ZipEntry{
		{Name: "file1.txt", Content: "content1"},
		{Name: "file2.txt", Content: "content2"},
	}
	zip1SoongReader, err := createInMemorySoongZip(zip1Entries)
	if err != nil {
		t.Fatalf("Failed to create zip1: %v", err)
	}
	zip2Entries := []ZipEntry{
		{Name: "file3.txt", Content: "content3"},
		{Name: "file4.txt", Content: "content4"},
	}
	zip2SoongReader, err := createInMemorySoongZip(zip2Entries)
	if err != nil {
		t.Fatalf("Failed to create zip2: %v", err)
	}
	readers := []zipReader{
		{reader: zip1SoongReader, zipPath: "zip1.zip"}, // reader field expects *zip.Reader
		{reader: zip2SoongReader, zipPath: "zip2.zip"},
	}
	writer1, writer1Buf := createTestAppZipWriter()
	writer2, writer2Buf := createTestAppZipWriter()
	writers := []zipWriter{writer1, writer2}

	processAndSplit(readers, writers)
	writer1.writer.Close()
	writer2.writer.Close()

	expectedSplit1Files := []string{"file1.txt", "file2.txt"}
	expectedSplit2Files := []string{"file3.txt", "file4.txt"}

	actualSplit1Files, err := getFilenamesFromSoongZipBuffer(writer1Buf)
	if err != nil {
		t.Fatalf("Error reading split1 output: %v", err)
	}
	actualSplit2Files, err := getFilenamesFromSoongZipBuffer(writer2Buf)
	if err != nil {
		t.Fatalf("Error reading split2 output: %v", err)
	}

	if !reflect.DeepEqual(actualSplit1Files, expectedSplit1Files) {
		t.Errorf("Split 1 files mismatch:\nExpected: %v\nActual:   %v", expectedSplit1Files, actualSplit1Files)
	}
	if !reflect.DeepEqual(actualSplit2Files, expectedSplit2Files) {
		t.Errorf("Split 2 files mismatch:\nExpected: %v\nActual:   %v", expectedSplit2Files, actualSplit2Files)
	}
}

// TestProcessAndSplit_SingleOutput tests splitting into a single output zip.
func TestProcessAndSplit_SingleOutput(t *testing.T) {
	zipEntries := []ZipEntry{
		{Name: "a.txt", Content: "a"},
		{Name: "b.txt", Content: "b"},
	}
	zip1SoongReader, _ := createInMemorySoongZip(zipEntries)
	readers := []zipReader{{reader: zip1SoongReader, zipPath: "zip1.zip"}}
	writer, writerBuf := createTestAppZipWriter()
	writers := []zipWriter{writer}

	processAndSplit(readers, writers)
	writer.writer.Close()

	expectedFiles := []string{"a.txt", "b.txt"}
	actualFiles, _ := getFilenamesFromSoongZipBuffer(writerBuf)
	if !reflect.DeepEqual(actualFiles, expectedFiles) {
		t.Errorf("Single output files mismatch:\nExpected: %v\nActual:   %v", expectedFiles, actualFiles)
	}
}

// TestProcessAndSplit_UnevenSplit tests splitting when files don't divide evenly.
func TestProcessAndSplit_UnevenSplit(t *testing.T) {
	zipEntries := []ZipEntry{
		{Name: "f1.txt", Content: "1"},
		{Name: "f2.txt", Content: "2"},
		{Name: "f3.txt", Content: "3"},
		{Name: "f4.txt", Content: "4"},
		{Name: "f5.txt", Content: "5"},
	}
	soongReader, _ := createInMemorySoongZip(zipEntries)
	readers := []zipReader{{reader: soongReader, zipPath: "zip.zip"}}
	writer1, writer1Buf := createTestAppZipWriter()
	writer2, writer2Buf := createTestAppZipWriter()
	writers := []zipWriter{writer1, writer2}

	processAndSplit(readers, writers)
	writer1.writer.Close()
	writer2.writer.Close()

	expectedSplit1 := []string{"f1.txt", "f2.txt", "f3.txt"}
	expectedSplit2 := []string{"f4.txt", "f5.txt"}

	actualSplit1, _ := getFilenamesFromSoongZipBuffer(writer1Buf)
	actualSplit2, _ := getFilenamesFromSoongZipBuffer(writer2Buf)

	if !reflect.DeepEqual(actualSplit1, expectedSplit1) {
		t.Errorf("Uneven Split 1 files mismatch:\nExpected: %v\nActual:   %v", expectedSplit1, actualSplit1)
	}
	if !reflect.DeepEqual(actualSplit2, expectedSplit2) {
		t.Errorf("Uneven Split 2 files mismatch:\nExpected: %v\nActual:   %v", expectedSplit2, actualSplit2)
	}
}

// TestProcessAndSplit_SkipsDirectories tests that directory entries are skipped.
func TestProcessAndSplit_SkipsDirectories(t *testing.T) {
	buf := new(bytes.Buffer)
	zw := zip.NewWriter(buf)
	_, err := zw.CreateHeader(&zip.FileHeader{Name: "dir1/", Method: zip.Store})
	if err != nil {
		t.Fatalf("Failed to create dir header: %v", err)
	}
	f, err := zw.CreateHeader(&zip.FileHeader{Name: "dir1/file_in_dir.txt", Method: zip.Deflate})
	if err != nil {
		t.Fatalf("Failed to create file header: %v", err)
	}
	_, err = io.WriteString(f, "content")
	if err != nil {
		t.Fatalf("Failed to write file content: %v", err)
	}
	zw.Close()
	soongReader, _ := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))

	readers := []zipReader{{reader: soongReader, zipPath: "zip.zip"}}
	writer, writerBuf := createTestAppZipWriter()
	writers := []zipWriter{writer}

	processAndSplit(readers, writers)
	writer.writer.Close()

	// --- Assert ---
	expectedFiles := []string{"dir1/file_in_dir.txt"}
	actualFiles, _ := getFilenamesFromSoongZipBuffer(writerBuf)

	if !reflect.DeepEqual(actualFiles, expectedFiles) {
		t.Errorf("Directory skip test files mismatch:\nExpected: %v\nActual:   %v", expectedFiles, actualFiles)
	}
}

// TestProcessAndSplit_DuplicateFiles_Panics tests that duplicate files across zips cause a panic.
func TestProcessAndSplit_DuplicateFiles_Panics(t *testing.T) {
	zip1Entries := []ZipEntry{{Name: "common.txt", Content: "v1"}}
	zip2Entries := []ZipEntry{{Name: "common.txt", Content: "v2"}} // Duplicate
	zip1SoongReader, _ := createInMemorySoongZip(zip1Entries)
	zip2SoongReader, _ := createInMemorySoongZip(zip2Entries)

	readers := []zipReader{
		{reader: zip1SoongReader, zipPath: "zip1.zip"},
		{reader: zip2SoongReader, zipPath: "zip2.zip"},
	}
	writer, _ := createTestAppZipWriter()
	writers := []zipWriter{writer}

	defer func() {
		if r := recover(); r == nil {
			t.Errorf("Expected processAndSplit to panic on duplicate files, but it did not")
		} else {
			errMsg := fmt.Sprintf("%v", r)
			if !strings.Contains(errMsg, "Duplicate file 'common.txt'") {
				t.Errorf("Panic message did not contain expected duplicate file info. Got: %s", errMsg)
			}
		}
	}()
	processAndSplit(readers, writers)
}

// TestProcessAndSplit_NoProcessableFiles_Panics tests panic when no files match filter or are directories.
func TestProcessAndSplit_NoProcessableFiles_Panics(t *testing.T) {
	buf := new(bytes.Buffer)
	zw := zip.NewWriter(buf)
	_, _ = zw.CreateHeader(&zip.FileHeader{Name: "onlydir/", Method: zip.Store})
	zw.Close()
	soongReader, _ := zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))

	readers := []zipReader{{reader: soongReader, zipPath: "zip.zip"}}
	writer, _ := createTestAppZipWriter()
	writers := []zipWriter{writer}
	originalFilter := filter
	filter = nil
	defer func() { filter = originalFilter }()

	defer func() {
		if r := recover(); r == nil {
			t.Errorf("Expected processAndSplit to panic with no processable files, but it did not")
		} else {
			errMsg := fmt.Sprintf("%v", r)
			if !strings.Contains(errMsg, "no processable files found") {
				t.Errorf("Panic message incorrect. Got: %s", errMsg)
			}
		}
	}()
	processAndSplit(readers, writers)
}

// TestProcessAndSplit_TooManySplits_Panics tests panic when numSplits > numFiles.
func TestProcessAndSplit_TooManySplits_Panics(t *testing.T) {
	zipEntries := []ZipEntry{{Name: "file1.txt", Content: "1"}} // Only one file
	soongReader, _ := createInMemorySoongZip(zipEntries)
	readers := []zipReader{{reader: soongReader, zipPath: "zip.zip"}}
	writer1, _ := createTestAppZipWriter()
	writer2, _ := createTestAppZipWriter()
	writers := []zipWriter{writer1, writer2}
	originalFilter := filter
	filter = nil
	defer func() { filter = originalFilter }()

	defer func() {
		if r := recover(); r == nil {
			t.Errorf("Expected processAndSplit to panic with too many splits, but it did not")
		} else {
			errMsg := fmt.Sprintf("%v", r)
			if !strings.Contains(errMsg, "number of splits '2' more than total files '1'") {
				t.Errorf("Panic message incorrect. Got: %s", errMsg)
			}
		}
	}()
	processAndSplit(readers, writers)
}

// TestProcessAndSplit_WithFilter tests file filtering.
func TestProcessAndSplit_WithFilter(t *testing.T) {
	zipEntries := []ZipEntry{
		{Name: "image.jpg", Content: "jpg_data"},
		{Name: "document.txt", Content: "txt_data"},
		{Name: "archive.zip", Content: "zip_data"},
	}
	soongReader, _ := createInMemorySoongZip(zipEntries)
	readers := []zipReader{{reader: soongReader, zipPath: "zip.zip"}}
	writer, writerBuf := createTestAppZipWriter()
	writers := []zipWriter{writer}

	originalFilter := filter
	filter = multiFlag{"*.txt", "*.zip"}
	defer func() { filter = originalFilter }()

	processAndSplit(readers, writers)
	writer.writer.Close()

	expectedFiles := []string{"archive.zip", "document.txt"}
	actualFiles, _ := getFilenamesFromSoongZipBuffer(writerBuf)
	if !reflect.DeepEqual(actualFiles, expectedFiles) {
		t.Errorf("Filter test files mismatch:\nExpected: %v\nActual:   %v", expectedFiles, actualFiles)
	}
}

// TestProcessAndSplit_EmptyInputZips_Panics tests panic when input zips are empty.
func TestProcessAndSplit_EmptyInputZips_Panics(t *testing.T) {
	soongReader, _ := createInMemorySoongZip([]ZipEntry{}) // Empty zip
	readers := []zipReader{{reader: soongReader, zipPath: "empty.zip"}}
	writer, _ := createTestAppZipWriter()
	writers := []zipWriter{writer}
	originalFilter := filter
	filter = nil
	defer func() { filter = originalFilter }()

	defer func() {
		if r := recover(); r == nil {
			t.Errorf("Expected processAndSplit to panic with empty input zips, but it did not")
		} else {
			errMsg := fmt.Sprintf("%v", r)
			if !strings.Contains(errMsg, "no processable files found") {
				t.Errorf("Panic message incorrect for empty input. Got: %s", errMsg)
			}
		}
	}()
	processAndSplit(readers, writers)
}

type ZipEntry struct {
	Name    string
	Content string
}

// Helper function to create an in-memory ZIP with specified files.
// Each fileContent entry should be a ZipEntry content.
func createInMemorySoongZip(entries []ZipEntry) (*zip.Reader, error) { // Returns *zip.Reader from soong
	buf := new(bytes.Buffer)
	zipWriter := zip.NewWriter(buf)

	for _, entry := range entries {
		name := filepath.ToSlash(entry.Name)
		fh := &zip.FileHeader{
			Name:   name,
			Method: zip.Deflate,
		}

		f, err := zipWriter.CreateHeader(fh)
		if err != nil {
			return nil, err
		}
		_, err = io.WriteString(f, entry.Content)
		if err != nil {
			return nil, err
		}
	}
	if err := zipWriter.Close(); err != nil {
		return nil, err
	}

	return zip.NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
}

// Helper function to create a zipWriter for testing.
func createTestAppZipWriter() (zipWriter, *bytes.Buffer) {
	outputBuf := &bytes.Buffer{}
	return zipWriter{
		writer:     zip.NewWriter(outputBuf),
		zipBytes:   outputBuf,
		outputPath: "test_output.zip",
	}, outputBuf
}

// Helper function to get filenames from a zip.Reader (from its buffer)
func getFilenamesFromSoongZipBuffer(zipBuf *bytes.Buffer) ([]string, error) {
	if zipBuf.Len() == 0 {
		return []string{}, nil
	}
	r, err := zip.NewReader(bytes.NewReader(zipBuf.Bytes()), int64(zipBuf.Len()))
	if err != nil {
		return nil, err
	}
	var names []string
	for _, f := range r.File {
		names = append(names, f.Name)
	}
	sort.Strings(names)
	return names, nil
}