// Package fileutil provides secure file operation utilities to prevent path traversal attacks.
package fileutil

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// FileType represents a known file type with its magic bytes
type FileType struct {
	Name        string
	Description string
	MagicBytes  []byte
	Extensions  []string
}

// Known file types for ML artifacts
var (
	// SafeTensor uses ZIP format
	SafeTensors = FileType{
		Name:        "safetensors",
		MagicBytes:  []byte{0x50, 0x4B, 0x03, 0x04}, // ZIP header
		Extensions:  []string{".safetensors"},
		Description: "SafeTensors model format",
	}

	// GGUF format
	GGUF = FileType{
		Name:        "gguf",
		MagicBytes:  []byte{0x47, 0x47, 0x55, 0x46}, // "GGUF"
		Extensions:  []string{".gguf"},
		Description: "GGML/GGUF model format",
	}

	// HDF5 format
	HDF5 = FileType{
		Name:        "hdf5",
		MagicBytes:  []byte{0x89, 0x48, 0x44, 0x46}, // HDF5 signature
		Extensions:  []string{".h5", ".hdf5", ".hdf"},
		Description: "HDF5 data format",
	}

	// NumPy format
	NumPy = FileType{
		Name:        "numpy",
		MagicBytes:  []byte{0x93, 0x4E, 0x55, 0x4D}, // NUMPY magic
		Extensions:  []string{".npy"},
		Description: "NumPy array format",
	}

	// JSON format
	JSON = FileType{
		Name:        "json",
		MagicBytes:  []byte{0x7B}, // "{"
		Extensions:  []string{".json"},
		Description: "JSON data format",
	}

	// CSV format (text-based, no reliable magic bytes)
	CSV = FileType{
		Name:        "csv",
		MagicBytes:  nil, // Text-based, validated by content inspection
		Extensions:  []string{".csv"},
		Description: "CSV data format",
	}

	// YAML format (text-based)
	YAML = FileType{
		Name:        "yaml",
		MagicBytes:  nil, // Text-based
		Extensions:  []string{".yaml", ".yml"},
		Description: "YAML configuration format",
	}

	// Text format
	Text = FileType{
		Name:        "text",
		MagicBytes:  nil, // Text-based
		Extensions:  []string{".txt", ".md", ".rst"},
		Description: "Plain text format",
	}

	// AllAllowedTypes contains all types that are permitted for upload
	AllAllowedTypes = []FileType{SafeTensors, GGUF, HDF5, NumPy, JSON, CSV, YAML, Text}

	// BinaryModelTypes contains binary model formats only
	BinaryModelTypes = []FileType{SafeTensors, GGUF, HDF5, NumPy}
)

// DangerousExtensions are file extensions that should be rejected immediately
var DangerousExtensions = []string{
	".pt", ".pkl", ".pickle", // PyTorch pickle - arbitrary code execution
	".pth",                          // PyTorch state dict (often pickle-based)
	".joblib",                       // scikit-learn pickle format
	".exe", ".dll", ".so", ".dylib", // Executables
	".sh", ".bat", ".cmd", ".ps1", // Scripts
	".zip", ".tar", ".gz", ".bz2", ".xz", // Archives (may contain malicious files)
	".rar", ".7z",
}

// ValidateFileType checks if a file matches an allowed type using magic bytes validation.
// Returns the detected type or an error if the file is not allowed.
func ValidateFileType(filePath string, allowedTypes []FileType) (*FileType, error) {
	// First check extension-based rejection
	ext := strings.ToLower(filepath.Ext(filePath))
	for _, dangerous := range DangerousExtensions {
		if ext == dangerous {
			return nil, fmt.Errorf("file type not allowed (dangerous extension): %s", ext)
		}
	}

	// Open and read the file
	file, err := os.Open(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to open file for type validation: %w", err)
	}
	defer file.Close()

	// Read first 8 bytes for magic byte detection
	header := make([]byte, 8)
	n, err := file.Read(header)
	if err != nil {
		return nil, fmt.Errorf("failed to read file header: %w", err)
	}
	header = header[:n]

	// Try to match by magic bytes first
	for _, ft := range allowedTypes {
		if len(ft.MagicBytes) > 0 && len(header) >= len(ft.MagicBytes) {
			if bytes.Equal(header[:len(ft.MagicBytes)], ft.MagicBytes) {
				return &ft, nil
			}
		}
	}

	// For text-based formats, validate by extension and content
	for _, ft := range allowedTypes {
		if ft.MagicBytes == nil {
			// Check if extension matches
			for _, allowedExt := range ft.Extensions {
				if ext == allowedExt {
					// Additional content validation for text files
					if err := validateTextContent(filePath, ft); err != nil {
						return nil, err
					}
					return &ft, nil
				}
			}
		}
	}

	return nil, fmt.Errorf("file type not recognized or not in allowed list")
}

// validateTextContent performs basic validation on text files
func validateTextContent(filePath string, ft FileType) error {
	// Read a sample of the file
	data, err := os.ReadFile(filePath)
	if err != nil {
		return fmt.Errorf("failed to read text file: %w", err)
	}

	// Check for null bytes (indicates binary content)
	if bytes.Contains(data, []byte{0x00}) {
		return fmt.Errorf("file contains null bytes, not valid %s", ft.Name)
	}

	// For JSON, validate it can be parsed
	if ft.Name == "json" {
		// Basic JSON validation - check for valid JSON structure
		trimmed := bytes.TrimSpace(data)
		if len(trimmed) == 0 {
			return fmt.Errorf("empty JSON file")
		}
		if (trimmed[0] != '{' && trimmed[0] != '[') ||
			(trimmed[len(trimmed)-1] != '}' && trimmed[len(trimmed)-1] != ']') {
			return fmt.Errorf("invalid JSON structure")
		}
	}

	return nil
}

// IsAllowedExtension checks if a file extension is in the allowed list
func IsAllowedExtension(filePath string, allowedTypes []FileType) bool {
	ext := strings.ToLower(filepath.Ext(filePath))

	// Check against dangerous extensions first
	for _, dangerous := range DangerousExtensions {
		if ext == dangerous {
			return false
		}
	}

	// Check against allowed types
	for _, ft := range allowedTypes {
		for _, allowedExt := range ft.Extensions {
			if ext == allowedExt {
				return true
			}
		}
	}

	return false
}

// ValidateDatasetFile validates a dataset file for safe formats
func ValidateDatasetFile(filePath string) error {
	_, err := ValidateFileType(filePath, AllAllowedTypes)
	return err
}

// ValidateModelFile validates a model file for safe binary formats only
func ValidateModelFile(filePath string) error {
	ft, err := ValidateFileType(filePath, BinaryModelTypes)
	if err != nil {
		return err
	}

	// Additional check: ensure it's actually a model format, not just a matching extension
	if ft == nil {
		return fmt.Errorf("file type validation returned nil type")
	}

	return nil
}
