package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
)

func processLines(path string, f func(string) (stop bool, err error)) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		stop, err := f(scanner.Text())
		if stop || err != nil {
			return err
		}
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	return nil
}

func processGroups(path string, f func([]string) (stop bool, err error)) error {
	group := []string{}
	if err := processLines(path, func(line string) (stop bool, err error) {
		if strings.HasPrefix(line, "mask") {
			stop, err = f(group)
			if stop || err != nil {
				return
			}
			group = []string{}
		} else {
			group = append(group, line)
		}

		return
	}); err != nil {
		return err
	}

	_, err := f(group)

	return err
}

// StringSet is a set of strings
type StringSet struct {
	d map[string]bool
}

// Add an item to a StringSet
func (s *StringSet) Add(v string) {
	s.d[v] = true
}

// Len gives the length of a StringSet
func (s StringSet) Len() int {
	return len(s.d)
}

// NewStringSet creates an empty string set
func NewStringSet() StringSet {
	return StringSet{map[string]bool{}}
}

// Union returns the union of two SringSets
func Union(s1, s2 StringSet) StringSet {
	union := NewStringSet()
	for v := range s1.d {
		union.Add(v)
	}
	for v := range s2.d {
		union.Add(v)
	}
	return union
}

// Intersection returns the union of two SringSets
func Intersection(s1, s2 StringSet) StringSet {
	intersect := NewStringSet()
	for v := range s1.d {
		if _, ok := s2.d[v]; ok {
			intersect.Add(v)
		}
	}
	return intersect
}

func part1() {
	total := 0
	err := processGroups("input.txt", func(group []string) (stop bool, err error) {
		questions := NewStringSet()
		for _, person := range group {
			for _, q := range person {
				questions.Add(string(q))
			}
		}
		total += questions.Len()
		return
	})
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("Part 1 answer %d\n", total)
}

func part2() {
	total := 0
	err := processGroups("input.txt", func(group []string) (stop bool, err error) {
		result := NewStringSet()
		for _, q := range group[0] {
			result.Add(string(q))
		}
		for _, person := range group[1:] {
			questions := NewStringSet()
			for _, q := range person {
				questions.Add(string(q))
			}
			result = Intersection(result, questions)
		}
		total += result.Len()
		return
	})
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("Part 2 answer %d\n", total)
}

func main() {
	part1()
	part2()
}