package log

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"

	"github.com/charmbracelet/log"
	"go.opentelemetry.io/contrib/bridges/otelslog"
	otellog "go.opentelemetry.io/otel/log"
	"go.opentelemetry.io/otel/trace"
)

type handlerOperation struct {
	attrs []slog.Attr
	group string
}

type namedHandler struct {
	slog.Handler
	name       string
	format     string
	writer     io.Writer
	hasTraceID bool
	hasSpanID  bool
	operations []handlerOperation
}

func (h *namedHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	hasTraceID := h.hasTraceID
	hasSpanID := h.hasSpanID
	for _, attr := range attrs {
		switch attr.Key {
		case "trace_id":
			hasTraceID = true
		case "span_id":
			hasSpanID = true
		}
	}
	operations := append(h.operations[:len(h.operations):len(h.operations)], handlerOperation{
		attrs: append([]slog.Attr(nil), attrs...),
	})
	return &namedHandler{
		Handler:    h.Handler.WithAttrs(attrs),
		name:       h.name,
		format:     h.format,
		writer:     h.writer,
		hasTraceID: hasTraceID,
		hasSpanID:  hasSpanID,
		operations: operations,
	}
}

func (h *namedHandler) WithGroup(name string) slog.Handler {
	operations := append(h.operations[:len(h.operations):len(h.operations)], handlerOperation{group: name})
	return &namedHandler{
		Handler:    h.Handler.WithGroup(name),
		name:       h.name,
		format:     h.format,
		writer:     h.writer,
		hasTraceID: h.hasTraceID,
		hasSpanID:  h.hasSpanID,
		operations: operations,
	}
}

func (h *namedHandler) WithName(name string) slog.Handler {
	next, err := newHandler(h.writer, name, h.format)
	if err != nil {
		return h
	}
	for _, operation := range h.operations {
		if operation.group != "" {
			next = next.WithGroup(operation.group)
		} else {
			next = next.WithAttrs(operation.attrs)
		}
	}
	return next
}

func (h *namedHandler) Handle(ctx context.Context, r slog.Record) error {
	spanContext := trace.SpanContextFromContext(ctx)
	if spanContext.IsValid() && spanContext.IsSampled() {
		hasTraceID := h.hasTraceID
		hasSpanID := h.hasSpanID
		if !hasTraceID || !hasSpanID {
			r.Attrs(func(a slog.Attr) bool {
				if a.Key == "trace_id" {
					hasTraceID = true
				}
				if a.Key == "span_id" {
					hasSpanID = true
				}
				return !hasTraceID || !hasSpanID
			})
		}
		var attrs []slog.Attr
		if !hasTraceID {
			attrs = append(attrs, slog.String("trace_id", spanContext.TraceID().String()))
		}
		if !hasSpanID {
			attrs = append(attrs, slog.String("span_id", spanContext.SpanID().String()))
		}
		if len(attrs) > 0 {
			r.AddAttrs(attrs...)
		}
	}
	return h.Handler.Handle(ctx, r)
}

type fanoutHandler struct {
	primary slog.Handler
	others  []slog.Handler
}

func (h *fanoutHandler) Enabled(ctx context.Context, level slog.Level) bool {
	if h.primary.Enabled(ctx, level) {
		return true
	}
	for _, other := range h.others {
		if other.Enabled(ctx, level) {
			return true
		}
	}
	return false
}

func (h *fanoutHandler) Handle(ctx context.Context, r slog.Record) error {
	var errs []error
	if h.primary.Enabled(ctx, r.Level) {
		if err := h.primary.Handle(ctx, r); err != nil {
			errs = append(errs, err)
		}
	}
	for _, other := range h.others {
		if other.Enabled(ctx, r.Level) {
			if err := other.Handle(ctx, r); err != nil {
				errs = append(errs, err)
			}
		}
	}
	if len(errs) > 0 {
		return errors.Join(errs...)
	}
	return nil
}

func (h *fanoutHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	others := make([]slog.Handler, len(h.others))
	for i, other := range h.others {
		others[i] = other.WithAttrs(attrs)
	}
	return &fanoutHandler{
		primary: h.primary.WithAttrs(attrs),
		others:  others,
	}
}

func (h *fanoutHandler) WithGroup(name string) slog.Handler {
	others := make([]slog.Handler, len(h.others))
	for i, other := range h.others {
		others[i] = other.WithGroup(name)
	}
	return &fanoutHandler{
		primary: h.primary.WithGroup(name),
		others:  others,
	}
}

func (h *fanoutHandler) WithName(name string) slog.Handler {
	newPrimary := h.primary
	if nh, ok := h.primary.(interface{ WithName(string) slog.Handler }); ok {
		newPrimary = nh.WithName(name)
	}
	others := make([]slog.Handler, len(h.others))
	for i, other := range h.others {
		if nh, ok := other.(interface{ WithName(string) slog.Handler }); ok {
			others[i] = nh.WithName(name)
		} else {
			others[i] = other
		}
	}
	return &fanoutHandler{
		primary: newPrimary,
		others:  others,
	}
}

