package middleware

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
)

type IdempotencyConfig struct {
	TTL              time.Duration `json:"ttl"`
	MaxEntries       int           `json:"maxEntries"`
	MaxBytes         int           `json:"maxBytes"`
	MaxRequestBytes  int64         `json:"maxRequestBytes"`
	MaxResponseBytes int           `json:"maxResponseBytes"`
	MaxInFlight      int           `json:"maxInFlight"`
	// Scope must return a server-verified tenant/user identity. Register auth
	// before this middleware. Client-provided identity is not authorization.
	Scope func(*httpx.Context) (string, error) `json:"-"`
}

type replay struct {
	fingerprint     [32]byte
	status          int
	headers         http.Header
	body            []byte
	expires         time.Time
	running, failed bool
	size            int
}
type Idempotency struct {
	mu     sync.Mutex
	config IdempotencyConfig
	items  map[[32]byte]*replay
	bytes  int
	slots  chan struct{}
	now    func() time.Time
}

func NewIdempotency(cfg IdempotencyConfig) (*Idempotency, error) {
	if cfg.Scope == nil || cfg.TTL < 0 || cfg.MaxEntries < 0 || cfg.MaxBytes < 0 || cfg.MaxRequestBytes < 0 || cfg.MaxResponseBytes < 0 || cfg.MaxInFlight < 0 {
		return nil, errors.New("middleware: invalid idempotency configuration or missing scope")
	}
	if cfg.TTL == 0 {
		cfg.TTL = 10 * time.Minute
	}
	if cfg.MaxEntries == 0 {
		cfg.MaxEntries = 256
	}
	if cfg.MaxBytes == 0 {
		cfg.MaxBytes = 16 << 20
	}
	if cfg.MaxRequestBytes == 0 {
		cfg.MaxRequestBytes = 1 << 20
	}
	if cfg.MaxResponseBytes == 0 {
		cfg.MaxResponseBytes = 64 << 10
	}
	if cfg.MaxInFlight == 0 {
		cfg.MaxInFlight = 32
	}
	if cfg.MaxRequestBytes > 1<<30 {
		return nil, errors.New("middleware: maxRequestBytes exceeds 1 GiB")
	}
	return &Idempotency{config: cfg, items: make(map[[32]byte]*replay), slots: make(chan struct{}, cfg.MaxInFlight), now: time.Now}, nil
}

