package main

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"iter"
	"maps"
	"slices"
	"strings"
)

func Count(n int) iter.Seq[int] {
	return func(yield func(int) bool) {
		fmt.Println("start")
		defer fmt.Println("cleanup")
		for i := 0; i < n; i++ {
			fmt.Println("offer", i)
			if !yield(i) {
				fmt.Println("stop")
				return
			}
			fmt.Println("resume", i)
		}
	}
}

func Filter[T any](src iter.Seq[T], keep func(T) bool) iter.Seq[T] {
	return func(yield func(T) bool) {
		for v := range src {
			if keep(v) && !yield(v) {
				return
			}
		}
	}
}

// Lines consumes r; it neither rewinds nor closes the caller's reader.
func Lines(r io.Reader) iter.Seq2[string, error] {
	return func(yield func(string, error) bool) {
		sc := bufio.NewScanner(r)
		for sc.Scan() {
			if !yield(sc.Text(), nil) {
				return
			}
		}
		if err := sc.Err(); err != nil {
			yield("", err)
		}
	}
}

func UntilCanceled(ctx context.Context, n int) iter.Seq[int] {
	return func(yield func(int) bool) {
		for i := 0; i < n; i++ {
			if ctx.Err() != nil {
				return
			}
			if !yield(i) {
				return
			}
		}
	}
}

func main() {
	seq := Count(4)
	fmt.Println("created")
	for n := range seq {
		fmt.Println("body", n)
		if n == 1 {
			break
		}
	}
	fmt.Println("loop ended")

	fmt.Println("even", slices.Collect(Filter(slices.Values([]int{1, 2, 3, 4}), func(n int) bool { return n%2 == 0 })))
	for i, name := range slices.All([]string{"Ada", "Lin"}) {
		fmt.Println("pair", i, name)
	}
	m := map[string]int{"b": 2, "a": 1}
	fmt.Println("keys", slices.Sorted(maps.Keys(m)))
	copied := maps.Collect(maps.All(m))
	fmt.Println("map equal", maps.Equal(m, copied))
	xs := []int{10, 20}
	live := slices.Values(xs)
	xs[0] = 99
	fmt.Println("live", slices.Collect(live))

	lines := Lines(strings.NewReader("first\nsecond\n"))
	for line, err := range lines {
		fmt.Println("line", line, err)
	}
	again := 0
	for range lines {
		again++
	}
	fmt.Println("second traversal", again)

	next, stop := iter.Pull(Count(3))
	defer stop()
	v, ok := next()
	fmt.Println("next", v, ok)
	stop()
	v, ok = next()
	fmt.Println("after stop", v, ok)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	for n := range UntilCanceled(ctx, 100) {
		fmt.Println("cancel demo", n)
		cancel()
	}
}