func NewFanoutHandler(primary slog.Handler, others ...slog.Handler) slog.Handler {
	var validOthers []slog.Handler
	for _, other := range others {
		if other != nil {
			validOthers = append(validOthers, other)
		}
	}
	if len(validOthers) == 0 {
		return primary
	}
	return &fanoutHandler{
		primary: primary,
		others:  validOthers,
	}
}

type otelHandler struct {
	slog.Handler
	provider   otellog.LoggerProvider
	name       string
	operations []handlerOperation
}

func (h *otelHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	operations := append(h.operations[:len(h.operations):len(h.operations)], handlerOperation{
		attrs: append([]slog.Attr(nil), attrs...),
	})
	return &otelHandler{
		Handler:    h.Handler.WithAttrs(attrs),
		provider:   h.provider,
		name:       h.name,
		operations: operations,
	}
}

func (h *otelHandler) WithGroup(name string) slog.Handler {
	operations := append(h.operations[:len(h.operations):len(h.operations)], handlerOperation{group: name})
	return &otelHandler{
		Handler:    h.Handler.WithGroup(name),
		provider:   h.provider,
		name:       h.name,
		operations: operations,
	}
}

func (h *otelHandler) WithName(newName string) slog.Handler {
	var opts []otelslog.Option
	if h.provider != nil {
		opts = append(opts, otelslog.WithLoggerProvider(h.provider))
	}
	var handler slog.Handler = otelslog.NewHandler(newName, opts...)
	for _, operation := range h.operations {
		if operation.group != "" {
			handler = handler.WithGroup(operation.group)
		} else {
			handler = handler.WithAttrs(operation.attrs)
		}
	}
	return &otelHandler{
		Handler:    handler,
		provider:   h.provider,
		name:       newName,
		operations: h.operations,
	}
}

func NewOtelHandler(provider otellog.LoggerProvider, name string) slog.Handler {
	var opts []otelslog.Option
	if provider != nil {
		opts = append(opts, otelslog.WithLoggerProvider(provider))
	}
	h := otelslog.NewHandler(name, opts...)
	return &otelHandler{
		Handler:  h,
		provider: provider,
		name:     name,
	}
}

func newHandler(w io.Writer, name, format string) (slog.Handler, error) {
	var handler slog.Handler
	switch format {
	case "", "text":
		handler = log.NewWithOptions(w, log.Options{
			ReportTimestamp: true,
			Prefix:          name,
			Level:           log.DebugLevel,
		})
		format = "text"
	case "json":
		handler = slog.NewJSONHandler(w, &slog.HandlerOptions{
			Level: slog.LevelDebug,
		}).WithAttrs([]slog.Attr{slog.String("logger", name)})
	default:
		return nil, fmt.Errorf("unknown log format %q", format)
	}
	return &namedHandler{
		Handler: handler,
		name:    name,
		format:  format,
		writer:  w,
	}, nil
}

func NewHandler(name string) slog.Handler {
	handler, _ := newHandler(os.Stderr, name, "text")
	return handler
}

func New(name string) *slog.Logger {
	return slog.New(NewHandler(name))
}

func NewWithFormat(name, format string) (*slog.Logger, error) {
	return newWithFormatWriter(os.Stderr, name, format)
}

func newWithFormatWriter(w io.Writer, name, format string) (*slog.Logger, error) {
	handler, err := newHandler(w, name, format)
	if err != nil {
		return nil, err
	}
	return slog.New(handler), nil
}

func NewContext(ctx context.Context, name string) context.Context {
	return IntoContext(ctx, New(name))
}

type ctxKey struct{}

// IntoContext adds a logger to a context. Use FromContext to
// pull the logger out.
func IntoContext(ctx context.Context, l *slog.Logger) context.Context {
	return context.WithValue(ctx, ctxKey{}, l)
}

// FromContext returns a logger from a context.Context;
// if the passed context is nil, we return the default slog
// logger.
func FromContext(ctx context.Context) *slog.Logger {
	if ctx != nil {
		v := ctx.Value(ctxKey{})
		if v == nil {
			return slog.Default()
		}
		return v.(*slog.Logger)
	}

	return slog.Default()
}

// sublogger derives a new logger from an existing one by appending a suffix to its prefix.
func SubLogger(base *slog.Logger, suffix string) *slog.Logger {
	if base == nil {
		return New(suffix)
	}
	if handler, ok := base.Handler().(*namedHandler); ok {
		name := suffix
		if handler.name != "" {
			name = handler.name + "/" + suffix
		}
		return slog.New(handler.WithName(name))
	}
	if fh, ok := base.Handler().(*fanoutHandler); ok {
		var newName string
		if ph, ok := fh.primary.(*namedHandler); ok {
			if ph.name != "" {
				newName = ph.name + "/" + suffix
			} else {
				newName = suffix
			}
		} else {
			newName = suffix
		}
		return slog.New(fh.WithName(newName))
	}
	if handler, ok := base.Handler().(*log.Logger); ok {
		name := suffix
		if handler.GetPrefix() != "" {
			name = handler.GetPrefix() + "/" + suffix
		}
		return New(name)
	}
	return base.With("component", suffix)
}