package postgres

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"telesrv/internal/domain"
	"telesrv/internal/store"
	"telesrv/internal/store/postgres/sqlcgen"
)

const (
	plainPrivateSendBatchWorkers        = 8
	plainPrivateSendBatchMaxTasks       = 32
	plainPrivateSendBatchMinTasks       = 8
	plainPrivateSendBatchMailbox        = 8192
	plainPrivateSendBatchFlushInterval  = 8 * time.Millisecond
	plainPrivateSendBatchTimeout        = 10 * time.Second
	plainPrivateSendBatchMaxQueuedBytes = int64(64 << 20)
	plainPrivateSendTaskFixedBytes      = int64(1024)
)

var errPlainPrivateSendBatchOverloaded = errors.New("postgres: plain private send batch actor overloaded")

var processPlainPrivateSendBatcher = newPlainPrivateSendBatchActor()

type plainPrivateSendBatchTask struct {
	ctx         context.Context
	store       *MessageStore
	req         domain.SendPrivateTextRequest
	fingerprint []byte
	bytes       int64
	done        chan plainPrivateSendBatchResult
}

type plainPrivateSendBatchResult struct {
	result domain.SendPrivateTextResult
	err    error
}

type plainPrivateSendBatchCompletion struct {
	tasks   []*plainPrivateSendBatchTask
	results []plainPrivateSendBatchResult
}

type plainPrivateSendScope struct {
	store  *MessageStore
	userID int64
}

type plainPrivateSendBatchActor struct {
	submit      chan *plainPrivateSendBatchTask
	work        chan []*plainPrivateSendBatchTask
	completed   chan plainPrivateSendBatchCompletion
	queuedBytes atomic.Int64
	batches     atomic.Uint64
	tasks       atomic.Uint64
}

type plainPrivateSendBatchSnapshot struct {
	Batches uint64
	Tasks   uint64
}

func newPlainPrivateSendBatchActor() *plainPrivateSendBatchActor {
	a := &plainPrivateSendBatchActor{
		submit:    make(chan *plainPrivateSendBatchTask, plainPrivateSendBatchMailbox),
		work:      make(chan []*plainPrivateSendBatchTask, plainPrivateSendBatchWorkers),
		completed: make(chan plainPrivateSendBatchCompletion, plainPrivateSendBatchWorkers),
	}
	for range plainPrivateSendBatchWorkers {
		go a.worker()
	}
	go a.run()
	return a
}

func (a *plainPrivateSendBatchActor) Snapshot() plainPrivateSendBatchSnapshot {
	if a == nil {
		return plainPrivateSendBatchSnapshot{}
	}
	return plainPrivateSendBatchSnapshot{Batches: a.batches.Load(), Tasks: a.tasks.Load()}
}

func (a *plainPrivateSendBatchActor) Eligible(messageStore *MessageStore) bool {
	if a == nil || messageStore == nil || messageStore.boxIDs == nil {
		return false
	}
	if _, ok := messageStore.db.(*pgxpool.Pool); !ok {
		return false
	}
	// The process batch path is a distributed allocator + PostgreSQL boundary.
	// Local and test allocators cannot reserve gap-safe ids before the
	// transaction under cross-process concurrency and therefore remain on the
	// single-command transaction path.
	_, distributed := messageStore.boxIDs.(store.DistributedBoxIDAllocator)
	return distributed
}

func (a *plainPrivateSendBatchActor) Submit(
	ctx context.Context,
	messageStore *MessageStore,
	req domain.SendPrivateTextRequest,
	fingerprint []byte,
) (domain.SendPrivateTextResult, error) {
	if !a.Eligible(messageStore) || ctx == nil || len(fingerprint) == 0 {
		return domain.SendPrivateTextResult{}, fmt.Errorf("postgres: invalid plain private send batch submission")
	}
	retained := plainPrivateSendTaskFixedBytes + int64(len(req.Message)+len(fingerprint))
	if !reservePlainPrivateSendBatchBytes(&a.queuedBytes, retained) {
		return domain.SendPrivateTextResult{}, errPlainPrivateSendBatchOverloaded
	}
	task := &plainPrivateSendBatchTask{
		ctx: ctx, store: messageStore, req: req,
		fingerprint: append([]byte(nil), fingerprint...),
		bytes:       retained,
		done:        make(chan plainPrivateSendBatchResult, 1),
	}
	select {
	case a.submit <- task:
	case <-ctx.Done():
		a.queuedBytes.Add(-retained)
		return domain.SendPrivateTextResult{}, ctx.Err()
	}
	select {
	case result := <-task.done:
		return result.result, result.err
	case <-ctx.Done():
		// Once accepted, the bounded actor may commit after the caller stops
		// waiting. random_id is the durable receipt for an exact replay.
		return domain.SendPrivateTextResult{}, ctx.Err()
	}
}

func reservePlainPrivateSendBatchBytes(used *atomic.Int64, amount int64) bool {
	if used == nil || amount <= 0 || amount > plainPrivateSendBatchMaxQueuedBytes {
		return false
	}
	for {
		current := used.Load()
		if current > plainPrivateSendBatchMaxQueuedBytes-amount {
			return false
		}
		if used.CompareAndSwap(current, current+amount) {
			return true
		}
	}
}

