// Package fs provides parallel tar.gz extraction using pgzip
package fs

import (
	"archive/tar"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/klauspost/pgzip"
)

// CopyWithContext copies data from src to dst while checking for context cancellation.
// This allows Ctrl+C to interrupt large file extractions instead of blocking until complete.
// Checks context every 1MB of data copied for responsive interruption.
func CopyWithContext(ctx context.Context, dst io.Writer, src io.Reader) (int64, error) {
	buf := make([]byte, 1024*1024) // 1MB buffer - check context every 1MB
	var written int64
	for {
		// Check for cancellation before each read
		select {
		case <-ctx.Done():
			return written, ctx.Err()
		default:
		}

		nr, readErr := src.Read(buf)
		if nr > 0 {
			nw, writeErr := dst.Write(buf[:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if writeErr != nil {
				return written, writeErr
			}
			if nr != nw {
				return written, io.ErrShortWrite
			}
		}
		if readErr != nil {
			if readErr == io.EOF {
				return written, nil
			}
			return written, readErr
		}
	}
}

// ParallelGzipWriter wraps pgzip.Writer for streaming compression
type ParallelGzipWriter struct {
	*pgzip.Writer
}

// NewParallelGzipWriter creates a parallel gzip writer using all CPU cores
// This is 2-4x faster than standard gzip on multi-core systems
func NewParallelGzipWriter(w io.Writer, level int) (*ParallelGzipWriter, error) {
	gzWriter, err := pgzip.NewWriterLevel(w, level)
	if err != nil {
		return nil, fmt.Errorf("cannot create gzip writer: %w", err)
	}
	// Set block size and concurrency for parallel compression
	if err := gzWriter.SetConcurrency(1<<20, runtime.NumCPU()); err != nil {
		// Non-fatal, continue with defaults
	}
	return &ParallelGzipWriter{Writer: gzWriter}, nil
}

// ExtractProgress reports extraction progress
type ExtractProgress struct {
	CurrentFile  string
	BytesRead    int64
	TotalBytes   int64
	FilesCount   int
	CurrentIndex int
}

// ProgressCallback is called during extraction
type ProgressCallback func(progress ExtractProgress)

// ExtractTarGzParallel extracts a tar.gz archive using parallel gzip decompression
// This is 2-4x faster than standard gzip on multi-core systems
// Uses pgzip which decompresses in parallel using multiple goroutines
func ExtractTarGzParallel(ctx context.Context, archivePath, destDir string, progressCb ProgressCallback) error {
	// Open the archive
	file, err := os.Open(archivePath)
	if err != nil {
		return fmt.Errorf("cannot open archive: %w", err)
	}
	defer file.Close()

	// Get file size for progress
	stat, err := file.Stat()
	if err != nil {
		return fmt.Errorf("cannot stat archive: %w", err)
	}
	totalSize := stat.Size()

	// Create parallel gzip reader
	// Uses all available CPU cores for decompression
	gzReader, err := pgzip.NewReaderN(file, 1<<20, runtime.NumCPU()) // 1MB blocks
	if err != nil {
		return fmt.Errorf("cannot create gzip reader: %w", err)
	}
	defer gzReader.Close()

	// Create tar reader
	tarReader := tar.NewReader(gzReader)

	// Track progress
	var bytesRead int64
	var filesCount int

	// Extract each file
	for {
		// Check context
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("error reading tar: %w", err)
		}

		// Security: prevent path traversal
		targetPath := filepath.Join(destDir, header.Name)
		if !strings.HasPrefix(filepath.Clean(targetPath), filepath.Clean(destDir)) {
			return fmt.Errorf("path traversal detected: %s", header.Name)
		}

		filesCount++

		// Report progress
		if progressCb != nil {
			// Estimate bytes read from file position
			pos, _ := file.Seek(0, io.SeekCurrent)
			progressCb(ExtractProgress{
				CurrentFile:  header.Name,
				BytesRead:    pos,
				TotalBytes:   totalSize,
				FilesCount:   filesCount,
				CurrentIndex: filesCount,
			})
		}

		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(targetPath, 0700); err != nil {
				return fmt.Errorf("cannot create directory %s: %w", targetPath, err)
			}

		case tar.TypeReg:
			// Ensure parent directory exists
			if err := os.MkdirAll(filepath.Dir(targetPath), 0700); err != nil {
				return fmt.Errorf("cannot create parent directory: %w", err)
			}

			// Create file with secure permissions
			outFile, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
			if err != nil {
				return fmt.Errorf("cannot create file %s: %w", targetPath, err)
			}

			// Copy with context awareness to allow Ctrl+C interruption during large file extraction
			written, err := CopyWithContext(ctx, outFile, tarReader)
			outFile.Close()

			if err != nil {
				// Clean up partial file on error
				os.Remove(targetPath)
				return fmt.Errorf("error writing %s: %w", targetPath, err)
			}

			bytesRead += written

		case tar.TypeSymlink:
			// Handle symlinks (validate target is within destDir)
			linkTarget := header.Linkname
			absTarget := filepath.Join(filepath.Dir(targetPath), linkTarget)
			if !strings.HasPrefix(filepath.Clean(absTarget), filepath.Clean(destDir)) {
				// Skip symlinks that point outside
				continue
			}
			if err := os.Symlink(linkTarget, targetPath); err != nil {
				// Ignore symlink errors (may not be supported)
				continue
			}

		default:
			// Skip other types (devices, etc.)
			continue
		}
	}

	return nil
}

