package gcp

import (
	"context"
	"io"
	"log/slog"
	"os"
)

// LevelCritical is an extra log level supported by Cloud Logging.
const LevelCritical = slog.Level(12)

// Handler that outputs JSON understood by the structured log agent.
// See https://cloud.google.com/logging/docs/agent/logging/configuration#special-fields
type Handler struct {
	handler slog.Handler
}

// NewHandler returns a new Handler that writes to stderr.
func NewHandler(level slog.Level) *Handler {
	return NewHandlerForWriter(os.Stderr, level)
}

// NewHandlerForWriter returns a new Handler that writes to the given writer.
func NewHandlerForWriter(w io.Writer, level slog.Level) *Handler {
	return &Handler{handler: slog.NewJSONHandler(w, &slog.HandlerOptions{
		AddSource: true,
		Level:     level,
		ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
			if a.Key == slog.MessageKey {
				a.Key = "message"
			} else if a.Key == slog.SourceKey {
				a.Key = "logging.googleapis.com/sourceLocation"
			} else if a.Key == slog.LevelKey {
				a.Key = "severity"
				level, ok := a.Value.Any().(slog.Level)
				if ok {
				switch {
					case level == LevelCritical:
						a.Value = slog.StringValue("CRITICAL")
					case level == slog.LevelWarn:
						// Convert WARN -> WARNING for Cloud Batch
						a.Value = slog.StringValue("WARNING")
					}
				}
			}
			return a
		},
	})}
}

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

func (h *Handler) Handle(ctx context.Context, rec slog.Record) error {
	if trace := TraceFromContext(ctx); trace != "" {
		rec = rec.Clone()
		// Add trace ID	to the record so it is correlated with the request log
		// See https://cloud.google.com/trace/docs/trace-log-integration
		rec.Add("logging.googleapis.com/trace", slog.StringValue(trace))
	}

	return h.handler.Handle(ctx, rec)
}

func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return &Handler{handler: h.handler.WithAttrs(attrs)}
}

func (h *Handler) WithGroup(name string) slog.Handler {
	return &Handler{handler: h.handler.WithGroup(name)}
}
