// 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) {
	response, err, _ := c.generate(ctx, req, false)
	return response, err
}

// GenerateWithEvidence runs the built-in retry path and returns its factual
// attempt evidence alongside any response retained before cancellation.
func (c *Coordinator) GenerateWithEvidence(
	ctx context.Context,
	req domain.GenerateRequest,
) (*domain.GenerateResponse, error, GenerationEvidence) {
	return c.generate(ctx, req, true)
}

func (c *Coordinator) generate(
	ctx context.Context,
	req domain.GenerateRequest,
	captureEvidence bool,
) (*domain.GenerateResponse, error, GenerationEvidence) {
	if c.injected != nil {
		return c.generateInjected(ctx, req, captureEvidence)
	}
	prepared, err := c.builtIn.PrepareGeneration(req)
	if err != nil {
		return nil, err, GenerationEvidence{}
	}

	evidence := GenerationEvidence{}
	if captureEvidence {
		evidence.InitialEndpoint = prepared.InitialEndpoint()
	}
	var usage observation.UsageAccumulator
	finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
		if captureEvidence {
			evidence.Usage = usage.Observed()
		}
		return response, err, evidence
	}
	var attempts, retriesUsed uint64
	recorder := observation.FromContext(ctx)
	generation := req.Generation
	if generation == 0 {
		generation = 1
	}
	for {
		if err := ctx.Err(); err != nil {
			return finish(nil, err)
		}
		ordinal := attempts + 1
		outcome, dispatched, capacityWait, generationDuration, startedAt := c.execute(ctx, req.Target.BackendID, prepared, captureEvidence)
		if dispatched {
			attempts++
		}
		if captureEvidence {
			attempt := AttemptEvidence{
				StartedAt:           startedAt.UTC(),
				Dispatched:          dispatched,
				ResponseReceived:    outcome.Evidence.ResponseReceived,
				HTTP:                cloneHTTPEvidence(outcome.Evidence),
				CapacityWait:        capacityWait,
				RequestDuration:     generationDuration,
				Failure:             ClassifyFailure(outcome.Err),
				ResponseDecoded:     outcome.ResponseDecoded,
				GenerationCompleted: outcome.GenerationCompleted,
			}
			if outcome.StructuredOutputRequested != nil {
				requested := *outcome.StructuredOutputRequested
				attempt.StructuredOutputRequested = &requested
			}
			if dispatched {
				usage.Add(outcome.Usage, outcome.UsagePresence)
				var attemptUsage observation.UsageAccumulator
				attemptUsage.Add(outcome.Usage, outcome.UsagePresence)
				attempt.Usage = attemptUsage.Observed()
			}
			evidence.Attempts = append(evidence.Attempts, attempt)
		}
		if recorder != nil {
			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
				Dispatched: dispatched, CapacityWait: capacityWait, GenerationDuration: generationDuration,
				Outcome: attemptOutcome(outcome.Err), Failure: ClassifyFailure(outcome.Err),
				StatusCode: outcome.StatusCode, Usage: outcome.Usage, UsagePresence: outcome.UsagePresence})
		}
		if outcome.Err == nil {
			if err := ctx.Err(); err != nil {
				if captureEvidence {
					return finish(outcome.Response, err)
				}
				return finish(nil, err)
			}
			return finish(outcome.Response, nil)
		}
		if !dispatched || !eligibleStatus(outcome) {
			if captureEvidence {
				return finish(outcome.Response, outcome.Err)
			}
			return finish(nil, outcome.Err)
		}
		if retriesUsed == req.TransientRetries {
			if retriesUsed == 0 {
				return finish(nil, outcome.Err)
			}
			return finish(nil, &ExhaustedError{attempts: attempts, cause: outcome.Err})
		}
		if err := ctx.Err(); err != nil {
			return finish(nil, err)
		}
		retryOrdinal := retriesUsed + 1
		delay := retryDelay(c.backoff(retryOrdinal), outcome.RetryNotBefore, c.now())
		waitStart := c.now()
		waitErr := c.wait(ctx, delay)
		waitDuration := elapsed(waitStart, c.now())
		if captureEvidence {
			evidence.RetryDelays = append(evidence.RetryDelays, waitDuration)
		}
		if recorder != nil {
			recorder.Delay(observation.DelayData{Generation: generation, Attempt: attempts,
				Elapsed: waitDuration, Failure: ClassifyFailure(waitErr)})
		}
		if waitErr != nil {
			return finish(nil, waitErr)
		}
		if err := ctx.Err(); err != nil {
			return finish(nil, err)
		}
		retriesUsed++
	}
}

func (c *Coordinator) generateInjected(
	ctx context.Context,
	req domain.GenerateRequest,
	captureEvidence bool,
) (*domain.GenerateResponse, error, GenerationEvidence) {
	evidence := GenerationEvidence{}
	var usage observation.UsageAccumulator
	finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
		if captureEvidence {
			evidence.Usage = usage.Observed()
		}
		return response, err, evidence
	}
	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 captureEvidence {
			evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
				StartedAt:    start.UTC(),
				CapacityWait: capacityWait,
				Failure:      ClassifyFailure(err),
			})
		}
		if recorder != nil {
			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
		}
		return finish(nil, err)
	}
	defer release()
	if err := ctx.Err(); err != nil {
		release()
		if captureEvidence {
			evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
				StartedAt:    start.UTC(),
				CapacityWait: capacityWait,
				Failure:      ClassifyFailure(err),
			})
		}
		if recorder != nil {
			recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
		}
		return finish(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
	}
	if captureEvidence {
		attempt := AttemptEvidence{
			StartedAt:           start.UTC(),
			Dispatched:          true,
			ResponseReceived:    response != nil,
			RequestDuration:     duration,
			Failure:             ClassifyFailure(err),
			GenerationCompleted: err == nil && response != nil,
		}
		requested := req.StructuredOutput != nil
		attempt.StructuredOutputRequested = &requested
		var reportedUsage domain.TokenUsage
		var reportedPresence *domain.TokenUsagePresence
		if err == nil && response != nil {
			reportedUsage = response.Usage
			reportedPresence = response.UsagePresence
		}
		usage.Add(reportedUsage, reportedPresence)
		var attemptUsage observation.UsageAccumulator
		attemptUsage.Add(reportedUsage, reportedPresence)
		attempt.Usage = attemptUsage.Observed()
		evidence.Attempts = append(evidence.Attempts, attempt)
	}
	release()
	if recorder != nil {
		data := observation.AttemptData{Generation: generation, Attempt: 1, Dispatched: true, CapacityWait: capacityWait, GenerationDuration: duration,
			Outcome: attemptOutcome(err), Failure: ClassifyFailure(err)}
		if err == nil {
			data.Usage = response.Usage
			data.UsagePresence = response.UsagePresence
		}
		recorder.Attempt(data)
	}
	if err != nil {
		return finish(nil, err)
	}
	if err := ctx.Err(); err != nil {
		if captureEvidence {
			return finish(response, err)
		}
		return finish(nil, err)
	}
	return finish(response, nil)
}

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

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
}

// ClassifyFailure maps an execution error to the shared content-free category.
func ClassifyFailure(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()
	}
}