func (i *Idempotency) Middleware() httpx.Middleware {
	return func(next httpx.Handler) httpx.Handler {
		return func(c *httpx.Context) error {
			r := c.Request()
			if r.Method == "GET" || r.Method == "HEAD" || r.Method == "OPTIONS" {
				return next(c)
			}
			select {
			case i.slots <- struct{}{}:
				defer func() { <-i.slots }()
			default:
				return apierrors.New(503, "IDEMPOTENCY_BUSY", "idempotency capacity reached")
			}
			key := r.Header.Get("Idempotency-Key")
			if len(r.Header.Values("Idempotency-Key")) != 1 || len(key) < 1 || len(key) > 128 || strings.IndexFunc(key, func(r rune) bool { return r < 33 || r > 126 }) >= 0 {
				return apierrors.BadRequest("INVALID_IDEMPOTENCY_KEY", "Idempotency-Key must contain 1..128 printable ASCII characters")
			}
			scope, err := i.config.Scope(c)
			if err != nil {
				return err
			}
			if scope == "" || len(scope) > 1024 {
				return apierrors.Unauthorized("IDEMPOTENCY_SCOPE_REQUIRED", "authenticated scope is required")
			}
			if r.ContentLength > i.config.MaxRequestBytes {
				return apierrors.New(413, "BODY_TOO_LARGE", "request body too large")
			}
			var body []byte
			if r.Body != nil {
				body, err = io.ReadAll(io.LimitReader(r.Body, i.config.MaxRequestBytes+1))
				if err != nil {
					return apierrors.BadRequest("INVALID_BODY", "unable to read request body").WithCause(err)
				}
				if int64(len(body)) > i.config.MaxRequestBytes {
					return apierrors.New(413, "BODY_TOO_LARGE", "request body too large")
				}
			}
			if err := c.Context().Err(); err != nil {
				return err
			}
			originalBody := r.Body
			r.Body = io.NopCloser(bytes.NewReader(body))
			defer func() { r.Body = originalBody }()
			parts, _ := json.Marshal([]string{r.Method, r.URL.RequestURI(), r.Header.Get("Content-Type")})
			hash := sha256.New()
			hash.Write(parts)
			hash.Write(body)
			var fingerprint [32]byte
			copy(fingerprint[:], hash.Sum(nil))
			namespace, _ := json.Marshal([]string{scope, key})
			id := sha256.Sum256(namespace)
			i.mu.Lock()
			now := i.now()
			for k, v := range i.items {
				if !v.running && !now.Before(v.expires) {
					delete(i.items, k)
					i.bytes -= v.size
				}
			}
			if saved := i.items[id]; saved != nil {
				if saved.fingerprint != fingerprint {
					i.mu.Unlock()
					return apierrors.New(409, "IDEMPOTENCY_CONFLICT", "key was used with a different request")
				}
				if saved.running || saved.failed {
					i.mu.Unlock()
					return apierrors.New(409, "IDEMPOTENCY_UNAVAILABLE", "request is in progress or its outcome cannot be replayed")
				}
				status, headers, data := saved.status, saved.headers.Clone(), append([]byte(nil), saved.body...)
				i.mu.Unlock()
				headers.Set("Idempotency-Replayed", "true")
				return c.WriteResponse(status, headers, data)
			}
			if len(i.items) >= i.config.MaxEntries {
				i.mu.Unlock()
				return apierrors.New(503, "IDEMPOTENCY_BUSY", "idempotency capacity reached")
			}
			saved := &replay{fingerprint: fingerprint, running: true}
			i.items[id] = saved
			i.mu.Unlock()
			buffer := &boundedResponse{headers: c.Raw().Header().Clone(), limit: i.config.MaxResponseBytes}
			// Always retain an outcome tombstone on panic/error/overflow. Deleting
			// the key could execute a partially committed business operation twice.
			defer func() {
				i.mu.Lock()
				defer i.mu.Unlock()
				saved.running = false
				saved.failed = true
				saved.expires = i.now().Add(i.config.TTL)
				size := buffer.body.Len()
				for k, values := range buffer.headers {
					size += len(k)
					for _, v := range values {
						size += len(v)
					}
				}
				if buffer.complete && buffer.err == nil && buffer.status < 500 && buffer.headers.Get("Set-Cookie") == "" && !strings.Contains(strings.ToLower(buffer.headers.Get("Cache-Control")), "no-store") && size <= i.config.MaxBytes-i.bytes {
					saved.failed = false
					saved.status = buffer.status
					saved.headers = buffer.headers.Clone()
					saved.body = append([]byte(nil), buffer.body.Bytes()...)
					saved.size = size
					i.bytes += size
				}
			}()
			child := c.Clone(c.Context(), buffer)
			err = next(child)
			// Keep an ambiguous outcome tombstone, but let an outer Timeout
			// middleware turn deadline expiry into 504 before committing bytes.
			if ctxErr := child.Context().Err(); ctxErr != nil {
				return ctxErr
			}
			if err != nil {
				child.WriteError(err)
			}
			if buffer.err != nil {
				return apierrors.Internal(buffer.err)
			}
			if buffer.status == 0 {
				buffer.status = 200
			}
			buffer.complete = true
			return c.WriteResponse(buffer.status, buffer.headers, buffer.body.Bytes())
		}
	}
}

// Intentionally lacks Flush/Hijack: only buffered, bounded endpoints are valid.
type boundedResponse struct {
	headers       http.Header
	status, limit int
	body          bytes.Buffer
	err           error
	complete      bool
}

func (b *boundedResponse) Header() http.Header { return b.headers }
func (b *boundedResponse) WriteHeader(status int) {
	if status < 200 {
		b.err = errors.New("idempotency: informational/upgrade responses are unsupported")
		return
	}
	if b.status == 0 {
		b.status = status
	}
}
func (b *boundedResponse) Write(data []byte) (int, error) {
	if b.status == 0 {
		b.status = 200
	}
	if b.err != nil {
		return 0, b.err
	}
	if len(data) > b.limit-b.body.Len() {
		b.err = fmt.Errorf("idempotency: response exceeds %d bytes", b.limit)
		return 0, b.err
	}
	return b.body.Write(data)
}