func (a *plainPrivateSendBatchActor) run() {
	ticker := time.NewTicker(plainPrivateSendBatchFlushInterval)
	defer ticker.Stop()
	pending := make([]*plainPrivateSendBatchTask, 0, plainPrivateSendBatchMailbox)
	busy := make(map[plainPrivateSendScope]struct{})
	available := plainPrivateSendBatchWorkers

	completeCanceled := func(task *plainPrivateSendBatchTask) {
		err := context.Canceled
		if task.ctx != nil && task.ctx.Err() != nil {
			err = task.ctx.Err()
		}
		task.done <- plainPrivateSendBatchResult{err: err}
		a.queuedBytes.Add(-task.bytes)
	}
	dispatch := func(flush bool) {
		for available > 0 && len(pending) > 0 {
			if !flush && len(pending) < plainPrivateSendBatchMinTasks {
				return
			}
			batch, remaining, canceled := selectPlainPrivateSendBatch(pending, busy)
			pending = remaining
			for _, task := range canceled {
				completeCanceled(task)
			}
			if len(batch) == 0 {
				return
			}
			for _, task := range batch {
				for _, userID := range plainPrivateSendTaskUsers(task) {
					busy[plainPrivateSendScope{store: task.store, userID: userID}] = struct{}{}
				}
			}
			available--
			a.work <- batch
		}
	}

	for {
		select {
		case task := <-a.submit:
			pending = append(pending, task)
			dispatch(false)
		case completion := <-a.completed:
			available++
			for i, task := range completion.tasks {
				for _, userID := range plainPrivateSendTaskUsers(task) {
					delete(busy, plainPrivateSendScope{store: task.store, userID: userID})
				}
				result := plainPrivateSendBatchResult{err: errors.New("postgres: missing plain private send batch result")}
				if i < len(completion.results) {
					result = completion.results[i]
				}
				task.done <- result
				a.queuedBytes.Add(-task.bytes)
			}
			dispatch(false)
		case <-ticker.C:
			dispatch(true)
		}
	}
}

func selectPlainPrivateSendBatch(
	pending []*plainPrivateSendBatchTask,
	busy map[plainPrivateSendScope]struct{},
) (batch, remaining, canceled []*plainPrivateSendBatchTask) {
	selected := make(map[plainPrivateSendScope]struct{}, plainPrivateSendBatchMaxTasks*2)
	blockedByOlder := make(map[plainPrivateSendScope]struct{}, plainPrivateSendBatchMaxTasks*2)
	var selectedStore *MessageStore
	remaining = make([]*plainPrivateSendBatchTask, 0, len(pending))
	for _, task := range pending {
		if task == nil || task.ctx == nil || task.ctx.Err() != nil {
			if task != nil {
				canceled = append(canceled, task)
			}
			continue
		}
		users := plainPrivateSendTaskUsers(task)
		blocked := selectedStore != nil && task.store != selectedStore
		for _, userID := range users {
			key := plainPrivateSendScope{store: task.store, userID: userID}
			if _, ok := busy[key]; ok {
				blocked = true
			}
			if _, ok := blockedByOlder[key]; ok {
				blocked = true
			}
			if _, ok := selected[key]; ok {
				blocked = true
			}
		}
		if blocked || len(batch) >= plainPrivateSendBatchMaxTasks {
			remaining = append(remaining, task)
			for _, userID := range users {
				blockedByOlder[plainPrivateSendScope{store: task.store, userID: userID}] = struct{}{}
			}
			continue
		}
		if selectedStore == nil {
			selectedStore = task.store
		}
		batch = append(batch, task)
		for _, userID := range users {
			selected[plainPrivateSendScope{store: task.store, userID: userID}] = struct{}{}
		}
	}
	return batch, remaining, canceled
}

func plainPrivateSendTaskUsers(task *plainPrivateSendBatchTask) []int64 {
	if task == nil {
		return nil
	}
	return normalizedUserLaneIDs([]int64{task.req.SenderUserID, task.req.RecipientUserID})
}

func (a *plainPrivateSendBatchActor) worker() {
	for batch := range a.work {
		results := executePlainPrivateSendBatch(batch)
		a.batches.Add(1)
		a.tasks.Add(uint64(len(batch)))
		a.completed <- plainPrivateSendBatchCompletion{tasks: batch, results: results}
	}
}

