// Package accesslog writes one JSON line per notable DNS query to a stream
// (stdout by default), for whatever collects the pod's logs to index.
//
// It is deliberately not a log store: nothing is kept, searched or aggregated
// here. Counting belongs to the metrics, searching to the log collector. What
// this package guarantees is that logging can never slow or stop a query:
// lines go through a bounded buffer to a single writer goroutine, and when the
// buffer is full (a blocked pipe, a stalled collector) the line is dropped and
// counted rather than waited for.
package accesslog

import (
	"context"
	"io"
	"log/slog"
	"math/rand/v2"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/metrics"
)

// Action is the outcome classification of a query.
type Action string

const (
	// ActionAllow is a normal resolution.
	ActionAllow Action = "allow"
	// ActionBlock is a blocklist hit.
	ActionBlock Action = "block"
	// ActionRedirect is a blocklist rewrite to an address or CNAME target.
	ActionRedirect Action = "redirect"
	// ActionCached is a cache hit.
	ActionCached Action = "cached"
	// ActionError is a SERVFAIL or other failure.
	ActionError Action = "error"
	// ActionDrop is a query answered with silence.
	ActionDrop Action = "drop"
	// ActionRefused is an ACL refusal.
	ActionRefused Action = "refused"
)

// Entry is one logged query.
type Entry struct {
	TS       time.Time
	Client   string
	Group    string
	Name     string
	QType    uint16
	RCode    int
	Action   Action
	Duration time.Duration
	Upstream string
	List     string
	Protocol string
}

// FromMsg fills the question-derived fields from a DNS message.
func (e *Entry) FromMsg(m *dns.Msg) {
	if len(m.Question) > 0 {
		q := m.Question[0]
		e.Name = q.Name
		e.QType = q.Qtype
	}
	e.RCode = m.Rcode
}

// Options configures a Logger.
type Options struct {
	// Writer receives the lines; nil means os.Stdout.
	Writer io.Writer
	// Buffer is how many lines may wait for the writer; beyond that they are
	// dropped. Zero means 4096.
	Buffer int
	// LogBlocked, LogErrors and LogCached select whole categories. Refusals
	// and drops count with errors, redirects with blocks.
	LogBlocked bool
	LogErrors  bool
	LogCached  bool
	// SampleRate is the fraction of ordinary answers (allowed, and cached when
	// LogCached is off) to log. 0 logs none, 1 logs all.
	SampleRate float64
}

// Logger is the access log.
type Logger struct {
	opt     Options
	h       slog.Handler
	w       *asyncWriter
	written atomic.Uint64
}

// New starts a Logger. Close it to flush.
func New(opt Options) *Logger {
	if opt.Writer == nil {
		opt.Writer = os.Stdout
	}
	aw := newAsyncWriter(opt.Writer, opt.Buffer)
	h := slog.NewJSONHandler(aw, &slog.HandlerOptions{
		Level: slog.LevelInfo,
		// The level carries no information here, and a collector would index
		// a constant. "log":"access" says which stream the line belongs to.
		ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
			if len(groups) == 0 && a.Key == slog.LevelKey {
				return slog.Attr{}
			}
			return a
		},
	}).WithAttrs([]slog.Attr{slog.String("log", "access")})
	return &Logger{opt: opt, h: h, w: aw}
}

// ShouldLog reports whether a query with this outcome is logged. A nil Logger
// logs nothing.
func (l *Logger) ShouldLog(a Action) bool {
	if l == nil {
		return false
	}
	switch a {
	case ActionBlock, ActionRedirect:
		return l.opt.LogBlocked
	case ActionError, ActionDrop, ActionRefused:
		return l.opt.LogErrors
	case ActionCached:
		if l.opt.LogCached {
			return true
		}
	}
	switch {
	case l.opt.SampleRate >= 1:
		return true
	case l.opt.SampleRate <= 0:
		return false
	default:
		return rand.Float64() < l.opt.SampleRate
	}
}

