package logging

import (
	"encoding/json"
	"fmt"
	"iter"
	"log/slog"
	"runtime"
	"slices"
	"strings"
)

func RecordAll(r slog.Record, replaceAttr func(groups []string, a slog.Attr) slog.Attr) iter.Seq2[[]string, slog.Attr] {
	return func(yield func([]string, slog.Attr) bool) {
		var walk func([]string, slog.Attr) bool
		walk = func(groups []string, a slog.Attr) bool {
			if replaceAttr != nil {
				a = replaceAttr(groups, a)
			}

			if a.Key == "" {
				return true
			}

			a.Value = a.Value.Resolve()

			if a.Value.Kind() == slog.KindGroup {
				newGroups := append(slices.Clone(groups), a.Key)
				for _, child := range a.Value.Group() {
					if !walk(newGroups, child) {
						return false
					}
				}
				return true
			}

			return yield(groups, a)
		}

		if !r.Time.IsZero() {
			if !walk([]string{}, slog.Time(slog.TimeKey, r.Time)) {
				return
			}
		}

		if !walk([]string{}, slog.Any(slog.LevelKey, r.Level)) {
			return
		}
		if !walk([]string{}, slog.String(slog.MessageKey, r.Message)) {
			return
		}

		if r.PC != 0 {
			if !walk([]string{}, slog.Uint64(slog.SourceKey, uint64(r.PC))) {
				return
			}
		}

		r.Attrs(func(attr slog.Attr) bool {
			return walk([]string{}, attr)
		})
	}
}

type Encoder func(iter.Seq2[[]string, slog.Attr]) (string, error)

func LogFmtEncoder(attrs iter.Seq2[[]string, slog.Attr]) (string, error) {
	str := []string{}
	for groups, attr := range attrs {
		if len(groups) == 0 && attr.Key == slog.SourceKey {
			pc := uintptr(attr.Value.Uint64())
			fs := runtime.CallersFrames([]uintptr{pc})
			f, _ := fs.Next()
			attr.Value = slog.StringValue(fmt.Sprintf("%s:%d", f.File, f.Line))
		}

		str = append(str, fmt.Sprintf("%s=%q", strings.Join(append(groups, attr.Key), "."), attr.Value))
	}

	return strings.Join(str, " "), nil
}

func ToMap(seq iter.Seq2[[]string, slog.Attr]) map[string]any {
	out := make(map[string]any)
	for groups, attr := range seq {
		current := out
		for _, group := range groups {
			if next, ok := current[group].(map[string]any); ok {
				current = next
			} else {
				newMap := make(map[string]any)
				current[group] = newMap
				current = newMap
			}
		}
		current[attr.Key] = attr.Value.Any()
	}
	return out
}

func JSONEncoder(attrs iter.Seq2[[]string, slog.Attr]) (string, error) {
	m := ToMap(attrs)
	if v, ok := m[slog.SourceKey]; ok {
		pc := v.(uint64)
		fs := runtime.CallersFrames([]uintptr{uintptr(pc)})
		f, _ := fs.Next()
		m[slog.SourceKey] = fmt.Sprintf("%s:%d", f.File, f.Line)
	}

	b, err := json.Marshal(m)
	if err != nil {
		return "", err
	}

	return string(b), nil
}

func MessageEncoder(attrs iter.Seq2[[]string, slog.Attr]) (string, error) {
	for groups, attr := range attrs {
		if len(groups) == 0 && attr.Key == slog.MessageKey {
			return attr.Value.String(), nil
		}
	}
	return "", nil
}