package main

import "fmt"

// ------------------------------------------------

type F1[T any, R any] interface {
	Apply(T) R
}

type F2[T1 any, T2 any, R any] interface {
	Apply(T1, T2) R
	AsF1() F1[Tuple2[T1, T2], R]
}

type Named interface {
	Name() string
}

type NamedF1[T any, R any] struct {
	name string
	F1[T, R]
}

func (f NamedF1[T, R]) Name() string {
	return f.name
}

type NamedF2[T1 any, T2 any, R any] struct {
	name string
	F2[T1, T2, R]
}

func (f NamedF2[T1, T2, R]) Name() string {
	return f.name
}

func (f NamedF2[T1, T2, R]) AsF1() F1[Tuple2[T1, T2], R] {
	return NamedF1[Tuple2[T1, T2], R]{f.name, f.F2.AsF1()}
}

type F1Func[T any, R any] func(T) R

type F2Func[T1 any, T2 any, R any] func(T1, T2) R

func (f F1Func[T, R]) Named(name string) NamedF1[T, R] {
	return NamedF1[T, R]{name, f}
}

func (f F1Func[T, R]) Apply(t T) R {
	return f(t)
}

func (f F2Func[T1, T2, R]) Named(name string) NamedF2[T1, T2, R] {
	return NamedF2[T1, T2, R]{name, f}
}

func (f F2Func[T1, T2, R]) Apply(t1 T1, t2 T2) R {
	return f(t1, t2)
}

func (f F2Func[T1, T2, R]) AsF1() F1[Tuple2[T1, T2], R] {
	return Func1(func(t Tuple2[T1, T2]) R { return f(t._1, t._2) })
}

func Func1[T any, R any](f func(T) R) F1Func[T, R] {
	return F1Func[T, R](f)
}

func Func2[T1 any, T2 any, R any](f func(T1, T2) R) F2Func[T1, T2, R] {
	return F2Func[T1, T2, R](f)
}

func Compose[T1 any, T2 any, T3 any](f F1[T2, T3], g F1[T1, T2]) F1[T1, T3] {
	var name string
	if named, ok := f.(Named); ok {
		if named2, ok := g.(Named); ok {
			name = named2.Name() + "." + named.Name()
		}
	}
	f1 := Func1(func(x T1) T3 { return f.Apply(g.Apply(x)) })
	if name != "" {
		return f1.Named(name)
	}
	return f1
}

func JoinEq[T any, U any, R comparable](f F1[T, R], g F1[U, R]) F2[T, U, bool] {
	var name string
	if named, ok := f.(Named); ok {
		if named2, ok := g.(Named); ok {
			name = named.Name() + " == " + named2.Name()
		}
	}
	f2 := Func2(func(x T, y U) bool { return f.Apply(x) == g.Apply(y) })
	if name != "" {
		return f2.Named(name)
	}
	return f2
}

// ------------------------------------------------

type Iter[T any] interface {
	Next() (T, bool)
}

type IterFunc[T any] func() (T, bool)

func (f IterFunc[T]) Next() (T, bool) {
	return f()
}

type Basic[T any] interface {
	Iter() Iter[T]
}

type BasicT[T any] struct {
	Basic[T]
}

func (b BasicT[T]) Named(name string) NamedBasicT[T] {
	return NamedBasicT[T]{name, b}
}

type NamedBasicT[T any] struct {
	name string
	BasicT[T]
}

func (b BasicT[T]) ToSlice() []T {
	var elems []T
	iter := b.Iter()
	for {
		elem, ok := iter.Next()
		if !ok {
			break
		}
		elems = append(elems, elem)
	}
	return elems
}

type Basic2T[T1 any, T2 any] struct {
	BasicT[Tuple2[T1, T2]]
}

func (p Basic2T[T1, T2]) Where2(f F2[T1, T2, bool]) Basic2T[T1, T2] {
	return Basic2T[T1, T2]{BasicT[Tuple2[T1, T2]]{WhereImpl[Tuple2[T1, T2]]{p, f.AsF1()}}}
}

type WhereT[T any] struct {
	BasicT[T]
}

type WhereImpl[T any] struct {
	src Basic[T]
	f   F1[T, bool]
}

func (w WhereImpl[T]) Iter() Iter[T] {
	iter := w.src.Iter()
	return IterFunc[T](func() (T, bool) {
		for {
			elem, ok := iter.Next()
			if !ok {
				var empty T
				return empty, false
			}
			if w.f.Apply(elem) {
				return elem, true
			}
		}
	})
}

func (b BasicT[T]) Where(f F1[T, bool]) WhereT[T] {
	return WhereT[T]{BasicT[T]{WhereImpl[T]{b, f}}}
}

type FromT[T any] struct {
	BasicT[T]
}

type FromSliceImpl[T any] struct {
	elems []T
}