// ListTarGzContents lists the contents of a tar.gz archive without extracting
// Returns a slice of file paths in the archive
// Uses parallel gzip decompression for 2-4x faster listing on multi-core systems
func ListTarGzContents(ctx context.Context, archivePath string) ([]string, error) {
	// Open the archive
	file, err := os.Open(archivePath)
	if err != nil {
		return nil, fmt.Errorf("cannot open archive: %w", err)
	}
	defer file.Close()

	// Create parallel gzip reader
	gzReader, err := pgzip.NewReaderN(file, 1<<20, runtime.NumCPU())
	if err != nil {
		return nil, fmt.Errorf("cannot create gzip reader: %w", err)
	}
	defer gzReader.Close()

	// Create tar reader
	tarReader := tar.NewReader(gzReader)

	var files []string
	for {
		// Check for cancellation
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("tar read error: %w", err)
		}

		files = append(files, header.Name)
	}

	return files, nil
}

// ExtractTarGzFast is a convenience wrapper that chooses the best extraction method
// Uses parallel gzip if available, falls back to system tar if needed
func ExtractTarGzFast(ctx context.Context, archivePath, destDir string, progressCb ProgressCallback) error {
	// Always use parallel Go implementation - it's faster and more portable
	return ExtractTarGzParallel(ctx, archivePath, destDir, progressCb)
}

// CreateProgress reports archive creation progress
type CreateProgress struct {
	CurrentFile  string
	BytesWritten int64
	FilesCount   int
}

// CreateProgressCallback is called during archive creation
type CreateProgressCallback func(progress CreateProgress)

// CreateTarGzParallel creates a tar.gz archive using parallel gzip compression
// This is 2-4x faster than standard gzip on multi-core systems
// Uses pgzip which compresses in parallel using multiple goroutines
func CreateTarGzParallel(ctx context.Context, sourceDir, outputPath string, compressionLevel int, progressCb CreateProgressCallback) error {
	// Create output file
	outFile, err := os.Create(outputPath)
	if err != nil {
		return fmt.Errorf("cannot create archive: %w", err)
	}
	defer outFile.Close()

	// Create parallel gzip writer
	// Uses all available CPU cores for compression
	gzWriter, err := pgzip.NewWriterLevel(outFile, compressionLevel)
	if err != nil {
		return fmt.Errorf("cannot create gzip writer: %w", err)
	}
	// Set block size and concurrency for parallel compression
	if err := gzWriter.SetConcurrency(1<<20, runtime.NumCPU()); err != nil {
		// Non-fatal, continue with defaults
	}
	defer gzWriter.Close()

	// Create tar writer
	tarWriter := tar.NewWriter(gzWriter)
	defer tarWriter.Close()

	var bytesWritten int64
	var filesCount int

	// Walk the source directory
	err = filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
		// Check for cancellation
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		if err != nil {
			return err
		}

		// Get relative path
		relPath, err := filepath.Rel(sourceDir, path)
		if err != nil {
			return err
		}

		// Skip the root directory itself
		if relPath == "." {
			return nil
		}

		// Create tar header
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return fmt.Errorf("cannot create header for %s: %w", relPath, err)
		}

		// Use relative path in archive
		header.Name = relPath

		// Handle symlinks
		if info.Mode()&os.ModeSymlink != 0 {
			link, err := os.Readlink(path)
			if err != nil {
				return fmt.Errorf("cannot read symlink %s: %w", path, err)
			}
			header.Linkname = link
		}

		// Write header
		if err := tarWriter.WriteHeader(header); err != nil {
			return fmt.Errorf("cannot write header for %s: %w", relPath, err)
		}

		// If it's a regular file, write its contents
		if info.Mode().IsRegular() {
			file, err := os.Open(path)
			if err != nil {
				return fmt.Errorf("cannot open %s: %w", path, err)
			}
			defer file.Close()

			written, err := io.Copy(tarWriter, file)
			if err != nil {
				return fmt.Errorf("cannot write %s: %w", path, err)
			}
			bytesWritten += written
		}

		filesCount++

		// Report progress
		if progressCb != nil {
			progressCb(CreateProgress{
				CurrentFile:  relPath,
				BytesWritten: bytesWritten,
				FilesCount:   filesCount,
			})
		}

		return nil
	})

	if err != nil {
		// Clean up partial file on error
		outFile.Close()
		os.Remove(outputPath)
		return err
	}

	// Explicitly close tar and gzip to flush all data
	if err := tarWriter.Close(); err != nil {
		return fmt.Errorf("cannot close tar writer: %w", err)
	}
	if err := gzWriter.Close(); err != nil {
		return fmt.Errorf("cannot close gzip writer: %w", err)
	}

	return nil
}

// EstimateCompressionRatio samples the archive to estimate uncompressed size
// Returns a multiplier (e.g., 3.0 means uncompressed is ~3x the compressed size)
func EstimateCompressionRatio(archivePath string) (float64, error) {
	file, err := os.Open(archivePath)
	if err != nil {
		return 3.0, err // Default to 3x
	}
	defer file.Close()

	// Get compressed size
	stat, err := file.Stat()
	if err != nil {
		return 3.0, err
	}
	compressedSize := stat.Size()

	// Read first 1MB and measure decompression ratio
	gzReader, err := pgzip.NewReader(file)
	if err != nil {
		return 3.0, err
	}
	defer gzReader.Close()

	// Read up to 1MB of decompressed data
	buf := make([]byte, 1<<20)
	n, _ := io.ReadFull(gzReader, buf)

	if n < 1024 {
		return 3.0, nil // Not enough data, use default
	}

	// Estimate: decompressed / compressed
	// Based on sample of first 1MB
	compressedPortion := float64(compressedSize) * (float64(n) / float64(compressedSize))
	if compressedPortion > 0 {
		ratio := float64(n) / compressedPortion
		if ratio > 1.0 && ratio < 20.0 {
			return ratio, nil
		}
	}

	return 3.0, nil // Default
}
