package security

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

	"github.com/jfraeys/fetch_ml/internal/fileutil"
)

func TestValidateFileType_AllowedTypes(t *testing.T) {
	tempDir := t.TempDir()

	tests := []struct {
		name      string
		content   []byte
		ext       string
		wantType  string
		wantErr   bool
	}{
		{
			name:     "valid safetensors (ZIP magic)",
			content:  []byte{0x50, 0x4B, 0x03, 0x04, 0x00, 0x00, 0x00, 0x00},
			ext:     ".safetensors",
			wantType: "safetensors",
			wantErr:  false,
		},
		{
			name:     "valid GGUF",
			content:  []byte{0x47, 0x47, 0x55, 0x46, 0x00, 0x00, 0x00, 0x00},
			ext:     ".gguf",
			wantType: "gguf",
			wantErr:  false,
		},
		{
			name:     "valid JSON",
			content:  []byte(`{"key": "value"}`),
			ext:     ".json",
			wantType: "json",
			wantErr:  false,
		},
		{
			name:     "valid text file",
			content:  []byte("Hello, World!"),
			ext:     ".txt",
			wantType: "text",
			wantErr:  false,
		},
		{
			name:     "dangerous pickle extension",
			content:  []byte{0x00, 0x00, 0x00, 0x00},
			ext:     ".pkl",
			wantErr:  true,
		},
		{
			name:     "dangerous pytorch extension",
			content:  []byte{0x00, 0x00, 0x00, 0x00},
			ext:     ".pt",
			wantErr:  true,
		},
		{
			name:     "executable extension",
			content:  []byte{0x00, 0x00, 0x00, 0x00},
			ext:     ".exe",
			wantErr:  true,
		},
		{
			name:     "script extension",
			content:  []byte("#!/bin/bash"),
			ext:     ".sh",
			wantErr:  true,
		},
		{
			name:     "archive extension",
			content:  []byte{0x00, 0x00, 0x00, 0x00},
			ext:     ".zip",
			wantErr:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			filePath := filepath.Join(tempDir, "test"+tt.ext)
			if err := os.WriteFile(filePath, tt.content, 0644); err != nil {
				t.Fatalf("Failed to create test file: %v", err)
			}

			ft, err := fileutil.ValidateFileType(filePath, fileutil.AllAllowedTypes)
			if tt.wantErr {
				if err == nil {
					t.Errorf("ValidateFileType() error = nil, wantErr %v", tt.wantErr)
				}
			} else {
				if err != nil {
					t.Errorf("ValidateFileType() unexpected error = %v", err)
					return
				}
				if ft.Name != tt.wantType {
					t.Errorf("ValidateFileType() type = %v, want %v", ft.Name, tt.wantType)
				}
			}
		})
	}
}

func TestValidateModelFile(t *testing.T) {
	tempDir := t.TempDir()

	tests := []struct {
		name     string
		content  []byte
		ext      string
		wantErr  bool
	}{
		{
			name:    "valid model - safetensors",
			content: []byte{0x50, 0x4B, 0x03, 0x04},
			ext:     ".safetensors",
			wantErr: false,
		},
		{
			name:    "valid model - gguf",
			content: []byte{0x47, 0x47, 0x55, 0x46},
			ext:     ".gguf",
			wantErr: false,
		},
		{
			name:    "invalid - text file",
			content: []byte("not a model"),
			ext:     ".txt",
			wantErr: true, // Not in BinaryModelTypes
		},
		{
			name:    "invalid - JSON",
			content: []byte(`{"key": "value"}`),
			ext:     ".json",
			wantErr: true, // Not in BinaryModelTypes
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			filePath := filepath.Join(tempDir, "model"+tt.ext)
			if err := os.WriteFile(filePath, tt.content, 0644); err != nil {
				t.Fatalf("Failed to create test file: %v", err)
			}

			err := fileutil.ValidateModelFile(filePath)
			if tt.wantErr {
				if err == nil {
					t.Errorf("ValidateModelFile() error = nil, wantErr %v", tt.wantErr)
				}
			} else {
				if err != nil {
					t.Errorf("ValidateModelFile() unexpected error = %v", err)
				}
			}
		})
	}
}

func TestIsAllowedExtension(t *testing.T) {
	tests := []struct {
		path     string
		expected bool
	}{
		{"model.safetensors", true},
		{"data.json", true},
		{"config.yaml", true},
		{"readme.txt", true},
		{"script.sh", false},
		{"program.exe", false},
		{"archive.zip", false},
		{"pickle.pkl", false},
		{"pytorch.pt", false},
		{"model.pth", false},
	}

	for _, tt := range tests {
		t.Run(tt.path, func(t *testing.T) {
			got := fileutil.IsAllowedExtension(tt.path, fileutil.AllAllowedTypes)
			if got != tt.expected {
				t.Errorf("IsAllowedExtension(%q) = %v, want %v", tt.path, got, tt.expected)
			}
		})
	}
}
