package logger

import (
	"context"
	"errors"
	"log"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"

	"git.misaka.ren/M1saka/token_thief/db"
)

const (
	maxAttempts = 3
	BaseBackoff = 200 * time.Millisecond
)

type Batch interface {
	Append(v ...any) error
	Send() error
	Abort() error
}

type BatchPreparer interface {
	PrepareBatch(ctx context.Context, query string) (Batch, error)
}

type clickHouseBatchPreparer struct {
	conn driver.Conn
}

// The explicit adapter keeps the local Batch contract aligned with the real driver.
func (p clickHouseBatchPreparer) PrepareBatch(ctx context.Context, query string) (Batch, error) {
	return p.conn.PrepareBatch(ctx, query)
}

type Stats struct {
	Enqueued  uint64
	Dropped   uint64
	Failed    uint64
	Ambiguous uint64
	Bytes     int64
}

// FlushResult separates retry-safe Prepare failures from final or ambiguous failures.
type FlushResult struct {
	Retry     []*LogEntry
	Failed    int
	Ambiguous int
	Err       error
}

// Queue is an asynchronous, bounded logger. Submit never waits for database work.
type Queue struct {
	ch            chan *LogEntry
	pool          *db.Pool
	batchSize     int
	batchInterval time.Duration
	workers       int
	byteBudget    int64

	dropped   atomic.Uint64
	failed    atomic.Uint64
	ambiguous atomic.Uint64
	enq       atomic.Uint64
	bytes     atomic.Int64

	mu             sync.RWMutex
	started        bool
	stopped        bool
	shutdownCtx    context.Context
	reporterCancel context.CancelFunc
	workCancel     context.CancelFunc
	wg             sync.WaitGroup
	done           chan struct{}
	doneOnce       sync.Once
}

// NewQueue accepts an optional byte budget. A non-positive or omitted budget disables byte limiting.
func NewQueue(pool *db.Pool, queueSize, batchSize, workers int, batchInterval time.Duration, byteBudget ...int64) *Queue {
	var budget int64
	if len(byteBudget) > 0 {
		budget = byteBudget[0]
	}
	if queueSize < 0 {
		queueSize = 0
	}
	if batchSize < 1 {
		batchSize = 1
	}
	if workers < 1 {
		workers = 1
	}
	if batchInterval <= 0 {
		batchInterval = time.Second
	}
	return &Queue{
		ch:            make(chan *LogEntry, queueSize),
		pool:          pool,
		batchSize:     batchSize,
		batchInterval: batchInterval,
		workers:       workers,
		byteBudget:    budget,
		done:          make(chan struct{}),
	}
}

// EstimatedBytes covers all variable-size strings and byte slices retained by an entry.
func EstimatedBytes(e *LogEntry) int64 {
	if e == nil {
		return 0
	}
	return int64(len(e.RequestID) + len(e.Method) + len(e.Path) + len(e.Query) + len(e.ClientIP) + len(e.Error) +
		len(e.RequestHeaders) + len(e.RequestBody) + len(e.ResponseHeaders) + len(e.ResponseBody))
}

func (q *Queue) Stats() Stats {
	return Stats{
		Enqueued:  q.enq.Load(),
		Dropped:   q.dropped.Load(),
		Failed:    q.failed.Load(),
		Ambiguous: q.ambiguous.Load(),
		Bytes:     q.bytes.Load(),
	}
}

// Start starts workers once. The root context controls reporting only; Stop owns worker shutdown.
func (q *Queue) Start(root context.Context) {
	q.mu.Lock()
	if q.started || q.stopped {
		q.mu.Unlock()
		return
	}
	q.started = true
	workCtx, workCancel := context.WithCancel(context.Background())
	reportCtx, reporterCancel := context.WithCancel(root)
	q.workCancel = workCancel
	q.reporterCancel = reporterCancel
	for i := 0; i < q.workers; i++ {
		q.wg.Add(1)
		go q.run(workCtx)
	}
	q.wg.Add(1)
	go q.reportLoop(reportCtx)
	q.mu.Unlock()
}

