// 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"
	"time"
)

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
}

type StackFrame struct {
	Function string `json:"function"`
	File     string `json:"file"`
	Line     int    `json:"line"`
}

// 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) {
	logger := FromContext(ctx)
	handler := logger.Handler()

	if !handler.Enabled(ctx, slog.LevelError) {
		return
	}

	// Capture stack
	const depth = 32
	pcs := make([]uintptr, depth)
	n := runtime.Callers(3, pcs) // Skip: Callers, RecoverAndLog, defer-fn
	frames := runtime.CallersFrames(pcs[:n])

	var stack []StackFrame
	for {
		frame, more := frames.Next()
		stack = append(stack, StackFrame{
			Function: frame.Function,
			File:     frame.File,
			Line:     frame.Line,
		})
		if !more {
			break
		}
	}
	var pc uintptr
	if n > 0 {
		pc = pcs[0]
	}
	r := slog.NewRecord(time.Now(), slog.LevelError, msg, pc)
	r.Add("error", err, "stack", stack)
	handler.Handle(ctx, r)
}

// LogSink defines where log output should be written.
//
// It is used to configure the log destination in Setup.
type LogSink int

const (
	// SinkStdout directs logs to the standard output (os.Stdout).
	SinkStdout LogSink = iota

	// SinkStderr directs logs to the standard error (os.Stderr).
	SinkStderr

	// SinkFile directs logs to a file specified by the File field in LogConfig.
	//
	// If configured, the File field must be a valid path for writing logs.
	SinkFile
)

func (s LogSink) String() string {
	switch s {
	case SinkStdout:
		return "stdout"
	case SinkStderr:
		return "stderr"
	case SinkFile:
		return "file"
	default:
		return fmt.Sprintf("LogSink(%d)", s)
	}
}

func (s *LogSink) Set(str string) error {
	switch strings.ToLower(str) {
	case "stdout":
		*s = SinkStdout
	case "stderr":
		*s = SinkStderr
	case "file":
		*s = SinkFile
	default:
		return fmt.Errorf("invalid LogSink %q", str)
	}
	return nil
}

func (s LogSink) AppendText(b []byte) ([]byte, error) {
	return append(b, s.String()...), nil
}

func (s LogSink) MarshalText() ([]byte, error) {
	return s.AppendText(nil)
}

func (s *LogSink) UnmarshalText(b []byte) error {
	return s.Set(string(b))
}

type LogFormat int

const (
	// SinkStdout directs logs to the standard output (os.Stdout).
	FormatText LogFormat = iota

	// SinkStderr directs logs to the standard error (os.Stderr).
	FormatJSON
)

func (s LogFormat) String() string {
	switch s {
	case FormatJSON:
		return "json"
	case FormatText:
		return "text"
	default:
		return fmt.Sprintf("LogSink(%d)", s)
	}
}

func (s *LogFormat) Set(str string) error {
	switch strings.ToLower(str) {
	case "text":
		*s = FormatText
	case "json":
		*s = FormatJSON
	default:
		return fmt.Errorf("invalid LogSink %q", str)
	}
	return nil
}

func (s LogFormat) AppendText(b []byte) ([]byte, error) {
	return append(b, s.String()...), nil
}

func (s LogFormat) MarshalText() ([]byte, error) {
	return s.AppendText(nil)
}

func (s *LogFormat) UnmarshalText(b []byte) error {
	return s.Set(string(b))
}

// LogConfig holds configuration for the logging setup.
//
// Level specifies the minimum level for logs (e.g., slog.LevelInfo).
//
// Sink specifies the output destination (e.g., SinkStdout, SinkStderr, SinkFile).
//
// Format specifies the log format, either "json" or "text".
// If empty or "json", JSON format is used.
//
// File specifies the output file path when SinkFile is selected.
// It is required if SinkFile is used.
type LogConfig struct {
	Level  slog.Level `toml:"level"`
	Sink   LogSink    `toml:"sink"`
	Format LogFormat  `toml:"format"`
	File   string     `toml:"file"`
}

func init() {
	Setup(LogConfig{}) // Setup Sensible Defaults
}

// Setup initializes the default slog.Logger based on the provided LogConfig.
//
// It configures the log output sink, log format, and log level.
//
// Returns an error if the configuration is invalid or if opening the log file fails.
//
// Example:
//
//	cfg := LogConfig{
//	    Level:  slog.LevelInfo,
//	    Sink:   SinkFile,
//	    Format: FormatJSON,
//	    File:   "/var/log/myapp.log",
//	}
//	if err := Setup(cfg); err != nil {
//	    log.Fatalf("failed to configure logging: %v", err)
//	}
//
// After a successful call, the global slog logger will log accordingly.
func Setup(cfg LogConfig) error {
	logger, err := New(cfg)
	if err != nil {
		return nil
	}
	slog.SetDefault(logger)
	return nil
}

// Same as Setup but returns the logger instead of setting up the global logger
func New(cfg LogConfig) (*slog.Logger, error) {
	var w io.Writer

	switch cfg.Sink {
	case SinkStdout:
		w = os.Stdout
	case SinkStderr:
		w = os.Stderr
	case SinkFile:
		if cfg.File == "" {
			return nil, 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 nil, fmt.Errorf("failed to open log file: %w", err)
		}
		w = f
	default:
		return nil, fmt.Errorf("invalid LogSink %v", cfg.Sink)
	}

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

	switch cfg.Format {
	case FormatText:
		handler = slog.NewTextHandler(w, opts)
	case FormatJSON:
		handler = slog.NewJSONHandler(w, opts)
	default:
		return nil, fmt.Errorf("invalid LogFormat: %v", cfg.Format)
	}

	logger := slog.New(handler)
	return logger, nil
}