package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"golang.org/x/exp/rand"
)

var Error *log.Logger
var Warning *log.Logger

func init() {
	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
	logger := io.MultiWriter(os.Stdout)
	log.SetOutput(logger)

	Error = log.New(io.MultiWriter(os.Stderr, os.Stdout),
		fmt.Sprintf("%sERROR:%s ", "\033[0;101m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
	Warning = log.New(io.MultiWriter(os.Stdout),
		fmt.Sprintf("%sWarning:%s ", "\033[0;93m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
}

func main() {
	percent := flag.Int("p", 20, "Percent to mangle")
	unsafe := flag.Bool("unsafe", false, "Unsafe mode")
	w := flag.Int("w", 8, "Workers")
	flag.Parse()

	workers := make(chan struct{}, *w)
	wg := sync.WaitGroup{}

	log.Printf("Corrupting %d%% of each input file", *percent)
	log.Printf("Using %d workers", *w)

	for _, file := range flag.Args() {
		wg.Add(1)
		workers <- struct{}{}

		go func(file string) {
			defer wg.Done()
			defer func() { <-workers }()

			file = filepath.Clean(file)
			file = filepath.ToSlash(file)

			log.Printf("Corrupting %s", file)
			filehandle, err := os.OpenFile(file, os.O_RDWR, 0644)
			if err != nil {
				Error.Printf("Error opening %v: %v", file, err)
				return
			}
			defer filehandle.Close()

			// data, err := io.ReadAll(filehandle)
			// if err != nil {
			// 	Error.Printf("Error reading %v: %v", file, err)
			// 	return
			// }
			// log.Printf("Read %d bytes from %s", len(data), file)
			// filehandle.Seek(0, io.SeekStart)

			info, err := os.Stat(file)
			if err != nil {
				Error.Printf("Error getting file info for %v: %v", file, err)
				return
			}

			if info.IsDir() {
				Error.Printf("Skipping directory %s", file)
				return
			}

			if !*unsafe {
				err := backup(file)
				if err != nil {
					Error.Printf("Error backing up %v: %v", file, err)
					return
				}
			}

			filesize := info.Size()
			log.Printf("File size is %d bytes", filesize)

			toCorrupt := int64(float32(filesize) * (float32(*percent) / 100))
			indicesToCorrupt := selectRandomIndices(filesize, toCorrupt)
			sort.Slice(indicesToCorrupt, func(i, j int) bool { return indicesToCorrupt[i] < indicesToCorrupt[j] })
			log.Printf("Corrupting %d bytes in %s", toCorrupt, file)

			for _, index := range indicesToCorrupt {
				rbyte := byte(rand.Intn(256))
				filehandle.Seek(int64(index), io.SeekStart)
				filehandle.Write([]byte{rbyte})
			}

			log.Printf("Corrupted %d bytes in %s", toCorrupt, file)

			// n, err := filehandle.Write(data)
			// if err != nil {
			// 	Error.Printf("Error writing %v: %v", file, err)
			// 	return
			// }
			// log.Printf("Wrote %d bytes to %s", n, file)
		}(file)
	}
	wg.Wait()
}

func selectRandomIndices(size, n int64) []int64 {
	indices := make([]int64, n)
	for i := range indices {
		indices[i] = int64(i)
	}

	rand.Seed(uint64(time.Now().UnixNano()))

	for i := n; i < size; i++ {
		r := rand.Int63n(i + 1)
		if r < n {
			indices[r] = i
		}
	}

	return indices
}

func backup(file string) error {
	backupfile := file + ".bak"

	_, err := os.Stat(backupfile)
	if !os.IsNotExist(err) {
		return fmt.Errorf("backup file %s already exists %s.bak", backupfile)
	}

	cmd := exec.Command("cp", "-a", file, backupfile)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed creating backup file with return %d and out %s", err, string(out))
	}

	return nil
}