func (f FromSliceImpl[T]) Iter() Iter[T] {
	state := f.elems
	return IterFunc[T](func() (T, bool) {
		if len(state) == 0 {
			var empty T
			return empty, false
		}
		elem := state[0]
		state = state[1:]
		return elem, true
	})
}

func FromSlice[T any](elems []T) FromT[T] {
	return FromT[T]{BasicT[T]{FromSliceImpl[T]{elems}}}
}

type Product2T[T1 any, T2 any] struct {
	Basic2T[T1, T2]
}

type Product2Impl[T1 any, T2 any] struct {
	src1 Basic[T1]
	src2 Basic[T2]
}

type Tuple2[T1 any, T2 any] struct {
	_1 T1
	_2 T2
}

func Fst[T1 any, T2 any]() F1[Tuple2[T1, T2], T1] {
	return Func1(func(t Tuple2[T1, T2]) T1 { return t._1 }).Named("_1")
}

func Snd[T1 any, T2 any]() F1[Tuple2[T1, T2], T2] {
	return Func1(func(t Tuple2[T1, T2]) T2 { return t._2 }).Named("_2")
}

func (p Product2Impl[T1, T2]) Iter() Iter[Tuple2[T1, T2]] {
	var elem1 T1
	var iter1 Iter[T1]
	var iter2 Iter[T2]
	return IterFunc[Tuple2[T1, T2]](func() (Tuple2[T1, T2], bool) {
		if iter1 == nil {
			iter1 = p.src1.Iter()
			var ok1 bool
			elem1, ok1 = iter1.Next()
			if !ok1 {
				var empty Tuple2[T1, T2]
				return empty, false
			}
		}
		if iter2 == nil {
			iter2 = p.src2.Iter()
		}
		elem2, ok2 := iter2.Next()
		if !ok2 {
			var ok1 bool
			elem1, ok1 = iter1.Next()
			if !ok1 {
				var empty Tuple2[T1, T2]
				return empty, false
			}
			iter2 = p.src2.Iter()
			elem2, ok2 = iter2.Next()
			if !ok2 {
				var empty Tuple2[T1, T2]
				return empty, false
			}
		}
		return Tuple2[T1, T2]{elem1, elem2}, true
	})
}

func Product2[T1 any, T2 any](from1 Basic[T1], from2 Basic[T2]) Product2T[T1, T2] {
	return Product2T[T1, T2]{Basic2T[T1, T2]{BasicT[Tuple2[T1, T2]]{Product2Impl[T1, T2]{from1, from2}}}}
}

type GroupByMapReduceT[T any, K comparable, V any, U any] struct {
	Basic2T[K, U]
}

type GroupByMapReduceImpl[T any, K comparable, V any, U any] struct {
	src    Basic[T]
	key    F1[T, K]
	value  F1[T, V]
	reduce F2[K, []V, U]
}

// ugly: not lazy
func (g GroupByMapReduceImpl[T, K, V, U]) Iter() Iter[Tuple2[K, U]] {
	groups := make(map[K][]V)
	iter := g.src.Iter()
	for {
		elem, ok := iter.Next()
		if !ok {
			break
		}
		k := g.key.Apply(elem)
		v := g.value.Apply(elem)
		groups[k] = append(groups[k], v)
	}
	return IterFunc[Tuple2[K, U]](func() (Tuple2[K, U], bool) {
		for k, elems := range groups {
			delete(groups, k)
			return Tuple2[K, U]{k, g.reduce.Apply(k, elems)}, true
		}
		var empty Tuple2[K, U]
		return empty, false
	})
}

func GroupByMapReduceKey[T any, K comparable, V any, U any](src Basic[T], key F1[T, K], value F1[T, V], reduce F2[K, []V, U]) GroupByMapReduceT[T, K, V, U] {
	return GroupByMapReduceT[T, K, V, U]{Basic2T[K, U]{BasicT[Tuple2[K, U]]{GroupByMapReduceImpl[T, K, V, U]{src, key, value, reduce}}}}
}

func GroupByMapReduce[T any, K comparable, V any, U any](src Basic[T], key F1[T, K], value F1[T, V], reduce F1[[]V, U]) GroupByMapReduceT[T, K, V, U] {
	return GroupByMapReduceKey(src, key, value, Func2(func(_ K, vs []V) U { return reduce.Apply(vs) }))
}

// ------------------------------------------------

type User struct {
	ID   int
	Name string
}

func UserID() F1[User, int] {
	return Func1(func(u User) int { return u.ID }).Named("User.ID")
}

type Post struct {
	ID      int
	UserID  int
	Content string
	Rating  int
}

func PostUserID() F1[Post, int] {
	return Func1(func(p Post) int { return p.UserID }).Named("Post.UserID")
}

func PostRating() F1[Post, int] {
	return Func1(func(p Post) int { return p.Rating }).Named("Post.Rating")
}