// Stop closes submissions once and waits under the caller's single total deadline.
// The variadic form permits legacy Stop() calls while new callers should pass a context.
func (q *Queue) Stop(contexts ...context.Context) error {
	ctx := context.Background()
	if len(contexts) > 0 && contexts[0] != nil {
		ctx = contexts[0]
	}

	q.mu.Lock()
	if !q.stopped {
		q.stopped = true
		q.shutdownCtx = ctx
		close(q.ch)
		if q.reporterCancel != nil {
			q.reporterCancel()
		}
		if q.started {
			go func() {
				q.wg.Wait()
				q.doneOnce.Do(func() { close(q.done) })
			}()
		} else {
			q.discardQueued()
			q.doneOnce.Do(func() { close(q.done) })
		}
	}
	done := q.done
	workCancel := q.workCancel
	q.mu.Unlock()

	select {
	case <-done:
		return nil
	case <-ctx.Done():
		if workCancel != nil {
			workCancel()
		}
		return ctx.Err()
	}
}

// Submit reserves memory and enqueues without waiting; full, over-budget, and stopped queues drop.
func (q *Queue) Submit(e *LogEntry) {
	if e == nil {
		q.dropped.Add(1)
		return
	}
	size := EstimatedBytes(e)
	if !q.mu.TryRLock() {
		q.dropped.Add(1)
		return
	}
	if q.stopped || !q.reserve(size) {
		q.mu.RUnlock()
		q.dropped.Add(1)
		return
	}
	select {
	case q.ch <- e:
		q.enq.Add(1)
	default:
		q.release(size)
		q.dropped.Add(1)
	}
	q.mu.RUnlock()
}

func (q *Queue) reserve(size int64) bool {
	if q.byteBudget <= 0 {
		q.bytes.Add(size)
		return true
	}
	for {
		used := q.bytes.Load()
		if size > q.byteBudget-used {
			return false
		}
		if q.bytes.CompareAndSwap(used, used+size) {
			return true
		}
	}
}

func (q *Queue) release(size int64) {
	q.bytes.Add(-size)
}

func (q *Queue) discardQueued() {
	for e := range q.ch {
		q.release(EstimatedBytes(e))
		q.failed.Add(1)
	}
}

func (q *Queue) run(workCtx context.Context) {
	defer q.wg.Done()
	batch := make([]*LogEntry, 0, q.batchSize)
	ticker := time.NewTicker(q.batchInterval)
	defer ticker.Stop()

	flush := func() {
		if len(batch) == 0 {
			return
		}
		q.flush(q.flushContext(workCtx), batch)
		for _, e := range batch {
			q.release(EstimatedBytes(e))
		}
		clear(batch)
		batch = batch[:0]
	}

	for {
		select {
		case e, ok := <-q.ch:
			if !ok {
				flush()
				return
			}
			batch = append(batch, e)
			if len(batch) >= q.batchSize {
				flush()
			}
		case <-ticker.C:
			flush()
		case <-workCtx.Done():
			flush()
			q.discardQueued()
			return
		}
	}
}

func (q *Queue) flushContext(workCtx context.Context) context.Context {
	q.mu.RLock()
	defer q.mu.RUnlock()
	if q.shutdownCtx != nil {
		return q.shutdownCtx
	}
	return workCtx
}

func (q *Queue) flush(ctx context.Context, entries []*LogEntry) {
	if q.pool == nil || !q.pool.Healthy() {
		q.failed.Add(uint64(len(entries)))
		return
	}

	retry := entries
	var lastErr error
	for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
		if err := ctx.Err(); err != nil {
			lastErr = err
			break
		}
		conn := q.pool.Get()
		if conn == nil {
			lastErr = errors.New("pool nil")
			break
		}
		result := Flush(ctx, clickHouseBatchPreparer{conn: conn}, retry)
		q.failed.Add(uint64(result.Failed))
		q.ambiguous.Add(uint64(result.Ambiguous))
		if result.Ambiguous > 0 {
			q.pool.MarkUnhealthy()
		}
		lastErr = result.Err
		retry = result.Retry
		if len(retry) == 0 {
			return
		}
		if attempt < maxAttempts && !waitBackoff(ctx, attempt) {
			lastErr = ctx.Err()
			break
		}
	}

	if len(retry) > 0 {
		q.failed.Add(uint64(len(retry)))
		q.pool.MarkUnhealthy()
		log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
	}
}

