package main

import (
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
)

func main() {
	if len(os.Args) < 4 {
		fmt.Println("Usage: csvsplitter <file.csv> <lines_per_file> <has_header: yes/no> [<eol: lf/crlf>]")
		return
	}

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

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

	useCRLF := true
	if len(os.Args) > 4 {
		eolTypeArg := os.Args[4]
		if strings.EqualFold(eolTypeArg, "lf") {
			useCRLF = false
		}
	}

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

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

	reader := csv.NewReader(file)

	var header []string
	if hasHeader {
		header, err = reader.Read()
		if err != nil {
			return err
		}
	}

	partNumber := 1
	var recordsForFile [][]string
	if hasHeader {
		recordsForFile = append(recordsForFile, header) // Prepend the header for each new file if needed
	}

	lineCount := 0
	for {
		record, err := reader.Read()
		if err != nil {
			if err == io.EOF {
				// Check if there are any records to write to a new part file
				if lineCount > 0 {
					err = writePart(filePath, partNumber, recordsForFile, useCRLF)
					if err != nil {
						return err
					}
				}
				break
			}
			return err
		}

		lineCount++
		recordsForFile = append(recordsForFile, record)
		if lineCount == linesPerFile {
			err = writePart(filePath, partNumber, recordsForFile, useCRLF)
			if err != nil {
				return err
			}
			partNumber++
			recordsForFile = [][]string{}
			if hasHeader {
				recordsForFile = append(recordsForFile, header) // Reset for next file, re-adding header if necessary
			}
			lineCount = 0
		}
	}

	return nil
}

func writePart(filePath string, partNumber int, records [][]string, useCRLF bool) error {
	partFilePath := fmt.Sprintf("%s_part%d.csv", filePath, partNumber)
	partFile, err := os.Create(partFilePath)
	if err != nil {
		return err
	}
	defer partFile.Close()

	writer := csv.NewWriter(partFile)
	writer.UseCRLF = useCRLF

	err = writer.WriteAll(records)
	if err != nil {
		return err
	}

	writer.Flush()
	return writer.Error()
}