package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
)

func main() {
	// Define a boolean flag -l to count lines instead of words
	lines := flag.Bool("l", false, "count lines")
	// Define a boolean flag -b to count bytes instead of words
	bytes := flag.Bool("b", false, "count bytes")

	// Parse the flags provided by the user
	flag.Parse()
	// Calling the count function to count the number of words (or lines)
	// received from the standard input and printing it out
	fmt.Println(count(os.Stdin, *lines, *bytes))
}

func count(r io.Reader, countLines bool, countBytes bool) int {
	scanner := bufio.NewScanner(r)
	if !countLines {
		if countBytes {
			scanner.Split(bufio.ScanBytes)
		} else {
			// Split sets the split function for the Scanner.
			// The default split function is ScanLines.
			scanner.Split(bufio.ScanWords)
		}
	}

	wc := 0
	// For every word scanned, increment the counter
	for scanner.Scan() {
		wc++
	}
	return wc
}