package main

import (
    "archive/zip"
	"bufio"
	"encoding/csv"
	"fmt"
	"os"
	"path/filepath"
	"io"
	"strconv"
	"strings"
)

func main() {
	if len(os.Args) < 4 {
	    //                                filePath   linesPerFileArg  hasHeader          eolTypeArg      doZipArg
		fmt.Println("Usage: csvsplit <file.csv> <lines_per_file> <has_header: y/n> [<eol: lf/crlf>] [<zip: y/n]")
		fmt.Println("  Example:  csvsplit myfile.csv 1000 yes crlf no")
		fmt.Println("  The output files are named like the original file plus \"{file_name}_partN.csv\" where N is 1, 2, 3, ...")
		fmt.Println("  If zip option is yes, the output files are zipped into \"{file_name}.zip\"")
		return
	}

	filePath := os.Args[1]
	linesPerFileArg := os.Args[2]
	hasHeaderArg := os.Args[3]

	hasHeader := strings.HasPrefix(strings.ToLower(hasHeaderArg), "y")
	linesPerFile, err := strconv.Atoi(linesPerFileArg)
	if err != nil {
		fmt.Println("Error: lines_per_file must be an integer")
		return
	}

	useCRLF := true
    if os.PathSeparator == '/' { // unix style
        useCRLF = false
    }
	if len(os.Args) > 4 {
		eolTypeArg := os.Args[4]
		if strings.EqualFold(strings.ToLower(eolTypeArg), "lf") {
			useCRLF = false
		} else if strings.EqualFold(strings.ToLower(eolTypeArg), "crlf") {
		    useCRLF = true
		}
	}
	doZip := false
	if len(os.Args) > 5 {
	    doZipArg := os.Args[5]
	    doZip = strings.HasPrefix(strings.ToLower(doZipArg), "y")
	}
	
	// This argument is a problem in code ... not using for now
	skipLines := 0
	if len(os.Args) > 6 {
	    skipLinesArg := os.Args[6]
		skipLines, err = strconv.Atoi(skipLinesArg)
		skipLines = 0    // disable
		if err != nil {
			fmt.Println("Error: skip_lines must be an integer")
			return
		}
	}

	err = splitCSV(filePath, hasHeader, linesPerFile, useCRLF, doZip, skipLines)
	if err != nil {
		fmt.Printf("Error splitting CSV file: %v\n", err)
	}
}

// Assume outputFileNames is a slice containing the paths to the files you want to zip.
func zipFiles(zipFileName string, outputFileNames []string) error {
	// Create a new zip archive file.
	fmt.Println("Zip File name:", zipFileName)
	newZipFile, err := os.Create(zipFileName)
	if err != nil {
		return err
	}
	defer newZipFile.Close()

	zipWriter := zip.NewWriter(newZipFile)
	defer zipWriter.Close()

	// Add files to the zip archive.
	for _, fileName := range outputFileNames {
		fmt.Println("Add File name to zip:", fileName)
		err := addFileToZip(zipWriter, fileName)
		if err != nil {
			return err
		}
		// Optionally, delete the file after adding it to the zip.
		os.Remove(fileName)
	}

	return nil
}

func addFileToZip(zipWriter *zip.Writer, filename string) error {
	fileToZip, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fileToZip.Close()

	// Get the file information to populate the zip file header.
	info, err := fileToZip.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}
	header.Name = filepath.Base(filename)
	header.Method = zip.Deflate // Use deflate compression.

	writer, err := zipWriter.CreateHeader(header)
	if err != nil {
		return err
	}
	_, err = io.Copy(writer, fileToZip)
	return err
}

func splitCSV(filePath string, hasHeader bool, linesPerFile int, useCRLF bool, doZip bool, skipLines int) error {
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

    // Create a buffer to check for BOM
    bom := make([]byte, 3)
    _, err = file.Read(bom)
    if err != nil {
        fmt.Println("Error reading file:", err)
        return err
    }

    // Check if BOM is present; if not, seek back to the start of the file
    if bom[0] != 0xEF || bom[1] != 0xBB || bom[2] != 0xBF {
        _, err = file.Seek(0, os.SEEK_SET) // No BOM, start from the beginning
        if err != nil {
            fmt.Println("Error seeking file:", err)
            return err
        }
    }
	
	reader := csv.NewReader(bufio.NewReader(file))
	var writer *csv.Writer
	var outputFile *os.File
	var partNumber = 1
	fileNameList := []string{}

	// Handle header if present
	var header []string	
	
	if skipLines > 9999999999 {
	    fmt.Println("Skipping :", skipLines, " lines")
        scanner := bufio.NewScanner(file)
        // Skip the specified number of lines
        for i := 0; i < skipLines; i++ {
            if !scanner.Scan() {
                return fmt.Errorf("Reached EOF before skipping specified number of lines")
            }
        }
	}
	if hasHeader {
		header, err = reader.Read()
		if err != nil {
			return err
		}
	}

	lineCount := 0
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		// Open a new file when necessary
		if lineCount%linesPerFile == 0 {
			if outputFile != nil {
				writer.Flush()
				outputFile.Close()
			}
			outputFileName := fmt.Sprintf("%s_part%03d.csv", filePath, partNumber)
			fileNameList = append(fileNameList, outputFileName)
			outputFile, err = os.Create(outputFileName)
			if err != nil {
				return err
			}
			writer = csv.NewWriter(outputFile)
			if useCRLF {
				writer.UseCRLF = true
			}

			if hasHeader && (lineCount == 0 || partNumber > 1) {
				if err := writer.Write(header); err != nil {
					return err
				}
			}
			partNumber++
		}

		if err := writer.Write(record); err != nil {
			return err
		}

		lineCount++
	}

	if writer != nil {
		writer.Flush()
	}
	if outputFile != nil {
		outputFile.Close()
	}
    if doZip {
	    zipFileName := fmt.Sprintf("%s.zip", filePath)
	    zipFiles(zipFileName, fileNameList)
	}
	return nil
}