package logging

import (
	"context"
	"iter"
	"log/slog"
	"slices"
)

func attrAll(r slog.Record) iter.Seq[slog.Attr] {
	return func(yield func(slog.Attr) bool) {
		r.Attrs(func(attr slog.Attr) bool {
			return yield(attr)
		})
	}
}

// GetNormalizedAttrs extracts, resolves, filters, and flattens all attributes
// from a slog.Record into a ready-to-process slice of slog.Attr.
// This helper is essential for any custom slog.Handler implementation
// that needs to process the full set of attributes.
func GetNormalizedAttrs(r slog.Record) []slog.Attr {
	attrs := attrAll(r)

	return normalizeAttrsSlice(attrs)
}

// normalizeAttrsSlice contains the core logic for resolving, filtering, and flattening.
// This function must be private as it is a utility for GetNormalizedAttrs.
func normalizeAttrsSlice(attrs iter.Seq[slog.Attr]) []slog.Attr {
	var out []slog.Attr
	for attr := range attrs {
		if attr.Equal(slog.Attr{}) {
			continue
		}

		v := attr.Value.Resolve()

		switch v.Kind() {
		case slog.KindGroup:
			// Skip empty group
			if len(v.Group()) == 0 {
				continue
			}

			inner := normalizeAttrsSlice(slices.Values(v.Group()))

			// Flatten anonymous group (Key == "")
			if attr.Key == "" {
				out = append(out, inner...)
			} else {
				// Named group, keep as a GroupValue
				out = append(out, slog.Attr{
					Key:   attr.Key,
					Value: slog.GroupValue(inner...),
				})
			}
		default:
			// Regular attribute
			out = append(out, slog.Attr{
				Key:   attr.Key,
				Value: v,
			})
		}
	}
	return out
}

type chainingAttrHandler struct {
	parent slog.Handler
	attrs  []slog.Attr
}

var _ slog.Handler = (*chainingAttrHandler)(nil)

// ChainAttrs stores attributes/groups added via WithAttrs/WithGroup
// and ensures they are applied to the record before passing it to the next handler.
func ChainAttrs(h slog.Handler, attrs []slog.Attr) slog.Handler {
	if len(attrs) == 0 {
		return h
	}
	return &chainingAttrHandler{
		parent: h,
		attrs:  attrs,
	}
}

func (h *chainingAttrHandler) Enabled(ctx context.Context, level slog.Level) bool {
	return h.parent.Enabled(ctx, level)
}

func (h *chainingAttrHandler) Handle(ctx context.Context, r slog.Record) error {
	newRecord := slog.NewRecord(r.Time, r.Level, r.Message, r.PC)

	attrs := slices.Collect(attrAll(r))
	newRecord.AddAttrs(h.attrs...)
	newRecord.AddAttrs(attrs...)

	return h.parent.Handle(ctx, newRecord)
}

func (h *chainingAttrHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return ChainAttrs(h, attrs)

}

func (h *chainingAttrHandler) WithGroup(name string) slog.Handler {
	return ChainGroup(h, name)
}

type chainingGroupHandler struct {
	parent slog.Handler
	group  string
}

// ChainGroup stores attributes/groups added via WithAttrs/WithGroup
// and ensures they are applied to the record before passing it to the next handler.
func ChainGroup(h slog.Handler, name string) slog.Handler {
	if name == "" {
		return h
	}
	return &chainingGroupHandler{
		parent: h,
		group:  name,
	}
}

var _ slog.Handler = (*chainingGroupHandler)(nil)

func (h *chainingGroupHandler) Enabled(ctx context.Context, level slog.Level) bool {
	return h.parent.Enabled(ctx, level)
}

func (h *chainingGroupHandler) Handle(ctx context.Context, r slog.Record) error {
	newRecord := slog.NewRecord(r.Time, r.Level, r.Message, r.PC)

	newRecord.AddAttrs(slog.Attr{
		Key:   h.group,
		Value: slog.GroupValue(GetNormalizedAttrs(r)...),
	})

	return h.parent.Handle(ctx, newRecord)
}

func (h *chainingGroupHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return ChainAttrs(h, attrs)

}

func (h *chainingGroupHandler) WithGroup(name string) slog.Handler {
	return ChainGroup(h, name)
}