// Log writes one line. It never blocks.
func (l *Logger) Log(e Entry) {
	if l == nil {
		return
	}
	ts := e.TS
	if ts.IsZero() {
		ts = time.Now()
	}
	rec := slog.NewRecord(ts, slog.LevelInfo, "query", 0)
	rec.AddAttrs(
		slog.String("client", e.Client),
		slog.String("client_group", e.Group),
		slog.String("name", e.Name),
		slog.String("qtype", dns.Type(e.QType).String()),
		slog.String("rcode", rcodeName(e.RCode)),
		slog.String("action", string(e.Action)),
		slog.Float64("duration_ms", float64(e.Duration.Microseconds())/1000),
	)
	if e.Upstream != "" {
		rec.AddAttrs(slog.String("upstream", e.Upstream))
	}
	if e.List != "" {
		rec.AddAttrs(slog.String("list", e.List))
	}
	if e.Protocol != "" {
		rec.AddAttrs(slog.String("protocol", e.Protocol))
	}
	_ = l.h.Handle(context.Background(), rec)
	l.written.Add(1)
}

func rcodeName(c int) string {
	if s, ok := dns.RcodeToString[c]; ok {
		return s
	}
	return "RCODE" + strconv.Itoa(c)
}

// Dropped is how many lines were discarded because the writer fell behind.
func (l *Logger) Dropped() uint64 {
	if l == nil {
		return 0
	}
	return l.w.dropped.Load()
}

// Flush waits for the lines logged so far to be written.
func (l *Logger) Flush() {
	if l != nil {
		l.w.flush(2 * time.Second)
	}
}

// Close writes what is buffered (for at most two seconds) and stops.
func (l *Logger) Close() error {
	if l == nil {
		return nil
	}
	l.w.close()
	return nil
}

// asyncWriter hands each Write to one goroutine through a bounded channel. A
// Write never waits: a full channel drops the line.
type item struct {
	b   []byte
	ack chan struct{} // set on a flush marker: closed when everything before it is written
}

type asyncWriter struct {
	ch      chan item
	w       io.Writer
	dropped atomic.Uint64
	closed  atomic.Bool
	stop    chan struct{}
	done    chan struct{}
	once    sync.Once
}

func newAsyncWriter(w io.Writer, buffer int) *asyncWriter {
	if buffer <= 0 {
		buffer = 4096
	}
	a := &asyncWriter{
		ch: make(chan item, buffer), w: w,
		stop: make(chan struct{}), done: make(chan struct{}),
	}
	go a.run()
	return a
}

// Write implements io.Writer for slog's handler, which calls it once per line
// with a buffer it reuses, so the bytes are copied.
func (a *asyncWriter) Write(p []byte) (int, error) {
	if a.closed.Load() {
		return len(p), nil
	}
	b := append(make([]byte, 0, len(p)), p...)
	select {
	case a.ch <- item{b: b}:
	default:
		a.dropped.Add(1)
		metrics.RecordAccessLogDropped()
	}
	return len(p), nil
}

func (a *asyncWriter) run() {
	defer close(a.done)
	for {
		select {
		case it := <-a.ch:
			a.write(it)
		case <-a.stop:
			for {
				select {
				case it := <-a.ch:
					a.write(it)
				default:
					return
				}
			}
		}
	}
}

func (a *asyncWriter) write(it item) {
	if it.ack != nil {
		close(it.ack)
		return
	}
	_, _ = a.w.Write(it.b)
}

// flush waits until every line accepted before the call has been written, or
// the timeout passes. It is for tests and shutdown, never the query path.
func (a *asyncWriter) flush(timeout time.Duration) {
	if a.closed.Load() {
		return
	}
	ack := make(chan struct{})
	t := time.NewTimer(timeout)
	defer t.Stop()
	select {
	case a.ch <- item{ack: ack}:
	case <-t.C:
		return
	}
	select {
	case <-ack:
	case <-t.C:
	}
}

func (a *asyncWriter) close() {
	a.once.Do(func() {
		a.closed.Store(true)
		close(a.stop)
		select {
		case <-a.done:
		case <-time.After(2 * time.Second):
		}
	})
}