// Package logging provides helpers for context-aware logging using Go's slog package.
//
// It allows attaching loggers to context.Context for structured logging,
// recovering and logging panics with stack traces,
// and provides a simple package-level logger setup function.
package logging

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"os"
	"runtime"
	"strings"
)

type ctxLoggerKeyType struct{}

var ctxLoggerKey = ctxLoggerKeyType{}

// WithLogger returns a new context derived from ctx that carries the provided slog.Logger.
//
// Typical usage:
//
//	ctx = logging.WithLogger(ctx, logger)
func WithLogger(ctx context.Context, logger *slog.Logger) context.Context {
	return context.WithValue(ctx, ctxLoggerKey, logger)
}

// FromContext retrieves the slog.Logger stored in ctx if any.
// If no logger is attached to the context, returns slog.Default().
func FromContext(ctx context.Context) *slog.Logger {
	if logger, ok := ctx.Value(ctxLoggerKey).(*slog.Logger); ok {
		return logger
	}
	return slog.Default()
}

// With returns a new slog.Logger augmented with the passed key-value pairs,
// and a new context holding this logger.
//
// Typical usage:
//
//	logger, ctx := logging.With(ctx, "user_id", userID)
func With(ctx context.Context, args ...any) (*slog.Logger, context.Context) {
	logger := FromContext(ctx)
	logger = logger.With(args...)
	ctx = WithLogger(ctx, logger)
	return logger, ctx
}

// WithGroup returns a new slog.Logger grouped under name,
// and a new context holding this logger.
//
// Typical usage:
//
//	logger, ctx := logging.WithGroup(ctx, "http")
func WithGroup(ctx context.Context, name string) (*slog.Logger, context.Context) {
	logger := FromContext(ctx)
	logger = logger.WithGroup(name)
	ctx = WithLogger(ctx, logger)
	return logger, ctx
}

// RecoverAndLog logs a recovered panic or error with a stack trace using the logger from ctx.
//
// Typically used in a deferred recover block:
//
//	defer func() {
//	  if err := recover(); err != nil {
//	    logging.RecoverAndLog(ctx, "Panic recovered", err)
//	    os.Exit(1)
//	  }
//	}()
func RecoverAndLog(ctx context.Context, msg string, err any) {
	const size = 64 << 10 // 64 KB stack trace buffer
	buf := make([]byte, size)
	buf = buf[:runtime.Stack(buf, false)]
	logger := FromContext(ctx)
	logger.ErrorContext(ctx, msg, "Error", err, "stack", string(buf))
}

type LogSink string

const (
	SinkStdout LogSink = "stdout"
	SinkStderr LogSink = "stderr"
	SinkFile   LogSink = "file"
)

type LogConfig struct {
	Level  slog.Level `toml:"level"`
	Sink   LogSink    `toml:"sink"`
	Format string     `toml:"format"`
	File   string     `toml:"file"`
}

func Setup(cfg LogConfig) error {
	var w io.Writer

	switch cfg.Sink {
	case SinkStdout, "":
		w = os.Stdout
	case SinkStderr:
		w = os.Stderr
	case SinkFile:
		if cfg.File == "" {
			return fmt.Errorf("log sink set to 'file' but no file path provided")
		}
		f, err := os.OpenFile(cfg.File, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			return fmt.Errorf("failed to open log file: %w", err)
		}
		w = f
	default:
		return fmt.Errorf("unsupported log sink: %s", cfg.Sink)
	}

	var handler slog.Handler
	opts := &slog.HandlerOptions{
		Level: cfg.Level,
	}

	switch strings.ToLower(cfg.Format) {
	case "text":
		handler = slog.NewTextHandler(w, opts)
	case "json", "":
		handler = slog.NewJSONHandler(w, opts)
	default:
		return fmt.Errorf("unsupported log format: %s", cfg.Format)
	}

	slog.SetDefault(slog.New(handler))
	return nil
}