func executePlainPrivateSendBatch(tasks []*plainPrivateSendBatchTask) []plainPrivateSendBatchResult {
	results := make([]plainPrivateSendBatchResult, len(tasks))
	if len(tasks) == 0 || tasks[0] == nil || tasks[0].store == nil {
		return plainPrivateSendBatchErrorResults(results, errors.New("postgres: empty plain private send batch"))
	}
	messageStore := tasks[0].store
	pool, ok := messageStore.db.(*pgxpool.Pool)
	if !ok {
		return plainPrivateSendBatchErrorResults(results, errors.New("postgres: plain private send batch requires pgx pool"))
	}
	allocationUsers := make([]int64, 0, len(tasks)*2)
	lockUsers := make([]int64, 0, len(tasks)*2)
	seen := make(map[int64]struct{}, len(tasks)*2)
	for _, task := range tasks {
		if task == nil || task.store != messageStore {
			return plainPrivateSendBatchErrorResults(results, errors.New("postgres: mixed message stores in plain private send batch"))
		}
		for _, userID := range plainPrivateSendTaskUsers(task) {
			if _, exists := seen[userID]; exists {
				return plainPrivateSendBatchErrorResults(results, fmt.Errorf("postgres: overlapping user %d in plain private send batch", userID))
			}
			seen[userID] = struct{}{}
		}
		lockUsers = append(lockUsers, task.req.SenderUserID, task.req.RecipientUserID)
		allocationUsers = append(allocationUsers, task.req.SenderUserID)
		if task.req.SenderUserID != task.req.RecipientUserID && !task.req.RecipientBlocked {
			allocationUsers = append(allocationUsers, task.req.RecipientUserID)
		}
	}
	lockUsers = normalizedUserLaneIDs(lockUsers)
	allocationUsers = normalizedUserLaneIDs(allocationUsers)
	parent := context.Background()
	if tasks[0].ctx != nil {
		parent = context.WithoutCancel(tasks[0].ctx)
	}
	ctx, cancel := context.WithTimeout(parent, plainPrivateSendBatchTimeout)
	defer cancel()

	boxIDs, err := messageStore.boxIDs.NextBoxIDs(ctx, allocationUsers)
	if err != nil {
		return plainPrivateSendBatchErrorResults(results, fmt.Errorf("allocate plain private send batch box ids: %w", err))
	}
	releaseLanes, err := messageStore.privateSendLanes.acquire(ctx, lockUsers...)
	if err != nil {
		return plainPrivateSendBatchErrorResults(results, fmt.Errorf("admit plain private send batch: %w", err))
	}
	defer releaseLanes()

	tx, err := pool.Begin(ctx)
	if err != nil {
		return plainPrivateSendBatchErrorResults(results, fmt.Errorf("begin plain private send batch: %w", err))
	}
	committed := false
	defer func() {
		if !committed {
			_ = tx.Rollback(context.Background())
		}
	}()
	if err := lockUsersForUpdate(ctx, tx, lockUsers...); err != nil {
		return plainPrivateSendBatchErrorResults(results, fmt.Errorf("lock plain private send batch users: %w", err))
	}
	if err := lockDispatchOutboxAppendFences(ctx, tx, allocationUsers); err != nil {
		return plainPrivateSendBatchErrorResults(results, err)
	}
	created, err := createPlainPrivateMessageBatch(ctx, tx, tasks)
	if err != nil {
		return plainPrivateSendBatchErrorResults(results, err)
	}
	qtx := sqlcgen.New(tx)
	for i, task := range tasks {
		if created[i].inserted {
			continue
		}
		duplicate, found, duplicateErr := messageStore.duplicateSendResult(ctx, qtx, task.req, task.fingerprint)
		if duplicateErr != nil {
			return plainPrivateSendBatchErrorResults(results, duplicateErr)
		}
		if !found {
			return plainPrivateSendBatchErrorResults(results, errors.New("duplicate batched private message disappeared after unique conflict"))
		}
		duplicate.Duplicate = true
		results[i].result = duplicate
	}
	projections, err := persistPlainPrivateSendProjectionBatch(ctx, tx, tasks, created, boxIDs)
	if err != nil {
		return plainPrivateSendBatchErrorResults(results, err)
	}
	for i := range tasks {
		if !created[i].inserted {
			continue
		}
		projection, ok := projections[i]
		if !ok {
			return plainPrivateSendBatchErrorResults(results, fmt.Errorf("missing plain private send batch projection %d", i))
		}
		results[i].result = domain.SendPrivateTextResult{
			SenderMessage: projection.Sender, RecipientMessage: projection.Recipient,
			SenderEvent: eventFromMessage(projection.Sender), RecipientEvent: eventFromMessage(projection.Recipient),
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return plainPrivateSendBatchErrorResults(results, fmt.Errorf("commit plain private send batch: %w", err))
	}
	committed = true
	return results
}

func lockDispatchOutboxAppendFences(ctx context.Context, tx pgx.Tx, userIDs []int64) error {
	userIDs = normalizedUserLaneIDs(userIDs)
	if len(userIDs) == 0 {
		return nil
	}
	if _, err := tx.Exec(ctx, `
SELECT pg_advisory_xact_lock_shared(dispatch_outbox_lane_advisory_key(streams.target_user_id))
FROM unnest($1::bigint[]) AS streams(target_user_id)
ORDER BY streams.target_user_id`, userIDs); err != nil {
		return fmt.Errorf("lock dispatch outbox append fences: %w", err)
	}
	return nil
}

func plainPrivateSendBatchErrorResults(results []plainPrivateSendBatchResult, err error) []plainPrivateSendBatchResult {
	for i := range results {
		results[i] = plainPrivateSendBatchResult{err: err}
	}
	return results
}