package usecase

import (
	"context"
	"fmt"
	"sync"
	"time"

	"go.maximumdirect.net/promptkit/internal/domain"
	"go.maximumdirect.net/promptkit/internal/observation"
	"go.maximumdirect.net/promptkit/internal/validate"
)

type preparedExecutionState uint8

const (
	preparedExecutionReady preparedExecutionState = iota
	preparedExecutionClaimed
	preparedExecutionDiscarded
)

// PreparedExecution owns one frozen, single-use runner execution.
type PreparedExecution struct {
	owner   *runnerIdentity
	mu      sync.Mutex
	state   preparedExecutionState
	details *domain.PreparedRun
	payload *preparedExecutionPayload
}

type preparedExecutionPayload struct {
	prepared         *domain.PreparedRun
	validation       validate.PreparedValidation
	directKey        string
	transientRetries uint64
	observer         domain.ExecutionObserver
}

// PrepareExecution completes preparation without generation or admission and
// returns a runner-bound, single-use execution.
func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*PreparedExecution, error) {
	state, err := r.resolvePreparation(ctx, req, time.Now().UTC())
	if err != nil {
		return nil, err
	}

	operation, err := r.completePreparation(ctx, req, state)
	if err != nil {
		return nil, err
	}

	return &PreparedExecution{
		owner:   r.identity,
		state:   preparedExecutionReady,
		details: operation.run,
		payload: &preparedExecutionPayload{
			prepared:         operation.run,
			validation:       operation.validation,
			directKey:        state.effectiveModel.APIKey,
			transientRetries: state.transientRetries,
			observer:         req.Observer,
		},
	}, nil
}

// Details returns immutable, credential-redacted metadata borrowed from the
// handle. Callers must copy it before exposing it to mutable external code.
func (p *PreparedExecution) Details() *domain.PreparedRun {
	if p == nil {
		return nil
	}
	p.mu.Lock()
	detailsSnapshot := p.details
	p.mu.Unlock()

	return detailsSnapshot
}

// Discard invalidates an unclaimed execution and drops its private payload.
func (p *PreparedExecution) Discard() {
	if p == nil {
		return
	}

	p.mu.Lock()
	if p.state != preparedExecutionReady {
		p.mu.Unlock()
		return
	}
	p.state = preparedExecutionDiscarded
	payload := p.payload
	p.payload = nil
	p.mu.Unlock()

	payload.clear()
}

// RunPrepared claims and executes one prepared execution owned by this runner.
func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (*domain.RunResult, error) {
	payload, err := prepared.claim(r)
	if err != nil {
		return nil, err
	}
	defer payload.clear()

	var recorder *observation.Recorder
	var runID string
	if payload.observer != nil {
		var idErr error
		recorder, idErr = observation.New(payload.observer)
		observation.RegisterClaim(ctx, recorder)
		ctx = observation.WithRecorder(ctx, recorder)
		selected := payload.prepared
		recorder.UpdateSelection(observation.Selection{PromptID: selected.PromptID,
			ProfileID: selected.SelectedProfileID,
			BackendID: selected.SelectedBackendID, Model: selected.EffectiveModelParams.Model})
		recorder.Start()
		if idErr != nil {
			return nil, fmt.Errorf("failed to create run id: %w", idErr)
		}
		runID = recorder.RunID()
	} else {
		var err error
		runID, err = observation.NewRunID()
		if err != nil {
			return nil, fmt.Errorf("failed to create run id: %w", err)
		}
	}
	start := time.Now().UTC()

	target := payload.prepared.EffectiveModelParams
	if err := validateAPIKey(target.APIKeyEnv, payload.directKey, target.APIKeyRequired); err != nil {
		return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
	}

	release, err := r.admitRun(ctx, target.BackendID)
	if err != nil {
		return nil, err
	}
	defer release()

	return r.executePreparedRun(ctx, payload.prepared, payload.directKey, payload.transientRetries, runID, start, func(
		ctx context.Context,
		artifact *domain.Artifact,
		attemptsUsed int,
	) (domain.ValidationResult, error) {
		result, validationErr := payload.validation.Validate(ctx, artifact)
		if validationErr != nil {
			return domain.ValidationResult{}, validationErr
		}
		result.RepairAttempts = attemptsUsed
		return result, nil
	})
}

func (p *PreparedExecution) claim(owner *Runner) (*preparedExecutionPayload, error) {
	if p == nil || owner == nil || owner.identity == nil || p.owner != owner.identity {
		return nil, fmt.Errorf("%w: prepared execution does not belong to this runner", ErrInvalidRequest)
	}

	p.mu.Lock()
	defer p.mu.Unlock()
	if p.state != preparedExecutionReady || p.payload == nil {
		return nil, fmt.Errorf("%w: prepared execution is not ready", ErrInvalidRequest)
	}
	p.state = preparedExecutionClaimed
	payload := p.payload
	p.payload = nil
	return payload, nil
}

func (p *preparedExecutionPayload) clear() {
	if p == nil {
		return
	}
	p.prepared = nil
	p.validation = nil
	p.directKey = ""
	p.transientRetries = 0
	p.observer = nil
}

type noOpPreparedValidation struct {
	contract domain.OutputContract
}

func (p noOpPreparedValidation) Validate(
	ctx context.Context,
	_ *domain.Artifact,
) (domain.ValidationResult, error) {
	if err := ctx.Err(); err != nil {
		return domain.ValidationResult{}, err
	}
	return domain.ValidationResult{
		Status:         domain.ValidationSkipped,
		Mode:           p.contract.ValidationMode,
		SchemaPath:     p.contract.SchemaPath,
		RepairAttempts: p.contract.RepairAttempts,
		IsValid:        true,
	}, nil
}

func (noOpPreparedValidation) SchemaDocument() any {
	return nil
}