// Package generation coordinates model attempts, backend permits, and bounded
// built-in HTTP retries for one generation request.
package generation

import (
	"context"
	"errors"
	"fmt"
	"math/rand/v2"
	"time"

	"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/llm"
	"gitea.maximumdirect.net/eric/promptkit/internal/observation"
)

const (
	initialBackoff = 500 * time.Millisecond
	maximumBackoff = 8 * time.Second
)

// Coordinator implements the internal model client using the existing backend
// capacity manager. Exactly one of builtIn and injected is selected.
type Coordinator struct {
	manager  *capacity.Manager
	builtIn  *llm.OpenAICompatibleClient
	injected llm.Client
	now      func() time.Time
	jitter   func(time.Duration) time.Duration
	wait     func(context.Context, time.Duration) error
}

// NewBuiltIn selects the built-in HTTP retry path.
func NewBuiltIn(manager *capacity.Manager, client *llm.OpenAICompatibleClient) *Coordinator {
	return &Coordinator{manager: manager, builtIn: client, now: time.Now, jitter: randomJitter, wait: waitContext}
}

// NewInjected selects one opaque client dispatch with no automatic retry.
func NewInjected(manager *capacity.Manager, client llm.Client) *Coordinator {
	return &Coordinator{manager: manager, injected: client, now: time.Now, jitter: randomJitter, wait: waitContext}
}

// ExhaustedError retains the final eligible status failure and the number of
// attempts actually dispatched.
type ExhaustedError struct {
	attempts uint64
	cause    error
}

func (e *ExhaustedError) Error() string {
	if e == nil {
		return "transient model retries exhausted"
	}
	return fmt.Sprintf("transient model retries exhausted after %d attempts", e.attempts)
}

func (e *ExhaustedError) Unwrap() error {
	if e == nil {
		return nil
	}
	return e.cause
}

func (e *ExhaustedError) Attempts() uint64 {
	if e == nil {
		return 0
	}
	return e.attempts
}

// Generate executes either one opaque client call or a frozen built-in HTTP
// sequence with one permit per actual dispatch.
func (c *Coordinator) Generate(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
	if c.injected != nil {
		return c.generateInjected(ctx, req)
	}
	prepared, err := c.builtIn.PrepareGeneration(req)
	if err != nil {
		return nil, err
	}

	var attempts, retriesUsed uint64
	recorder := observation.FromContext(ctx)
	generation := req.Generation
	if generation == 0 {
		generation = 1
	}
	for {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		ordinal := attempts + 1
		outcome, dispatched, capacityWait, generationDuration := c.execute(ctx, req.Target.BackendID, prepared)
		if dispatched {
			attempts++
		}
		if recorder != nil {
			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
				Dispatched: dispatched, CapacityWait: capacityWait, GenerationDuration: generationDuration,
				Outcome: attemptOutcome(outcome.Err), Failure: attemptFailure(outcome.Err),
				StatusCode: outcome.StatusCode, Usage: outcome.Usage, UsagePresence: outcome.UsagePresence})
		}
		if outcome.Err == nil {
			if err := ctx.Err(); err != nil {
				return nil, err
			}
			return outcome.Response, nil
		}
		if !dispatched || !eligibleStatus(outcome) {
			return nil, outcome.Err
		}
		if retriesUsed == req.TransientRetries {
			if retriesUsed == 0 {
				return nil, outcome.Err
			}
			return nil, &ExhaustedError{attempts: attempts, cause: outcome.Err}
		}
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		retryOrdinal := retriesUsed + 1
		delay := retryDelay(c.backoff(retryOrdinal), outcome.RetryNotBefore, c.now())
		waitStart := c.now()
		waitErr := c.wait(ctx, delay)
		if recorder != nil {
			recorder.Delay(observation.DelayData{Generation: generation, Attempt: attempts,
				Elapsed: elapsed(waitStart, c.now()), Failure: attemptFailure(waitErr)})
		}
		if waitErr != nil {
			return nil, waitErr
		}
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		retriesUsed++
	}
}

