package logger

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

	"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)
}

// Backend is the database contract required by the asynchronous queue.
type Backend interface {
	BatchPreparer
	Healthy() bool
	MarkUnhealthy()
}

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 poolBackend struct {
	pool       *db.Pool
	conn       driver.Conn
	generation uint64
}

func (b poolBackend) PrepareBatch(ctx context.Context, query string) (Batch, error) {
	conn := b.conn
	if conn == nil {
		conn, _ = b.pool.GetWithGeneration()
	}
	if conn == nil {
		return nil, errors.New("pool nil")
	}
	return clickHouseBatchPreparer{conn: conn}.PrepareBatch(ctx, query)
}

func (b poolBackend) Healthy() bool  { return b.pool.Healthy() }
func (b poolBackend) MarkUnhealthy() { b.pool.MarkUnhealthyGeneration(b.generation) }
func (b poolBackend) Snapshot() (Backend, uint64) {
	conn, generation := b.pool.GetWithGeneration()
	return poolBackend{pool: b.pool, conn: conn, generation: generation}, generation
}
func (b poolBackend) MarkUnhealthyGeneration(generation uint64) {
	b.pool.MarkUnhealthyGeneration(generation)
}

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
}

type queuedEntry struct {
	entry *LogEntry
	size  int64
}

// Queue is an asynchronous, bounded logger. Submit never waits for database work.
type Queue struct {
	ch            chan queuedEntry
	backend       Backend
	batchSize     int
	batchInterval time.Duration
	workers       int
	entryBudget   int64
	byteBudget    int64

	dropped   atomic.Uint64
	failed    atomic.Uint64
	ambiguous atomic.Uint64
	enq       atomic.Uint64
	entries   atomic.Int64
	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 backend Backend
	if pool != nil {
		backend = poolBackend{pool: pool}
	}
	return NewQueueWithBackend(backend, queueSize, batchSize, workers, batchInterval, byteBudget...)
}

// NewQueueWithBackend builds a queue against the database operations used by workers.
func NewQueueWithBackend(backend Backend, 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 queuedEntry, queueSize),
		backend:       backend,
		batchSize:     batchSize,
		batchInterval: batchInterval,
		workers:       workers,
		entryBudget:   int64(queueSize),
		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(unsafe.Sizeof(*e)+unsafe.Sizeof(queuedEntry{})) + 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()
	shutdownOwner := false
	if !q.stopped {
		shutdownOwner = true
		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 shutdownOwner && workCancel != nil {
			workCancel()
			timer := time.NewTimer(100 * time.Millisecond)
			defer timer.Stop()
			select {
			case <-done:
			case <-timer.C:
			}
		}
		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
	}
	entry := cloneLogEntry(e)
	select {
	case q.ch <- queuedEntry{entry: entry, size: size}:
		q.enq.Add(1)
	default:
		q.release(size)
		q.dropped.Add(1)
	}
	q.mu.RUnlock()
}

func cloneLogEntry(entry *LogEntry) *LogEntry {
	clone := *entry
	clone.RequestHeaders = append([]byte(nil), entry.RequestHeaders...)
	clone.RequestBody = append([]byte(nil), entry.RequestBody...)
	clone.ResponseHeaders = append([]byte(nil), entry.ResponseHeaders...)
	clone.ResponseBody = append([]byte(nil), entry.ResponseBody...)
	return &clone
}

func (q *Queue) reserve(size int64) bool {
	for {
		used := q.entries.Load()
		if used >= q.entryBudget {
			return false
		}
		if q.entries.CompareAndSwap(used, used+1) {
			break
		}
	}

	if q.byteBudget <= 0 {
		q.bytes.Add(size)
		return true
	}
	for {
		used := q.bytes.Load()
		if size > q.byteBudget-used {
			q.entries.Add(-1)
			return false
		}
		if q.bytes.CompareAndSwap(used, used+size) {
			return true
		}
	}
}

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

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

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

	flush := func() {
		if len(batch) == 0 {
			return
		}
		entries := make([]*LogEntry, len(batch))
		for i, item := range batch {
			entries[i] = item.entry
		}
		q.flush(q.flushContext(workCtx), entries)
		for _, item := range batch {
			q.release(item.size)
		}
		clear(batch)
		batch = batch[:0]
	}

	for {
		select {
		case item, ok := <-q.ch:
			if !ok {
				flush()
				return
			}
			batch = append(batch, item)
			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.backend == nil || !q.backend.Healthy() {
		q.failed.Add(uint64(len(entries)))
		return
	}

	retry := entries
	var lastErr error
	var failedGeneration uint64
	for attempt := 1; attempt <= maxAttempts && len(retry) > 0; attempt++ {
		if err := ctx.Err(); err != nil {
			lastErr = err
			break
		}
		attemptBackend, generation := backendSnapshot(q.backend)
		failedGeneration = generation
		result := Flush(ctx, attemptBackend, retry)
		q.failed.Add(uint64(result.Failed))
		q.ambiguous.Add(uint64(result.Ambiguous))
		if result.Ambiguous > 0 {
			markBackendUnhealthy(q.backend, failedGeneration)
		}
		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)))
		markBackendUnhealthy(q.backend, failedGeneration)
		log.Printf("[logger] giving up %d retry-safe rows after %d attempts: %v", len(retry), maxAttempts, lastErr)
	}
}

type generationBackend interface {
	Snapshot() (Backend, uint64)
	MarkUnhealthyGeneration(uint64)
}

func backendSnapshot(backend Backend) (Backend, uint64) {
	if versioned, ok := backend.(generationBackend); ok {
		return versioned.Snapshot()
	}
	return backend, 0
}

func markBackendUnhealthy(backend Backend, generation uint64) {
	if versioned, ok := backend.(generationBackend); ok {
		versioned.MarkUnhealthyGeneration(generation)
		return
	}
	backend.MarkUnhealthy()
}

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 i, 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)
		if single.Ambiguous > 0 {
			if backend, ok := conn.(interface{ MarkUnhealthy() }); ok {
				backend.MarkUnhealthy()
			}
			combined.Failed += len(combined.Retry) + len(entries) - i - 1
			combined.Retry = nil
			break
		}
	}
	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.backend != nil && q.backend.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)
		}
	}
}