func waitBackoff(ctx context.Context, attempt int) bool {
	wait := BaseBackoff
	for i := 1; i < attempt; i++ {
		wait *= 3
	}
	jitter := time.Duration((rand.Float64()*0.4 - 0.2) * float64(wait))
	timer := time.NewTimer(wait + jitter)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-ctx.Done():
		return false
	}
}

const insertStatement = `INSERT INTO proxy_logs (
        request_id, method, path, query, client_ip,
        request_headers, request_body, request_truncated,
        status_code, response_headers, response_body, response_truncated,
        is_stream, latency_ms, started_at, finished_at, error
    )`

// Flush executes one batch. Only Prepare failures are retryable. Append failures are isolated by row;
// Send failures are ambiguous and therefore never replayed.
func Flush(ctx context.Context, conn BatchPreparer, entries []*LogEntry) FlushResult {
	if len(entries) == 0 {
		return FlushResult{}
	}
	result := flushOnce(ctx, conn, entries)
	if result.stage != flushAppend || len(entries) == 1 {
		return result.public(entries)
	}

	combined := FlushResult{}
	for _, entry := range entries {
		single := flushOnce(ctx, conn, []*LogEntry{entry}).public([]*LogEntry{entry})
		combined.Retry = append(combined.Retry, single.Retry...)
		combined.Failed += single.Failed
		combined.Ambiguous += single.Ambiguous
		combined.Err = errors.Join(combined.Err, single.Err)
	}
	return combined
}

type flushStage uint8

const (
	flushSuccess flushStage = iota
	flushPrepare
	flushAppend
	flushSend
)

type flushAttempt struct {
	stage flushStage
	err   error
}

func (r flushAttempt) public(entries []*LogEntry) FlushResult {
	switch r.stage {
	case flushSuccess:
		return FlushResult{}
	case flushPrepare:
		return FlushResult{Retry: entries, Err: r.err}
	case flushSend:
		return FlushResult{Failed: len(entries), Ambiguous: len(entries), Err: r.err}
	default:
		return FlushResult{Failed: len(entries), Err: r.err}
	}
}

func flushOnce(ctx context.Context, conn BatchPreparer, entries []*LogEntry) flushAttempt {
	batch, err := conn.PrepareBatch(ctx, insertStatement)
	if err != nil {
		return flushAttempt{stage: flushPrepare, err: err}
	}
	abort := func(cause error) error {
		return errors.Join(cause, batch.Abort())
	}
	for _, e := range entries {
		if err := batch.Append(
			e.RequestID, e.Method, e.Path, e.Query, e.ClientIP,
			string(e.RequestHeaders), string(e.RequestBody), e.RequestTruncated,
			int32(e.StatusCode), string(e.ResponseHeaders), string(e.ResponseBody), e.ResponseTruncated,
			e.IsStream, e.LatencyMS, e.StartedAt, e.FinishedAt, e.Error,
		); err != nil {
			return flushAttempt{stage: flushAppend, err: abort(err)}
		}
	}
	if err := batch.Send(); err != nil {
		return flushAttempt{stage: flushSend, err: abort(err)}
	}
	return flushAttempt{stage: flushSuccess}
}

func (q *Queue) reportLoop(ctx context.Context) {
	defer q.wg.Done()
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			stats := q.Stats()
			healthy := q.pool != nil && q.pool.Healthy()
			log.Printf("[logger] metrics: enq=%d dropped_queue_full=%d dropped_db_fail=%d ambiguous_send=%d queue_len=%d queue_bytes=%d db_healthy=%v",
				stats.Enqueued, stats.Dropped, stats.Failed, stats.Ambiguous, len(q.ch), stats.Bytes, healthy)
		}
	}
}