func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
	recorder := observation.FromContext(ctx)
	generation := req.Generation
	if generation == 0 {
		generation = 1
	}
	start := c.now()
	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
	capacityWait := c.capacityWait(req.Target.BackendID, start)
	if err != nil {
		if recorder != nil {
			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: attemptFailure(err)})
		}
		return nil, err
	}
	defer release()
	if err := ctx.Err(); err != nil {
		release()
		if recorder != nil {
			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: attemptFailure(err)})
		}
		return nil, err
	}
	dispatchStart := c.now()
	response, err := c.injected.Generate(ctx, req)
	duration := elapsed(dispatchStart, c.now())
	if err == nil && response == nil {
		err = llm.ErrMalformedResponse
	}
	release()
	if recorder != nil {
		data := observation.AttemptData{Generation: generation, Attempt: 1, Dispatched: true, CapacityWait: capacityWait, GenerationDuration: duration,
			Outcome: attemptOutcome(err), Failure: attemptFailure(err)}
		if err == nil {
			data.Usage = response.Usage
			data.UsagePresence = response.UsagePresence
		}
		recorder.Attempt(data)
	}
	if err != nil {
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	return response, nil
}

func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration) (outcome llm.AttemptOutcome, dispatched bool, capacityWait, generationDuration time.Duration) {
	start := c.now()
	release, err := c.manager.Acquire(ctx, backendID)
	capacityWait = c.capacityWait(backendID, start)
	if err != nil {
		return llm.AttemptOutcome{Err: err}, false, capacityWait, 0
	}
	defer release()
	if err := ctx.Err(); err != nil {
		return llm.AttemptOutcome{Err: err}, false, capacityWait, 0
	}
	dispatchStart := c.now()
	outcome = prepared.Execute(ctx)
	return outcome, true, capacityWait, elapsed(dispatchStart, c.now())
}

func (c *Coordinator) capacityWait(backendID string, start time.Time) time.Duration {
	if !c.manager.Limited(backendID) {
		return 0
	}
	return elapsed(start, c.now())
}

func elapsed(start, end time.Time) time.Duration {
	if d := end.Sub(start); d > 0 {
		return d
	}
	return 0
}

func attemptOutcome(err error) domain.ExecutionOutcome {
	if err == nil {
		return domain.ExecutionOutcomeSuccess
	}
	return domain.ExecutionOutcomeError
}

func attemptFailure(err error) domain.ExecutionFailureKind {
	if err == nil {
		return domain.ExecutionFailureNone
	}
	if errors.Is(err, context.Canceled) {
		return domain.ExecutionFailureCanceled
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return domain.ExecutionFailureDeadline
	}
	if errors.Is(err, llm.ErrUnexpectedStatus) {
		return domain.ExecutionFailureProviderStatus
	}
	if errors.Is(err, llm.ErrTransport) {
		return domain.ExecutionFailureTransport
	}
	if errors.Is(err, llm.ErrMalformedResponse) {
		return domain.ExecutionFailureMalformedResponse
	}
	if errors.Is(err, llm.ErrInvalidRequest) {
		return domain.ExecutionFailureInvalidRequest
	}
	return domain.ExecutionFailureGeneration
}

func eligibleStatus(outcome llm.AttemptOutcome) bool {
	if !errors.Is(outcome.Err, llm.ErrUnexpectedStatus) {
		return false
	}
	switch outcome.StatusCode {
	case 429, 500, 502, 503, 504:
		return true
	default:
		return false
	}
}

func (c *Coordinator) backoff(ordinal uint64) time.Duration {
	base := initialBackoff
	for step := uint64(1); step < ordinal && base < maximumBackoff; step++ {
		base *= 2
	}
	half := base / 2
	return half + c.jitter(half)
}

func retryDelay(local time.Duration, notBefore, now time.Time) time.Duration {
	if providerDelay := notBefore.Sub(now); providerDelay > local {
		return providerDelay
	}
	return local
}

func randomJitter(max time.Duration) time.Duration {
	return time.Duration(rand.Int64N(int64(max) + 1))
}

func waitContext(ctx context.Context, delay time.Duration) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return ctx.Err()
	}
}