type Numeric interface {
	~int | ~float64 // etc.
}

func SliceSum[T Numeric](elems []T) T {
	var sum T
	for _, elem := range elems {
		sum += elem
	}
	return sum
}

func AggregateSum[T Numeric]() F1[[]T, T] {
	return Func1[[]T, T](SliceSum[T]).Named("Sum")
}

// ------------------------------------------------

type SQNode interface{}

type SQNamed struct {
	Name string
	Node SQNode
}

type SQValues struct {
	Values interface{}
}

type SQFrom struct {
	Sources []SQNode
}

type SQCond struct {
	Cond interface{}
}

type SQWhere struct {
	Source SQNode
	Cond   SQNode
}

type SQGroupBy struct {
	Source SQNode
	Keys   []interface{}
}

// ------------------------------------------------

func TryName(f interface{}) string {
	if named, ok := f.(Named); ok {
		return named.Name()
	}
	return "(code)"
}

type SQL interface {
	Compile() SQNode
}

func (b BasicT[T]) Compile() SQNode {
	return b.Basic.(SQL).Compile()
}

func (f NamedBasicT[T]) Compile() SQNode {
	return SQNamed{f.name, f.Basic.(SQL).Compile()}
}

func (f FromT[T]) Compile() SQNode {
	return f.Basic.(SQL).Compile()
}

func (p Product2T[T1, T2]) Compile() SQNode {
	return p.Basic.(SQL).Compile()
}

func (g GroupByMapReduceT[T, K, V, U]) Compile() SQNode {
	return g.Basic.(SQL).Compile()
}

func (w WhereImpl[T]) Compile() SQNode {
	return SQWhere{w.src.(SQL).Compile(), SQCond{TryName(w.f)}}
}

func (p Product2Impl[T1, T2]) Compile() SQNode {
	return SQFrom{[]SQNode{p.src1.(SQL).Compile(), p.src2.(SQL).Compile()}}
}

func (s FromSliceImpl[T]) Compile() SQNode {
	return SQFrom{[]SQNode{SQValues{s.elems}}}
}

func (g GroupByMapReduceImpl[T, K, V, U]) Compile() SQNode {
	return SQGroupBy{g.src.(SQL).Compile(), []interface{}{TryName(g.key)}}
}

// ------------------------------------------------

func Join(ss []string, sep string) string {
	out := ""
	for i, s := range ss {
		if i > 0 {
			out += sep
		}
		out += s
	}
	return out
}

func FormatSQL(sql SQNode) string {
	switch sql := sql.(type) {
	case SQNamed:
		// to be pattern matched inside nameable sub-expressions
		return FormatSQL(sql.Node)
	case SQValues:
		return fmt.Sprintf("VALUES %v", sql.Values)
	case SQFrom:
		var sources []string
		for _, source := range sql.Sources {
			var alias string
			if named, ok := source.(SQNamed); ok {
				alias = named.Name
				source = named.Node
			}
			if from, ok := source.(SQFrom); ok {
				if len(from.Sources) == 1 {
					source = from.Sources[0]
				}
			}
			s := fmt.Sprintf("(%v)", FormatSQL(source))
			if alias != "" {
				s += " AS " + alias
			}
			sources = append(sources, s)
		}
		return fmt.Sprintf("FROM %v", Join(sources, ", "))
	case SQCond:
		return fmt.Sprintf("%v", sql.Cond)
	case SQWhere:
		return fmt.Sprintf("%v WHERE %v", FormatSQL(sql.Source), FormatSQL(sql.Cond))
	case SQGroupBy:
		keys := make([]string, len(sql.Keys))
		for i, key := range sql.Keys {
			keys[i] = fmt.Sprintf("%v", key)
		}
		return fmt.Sprintf("%v GROUP BY %v", FormatSQL(sql.Source), Join(keys, ", "))
	}
	return ""
}

// ------------------------------------------------

func main() {
	users := []User{
		{ID: 1, Name: "Alice"},
		{ID: 2, Name: "Bob"},
	}
	posts := []Post{
		{ID: 1, UserID: 1, Content: "Hello", Rating: 5},
		{ID: 2, UserID: 2, Content: "World", Rating: 3},
		{ID: 3, UserID: 1, Content: "Bye", Rating: 4},
		{ID: 4, UserID: 2, Content: "Moon", Rating: 2},
	}

	q := GroupByMapReduce(
		Product2(FromSlice(users).Named("users"), FromSlice(posts).Named("posts")).
			Where2(JoinEq(UserID(), PostUserID())),
		Compose(UserID(), Fst[User, Post]()),
		Compose(PostRating(), Snd[User, Post]()),
		AggregateSum[int]())

	fmt.Println(FormatSQL(q.Compile()))

	fmt.Printf("%v\n", q.ToSlice())
}