package pgstore

import (
	"bytes"
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"time"

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

// idempotency.go: the storage half of `Idempotency-Key`.

// ErrKeyReused is the same key presented for a DIFFERENT request. Told apart from a replay because
// the remedies differ: one is "your retry worked", the other is "your client has a bug".
var ErrKeyReused = errors.New("pgstore: this idempotency key was used for another request")

// ErrKeyInFlight is a repeat that arrived while the first attempt is still running. Waiting is the
// remedy, which is why the answer carries Retry-After.
var ErrKeyInFlight = errors.New("pgstore: the first request under this idempotency key is still running")

// IdempotencyKey identifies one claim: (principal, method, path), as the canon ratified — the same
// key on another operation is ANOTHER key, not a conflict.
//
// The path is KEYED by its digest and stored beside it: an address is unbounded, and an over-long
// one overflowed the index tuple, so the claim failed and the caller was answered 500 where the
// contract promises 409.
type IdempotencyKey struct {
	UserID      string
	Method      string
	Path        string
	Key         string
	Fingerprint []byte
}

// pathHash is what the index is keyed by. Derived rather than carried, so it cannot disagree with
// the path beside it.
func (k IdempotencyKey) pathHash() []byte {
	sum := sha256.Sum256([]byte(k.Path))
	return sum[:]
}

// IdempotentResponse is a completed first attempt, replayed verbatim.
type IdempotentResponse struct {
	Status   int
	Location string
	Body     []byte
	// ContentSHA256 is the digest of the body that attempt accepted, for a request whose identity its
	// FINGERPRINT cannot settle: a multipart upload declares its parts and not its bytes, so two
	// different books can arrive under one key with the same declaration. A caller holding one of
	// these may not replay until it has read the repeat's body and found the same digest.
	//
	// Absent where the fingerprint IS the whole request — a JSON write is compared whole.
	ContentSHA256 []byte
}

// ClaimStale is how long a claim may be outstanding before another attempt may take it over; without
// it a dead process locks the key until the retention sweep.
//
// ⚠ Must exceed the longest write here — an upload, bounded by its route's read deadline. Enforced
// at boot (config.Load): under a longer deadline a retry takes the claim from a LIVE attempt and the
// user gets the second book the header exists to prevent.
const ClaimStale = 30 * time.Minute

// ClaimToken says WHICH attempt holds a claim. The key identifies the request; this identifies the
// try, and without it every writer addressed the row by identity alone — so a late attempt released
// its successor's claim, and a superseded one wrote its own receipt over the row.
type ClaimToken string

// NewClaimToken mints one. It is an owner marker rather than a secret: what matters is that two
// attempts never share one.
func NewClaimToken() ClaimToken { return ClaimToken(newID("clm")) }

// errKeyVanished is the row disappearing between an attempt's insert and its re-read. Not an answer
// to anybody: the key is simply free again, and the claim is taken from the top.
var errKeyVanished = errors.New("pgstore: the idempotency key was released mid-claim")

// claimRounds bounds that retry. Two attempts of one key can hand it back and forth at most as fast
// as they can fail, and a caller that loses three rounds is meeting something other than a race.
const claimRounds = 3

// ClaimIdempotency takes the right to perform a request once.
//
// It answers one of four things, and the four are the whole semantics: a token to go ahead with, a
// stored response to replay, ErrKeyInFlight while the first attempt is still running, ErrKeyReused
// when the key was used for a different request. The token is what the caller must present to
// complete or release: it is the proof that the claim is still this attempt's.
func (s *Store) ClaimIdempotency(ctx context.Context, k IdempotencyKey, now time.Time) (*IdempotentResponse, ClaimToken, error) {
	// ⚠ Retried, because losing a race is not one of the four answers. Two attempts arrive together;
	// one wins the insert and then FAILS FAST — a 4xx gives the key straight back — and the loser's
	// re-read, on a fresh READ COMMITTED snapshot, finds nothing at all. Answered rather than retried
	// that was a 500 on precisely the case this header exists for, which is the same 500 the
	// `do nothing` below was written to close. Reproduced on a live database.
	for range claimRounds {
		out, token, err := s.claim(ctx, k, now)
		if errors.Is(err, errKeyVanished) {
			continue
		}
		return out, token, err
	}
	return nil, "", fmt.Errorf("pgstore: the idempotency key was claimed and released under us %d times", claimRounds)
}

func (s *Store) claim(ctx context.Context, k IdempotencyKey, now time.Time) (*IdempotentResponse, ClaimToken, error) {
	var out *IdempotentResponse
	token := NewClaimToken()
	err := s.inTx(ctx, func(tx pgx.Tx) error {
		var fingerprint, content []byte
		var status *int
		var location string
		var body []byte
		var claimedAt time.Time
		var finishedAt *time.Time
		const read = `
			select request_sha256, status, location, body, claimed_at, finished_at, content_sha256
			  from idempotency_keys
			 where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4 for update`
		err := tx.QueryRow(ctx, read, k.UserID, k.Method, k.pathHash(), k.Key).
			Scan(&fingerprint, &status, &location, &body, &claimedAt, &finishedAt, &content)
		if errors.Is(err, pgx.ErrNoRows) {
			// ⚠ `for update` on a missing row locks nothing, so two simultaneous first attempts both
			// land here. `do nothing` plus a re-read decides between them; a bare insert made the loser
			// raise a unique violation — a 500 on the one case this header exists for.
			tag, err := tx.Exec(ctx, `
				insert into idempotency_keys (user_id, method, path, path_sha256, key, request_sha256, claimed_at, claim_token)
				values ($1, $2, $3, $4, $5, $6, $7, $8)
				on conflict (user_id, method, path_sha256, key) do nothing`,
				k.UserID, k.Method, k.Path, k.pathHash(), k.Key, k.Fingerprint, now, token)
			if err != nil {
				return fmt.Errorf("pgstore: claim idempotency key: %w", err)
			}
			if tag.RowsAffected() == 1 {
				return nil
			}
			if err := tx.QueryRow(ctx, read, k.UserID, k.Method, k.pathHash(), k.Key).
				Scan(&fingerprint, &status, &location, &body, &claimedAt, &finishedAt, &content); err != nil {
				if errors.Is(err, pgx.ErrNoRows) {
					return errKeyVanished // the winner released it; the key is free and this attempt starts over
				}
				return fmt.Errorf("pgstore: re-read idempotency key: %w", err)
			}
		} else if err != nil {
			return fmt.Errorf("pgstore: read idempotency key: %w", err)
		}
		if !bytes.Equal(fingerprint, k.Fingerprint) {
			return ErrKeyReused
		}
		if finishedAt != nil && status != nil {
			out = &IdempotentResponse{Status: *status, Location: location, Body: body, ContentSHA256: content}
			return nil
		}
		if now.Sub(claimedAt) < ClaimStale {
			return ErrKeyInFlight
		}
		// The first attempt is gone: its process died, or it answered an error and released nothing.
		// Taking the claim over is what keeps a retry from meeting `key_in_flight` forever, and the NEW
		// token is what stops the old attempt — which may still be alive and merely slow — from
		// completing or releasing what is no longer its claim.
		if _, err := tx.Exec(ctx, `
			update idempotency_keys set claimed_at = $5, request_sha256 = $6, claim_token = $7
			 where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4`,
			k.UserID, k.Method, k.pathHash(), k.Key, now, k.Fingerprint, token); err != nil {
			return fmt.Errorf("pgstore: re-claim idempotency key: %w", err)
		}
		return nil
	})
	if err != nil {
		return nil, "", err
	}
	if out != nil {
		return out, "", nil // a replay holds no claim
	}
	return nil, token, nil
}

// ErrClaimLost is a completion whose row is gone or already finished. The work IS done, but nothing
// will replay it — a fact worth reporting rather than losing in a nil.
var ErrClaimLost = errors.New("pgstore: this attempt no longer holds the idempotency key")

// CompleteIdempotency records what the first attempt answered, so a repeat can be given the same
// thing without doing the work again.
//
// Only the attempt that still HOLDS the claim may write the receipt. An attempt whose claim was
// taken over has had its work done by somebody else; letting it answer would replay a `Location` for
// a book that was never created.
func (s *Store) CompleteIdempotency(ctx context.Context, k IdempotencyKey, token ClaimToken, resp IdempotentResponse, now time.Time) error {
	tag, err := s.pool.Exec(ctx, `
		update idempotency_keys set status = $5, location = $6, body = $7, finished_at = $8,
		                            content_sha256 = $10
		 where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4
		   and claim_token = $9 and finished_at is null`,
		k.UserID, k.Method, k.pathHash(), k.Key, resp.Status, resp.Location, resp.Body, now, token,
		resp.ContentSHA256)
	if err != nil {
		return fmt.Errorf("pgstore: complete idempotency key: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrClaimLost
	}
	return nil
}

// ReleaseIdempotency gives a key back after an attempt that did NOT complete.
//
// A failed attempt is not a completed one, and the contract names the case that matters: a `408` is
// not a completed attempt, so the same key may be presented again and the retry is a repeat of the
// first call rather than a second book.
//
// On the token for the same reason as the completion, and this is the half that cost a duplicate
// book: a slow attempt that finally failed deleted the row of the successor which had taken the key
// over and was already doing the work.
func (s *Store) ReleaseIdempotency(ctx context.Context, k IdempotencyKey, token ClaimToken) error {
	tag, err := s.pool.Exec(ctx, `
		delete from idempotency_keys
		 where user_id = $1 and method = $2 and path_sha256 = $3 and key = $4
		   and claim_token = $5 and finished_at is null`,
		k.UserID, k.Method, k.pathHash(), k.Key, token)
	if err != nil {
		return fmt.Errorf("pgstore: release idempotency key: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrClaimLost
	}
	return nil
}

// SweepIdempotency forgets records past the window the contract promises ("at least 24 hours").
func (s *Store) SweepIdempotency(ctx context.Context, before time.Time) (int64, error) {
	tag, err := s.pool.Exec(ctx, `delete from idempotency_keys where claimed_at < $1`, before)
	if err != nil {
		return 0, fmt.Errorf("pgstore: sweep idempotency keys: %w", err)
	}
	return tag.RowsAffected(), nil
}