package store

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"strings"
	"time"

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

type CapacityProfile struct {
	ProfileKey                  string         `json:"profileKey"`
	ReleaseSHA                  string         `json:"releaseSha"`
	ConfigHash                  string         `json:"configHash"`
	AcceptanceRunID             string         `json:"acceptanceRunId"`
	Kind                        string         `json:"kind"`
	Model                       string         `json:"model"`
	CapacityProfile             string         `json:"capacityProfile"`
	StableThroughputPerSecond   float64        `json:"stableThroughputPerSecond"`
	AdmittedThroughputPerSecond float64        `json:"admittedThroughputPerSecond"`
	MaxRunning                  int            `json:"maxRunning"`
	MaxQueued                   int            `json:"maxQueued"`
	MaxPredictedWaitSeconds     int            `json:"maxPredictedWaitSeconds"`
	Enabled                     bool           `json:"enabled"`
	Metadata                    map[string]any `json:"metadata"`
	CreatedAt                   time.Time      `json:"createdAt"`
	UpdatedAt                   time.Time      `json:"updatedAt"`
}

type CapacityProfileInput struct {
	ProfileKey                  string         `json:"profileKey,omitempty"`
	Kind                        string         `json:"kind"`
	Model                       string         `json:"model,omitempty"`
	StableThroughputPerSecond   float64        `json:"stableThroughputPerSecond"`
	AdmittedThroughputPerSecond float64        `json:"admittedThroughputPerSecond"`
	MaxRunning                  int            `json:"maxRunning"`
	MaxQueued                   int            `json:"maxQueued"`
	MaxPredictedWaitSeconds     int            `json:"maxPredictedWaitSeconds"`
	Metadata                    map[string]any `json:"metadata,omitempty"`
}

type StageAcceptanceCapacityProfilesInput struct {
	RunID            string                 `json:"-"`
	ConfigHash       string                 `json:"configHash"`
	CapacityProfiles []CapacityProfileInput `json:"capacityProfiles"`
}

func (s *Store) ListCapacityProfiles(ctx context.Context) ([]CapacityProfile, error) {
	rows, err := s.pool.Query(ctx, `
SELECT profile_key, release_sha, config_hash, COALESCE(acceptance_run_id::text, ''),
       kind, model, capacity_profile,
       stable_throughput_per_second::float8,
       admitted_throughput_per_second::float8,
       max_running, max_queued, max_predicted_wait_seconds, enabled,
       metadata, created_at, updated_at
FROM gateway_capacity_profiles
ORDER BY enabled DESC, kind, model, updated_at DESC`)
	if err != nil {
		if IsUndefinedDatabaseObject(err) {
			return []CapacityProfile{}, nil
		}
		return nil, err
	}
	defer rows.Close()
	items := make([]CapacityProfile, 0)
	for rows.Next() {
		item, scanErr := scanCapacityProfile(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func replaceCapacityProfilesTx(
	ctx context.Context,
	tx pgx.Tx,
	run AcceptanceRun,
	configHash string,
	inputs []CapacityProfileInput,
) error {
	configHash = strings.ToLower(strings.TrimSpace(configHash))
	if !capacityConfigHashPattern.MatchString(configHash) {
		return errors.New("configHash must be a lowercase SHA-256 value")
	}
	if err := validateCapacityProfileInputs(inputs); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `
UPDATE gateway_capacity_profiles
SET enabled = false, updated_at = now()
WHERE enabled = true`); err != nil {
		return err
	}
	for _, input := range inputs {
		kind := strings.ToLower(strings.TrimSpace(input.Kind))
		model := strings.TrimSpace(input.Model)
		if model == "" {
			model = "*"
		}
		profileKey := strings.TrimSpace(input.ProfileKey)
		if profileKey == "" {
			profileKey = fmt.Sprintf("%s:%s:%s:%s", run.ReleaseSHA[:12], strings.ToLower(run.CapacityProfile), kind, model)
		}
		metadata, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(input.Metadata)))
		if _, err := tx.Exec(ctx, `
INSERT INTO gateway_capacity_profiles (
  profile_key, release_sha, config_hash, acceptance_run_id, kind, model,
  capacity_profile, stable_throughput_per_second,
  admitted_throughput_per_second, max_running, max_queued,
  max_predicted_wait_seconds, enabled, metadata
)
VALUES ($1, $2, $3, $4::uuid, $5, $6, $7, $8, $9, $10, $11, $12, true, $13::jsonb)
ON CONFLICT (profile_key) DO UPDATE
SET release_sha = EXCLUDED.release_sha,
    config_hash = EXCLUDED.config_hash,
    acceptance_run_id = EXCLUDED.acceptance_run_id,
    kind = EXCLUDED.kind,
    model = EXCLUDED.model,
    capacity_profile = EXCLUDED.capacity_profile,
    stable_throughput_per_second = EXCLUDED.stable_throughput_per_second,
    admitted_throughput_per_second = EXCLUDED.admitted_throughput_per_second,
    max_running = EXCLUDED.max_running,
    max_queued = EXCLUDED.max_queued,
    max_predicted_wait_seconds = EXCLUDED.max_predicted_wait_seconds,
    enabled = true,
    metadata = EXCLUDED.metadata,
    updated_at = now()`,
			profileKey, run.ReleaseSHA, configHash, run.ID, kind, model,
			run.CapacityProfile, input.StableThroughputPerSecond,
			input.AdmittedThroughputPerSecond, input.MaxRunning, input.MaxQueued,
			input.MaxPredictedWaitSeconds, metadata,
		); err != nil {
			return err
		}
	}
	return nil
}

func validateCapacityProfileInputs(inputs []CapacityProfileInput) error {
	if len(inputs) == 0 {
		return errors.New("at least one certified capacity profile is required")
	}
	hasImage := false
	hasVideo := false
	keys := make(map[string]struct{}, len(inputs))
	for _, input := range inputs {
		kind := strings.ToLower(strings.TrimSpace(input.Kind))
		switch kind {
		case "image":
			hasImage = true
		case "video":
			hasVideo = true
		case "mixed":
			hasImage = true
			hasVideo = true
		default:
			return fmt.Errorf("unsupported capacity profile kind %q", input.Kind)
		}
		if input.StableThroughputPerSecond <= 0 ||
			input.AdmittedThroughputPerSecond <= 0 ||
			input.AdmittedThroughputPerSecond > input.StableThroughputPerSecond ||
			input.MaxRunning <= 0 ||
			input.MaxQueued < 0 ||
			input.MaxPredictedWaitSeconds <= 0 {
			return fmt.Errorf("invalid capacity limits for %s/%s", kind, input.Model)
		}
		if input.AdmittedThroughputPerSecond > input.StableThroughputPerSecond*0.800001 {
			return fmt.Errorf("admitted throughput for %s/%s exceeds the 80%% safety envelope", kind, input.Model)
		}
		model := strings.TrimSpace(input.Model)
		if model == "" {
			model = "*"
		}
		key := kind + "\x00" + model
		if _, exists := keys[key]; exists {
			return fmt.Errorf("duplicate capacity profile for %s/%s", kind, model)
		}
		keys[key] = struct{}{}
	}
	if !hasImage || !hasVideo {
		return errors.New("certified capacity profiles must cover both image and video traffic")
	}
	return nil
}

func (s *Store) StageAcceptanceCapacityProfiles(
	ctx context.Context,
	input StageAcceptanceCapacityProfilesInput,
) (AcceptanceRun, error) {
	input.RunID = strings.TrimSpace(input.RunID)
	input.ConfigHash = strings.ToLower(strings.TrimSpace(input.ConfigHash))
	if !capacityConfigHashPattern.MatchString(input.ConfigHash) {
		return AcceptanceRun{}, errors.New("configHash must be a lowercase SHA-256 value")
	}
	if err := validateCapacityProfileInputs(input.CapacityProfiles); err != nil {
		return AcceptanceRun{}, err
	}
	rawProfiles, _ := json.Marshal(input.CapacityProfiles)
	var profileValues any
	_ = json.Unmarshal(rawProfiles, &profileValues)
	profiles, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(profileValues)))
	profilesHash := fmt.Sprintf("%x", sha256.Sum256(profiles))
	return scanAcceptanceRun(s.pool.QueryRow(ctx, `
UPDATE gateway_acceptance_runs
SET config = jsonb_set(
      jsonb_set(
        jsonb_set(config, '{validationCapacityProfiles}', $2::jsonb, true),
        '{validationConfigHash}', to_jsonb($3::text), true
      ),
      '{validationProfilesHash}', to_jsonb($4::text), true
    ),
    updated_at = now()
WHERE id = $1::uuid
  AND status = 'running'
RETURNING `+acceptanceRunColumns,
		input.RunID, profiles, input.ConfigHash, profilesHash,
	))
}

func capacityProfileInputsHash(inputs []CapacityProfileInput) string {
	raw, _ := json.Marshal(inputs)
	var value any
	_ = json.Unmarshal(raw, &value)
	canonical, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(value)))
	return fmt.Sprintf("%x", sha256.Sum256(canonical))
}

func enforceTaskCapacityTx(ctx context.Context, tx pgx.Tx, input CreateTaskInput, runMode string) error {
	if runMode != "production" && runMode != "acceptance" {
		return nil
	}
	kind := capacityKindForTask(input.Kind)
	if kind == "" {
		return nil
	}
	model := strings.TrimSpace(input.Model)
	profile, applies, err := capacityProfileForTaskTx(ctx, tx, input, runMode, kind, model)
	if err != nil {
		return err
	}
	if !applies {
		return nil
	}
	var admissionLock bool
	if err := tx.QueryRow(ctx, `
SELECT pg_try_advisory_xact_lock(
  hashtextextended('capacity-profile:' || $1, 0)
)`, profile.ProfileKey).Scan(&admissionLock); err != nil {
		return err
	}
	if !admissionLock {
		return capacityExceededError(profile, "capacity_admission_busy", "admission_lock", 1, 1, 0, 0)
	}

	var queued int
	var running int
	var recent int
	if err := tx.QueryRow(ctx, `
SELECT
  count(*) FILTER (WHERE status = 'queued')::int,
  count(*) FILTER (WHERE status = 'running')::int,
  count(*) FILTER (WHERE created_at >= now() - interval '60 seconds')::int
FROM gateway_tasks
WHERE (
    ($3 = 'production' AND run_mode = 'production')
    OR (
      $3 = 'acceptance'
      AND acceptance_run_id = NULLIF($4, '')::uuid
      AND run_mode = 'acceptance'
    )
  )
  AND (
    ($1 = 'image' AND kind IN ('images.generations', 'images.edits', 'images.vectorize'))
    OR ($1 = 'video' AND kind IN ('videos.generations', 'videos.upscales'))
    OR ($1 = 'mixed' AND kind IN (
      'images.generations', 'images.edits', 'images.vectorize',
      'videos.generations', 'videos.upscales'
    ))
  )
  AND ($2 = '*' OR model = $2)
  AND (status IN ('queued', 'running') OR created_at >= now() - interval '60 seconds')`,
		profile.Kind, profile.Model, runMode, strings.TrimSpace(input.AcceptanceRunID),
	).Scan(&queued, &running, &recent); err != nil {
		return err
	}

	nextQueued := queued + 1
	nextRecentRate := float64(recent+1) / 60
	predictedWait := float64(nextQueued) / profile.AdmittedThroughputPerSecond
	reason := ""
	metric := ""
	current := float64(nextQueued)
	limit := float64(profile.MaxQueued)
	switch {
	case nextQueued > profile.MaxQueued:
		reason, metric = "capacity_queue_limit", "queue_depth"
	case predictedWait > float64(profile.MaxPredictedWaitSeconds):
		reason, metric = "capacity_predicted_wait_limit", "predicted_wait_seconds"
		current, limit = predictedWait, float64(profile.MaxPredictedWaitSeconds)
	case nextRecentRate > profile.AdmittedThroughputPerSecond:
		reason, metric = "capacity_throughput_limit", "throughput_per_second"
		current, limit = nextRecentRate, profile.AdmittedThroughputPerSecond
	}
	if reason == "" {
		return nil
	}
	return capacityExceededError(profile, reason, metric, current, limit, queued, running)
}

func capacityProfileForTaskTx(
	ctx context.Context,
	tx pgx.Tx,
	input CreateTaskInput,
	runMode string,
	kind string,
	model string,
) (CapacityProfile, bool, error) {
	if runMode == "acceptance" {
		var profilesJSON []byte
		var releaseSHA string
		var capacityProfile string
		var configHash string
		err := tx.QueryRow(ctx, `
SELECT COALESCE(config->'validationCapacityProfiles', 'null'::jsonb),
       release_sha,
       capacity_profile,
       COALESCE(config->>'validationConfigHash', '')
FROM gateway_acceptance_runs
WHERE id = NULLIF($1, '')::uuid
  AND status = 'running'`, strings.TrimSpace(input.AcceptanceRunID)).Scan(
			&profilesJSON, &releaseSHA, &capacityProfile, &configHash,
		)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				return CapacityProfile{}, false, nil
			}
			return CapacityProfile{}, false, err
		}
		if string(profilesJSON) == "null" {
			return CapacityProfile{}, false, nil
		}
		var inputs []CapacityProfileInput
		if err := json.Unmarshal(profilesJSON, &inputs); err != nil {
			return CapacityProfile{}, false, fmt.Errorf("decode staged capacity profiles: %w", err)
		}
		best := -1
		bestRank := 100
		for index, candidate := range inputs {
			candidateKind := strings.ToLower(strings.TrimSpace(candidate.Kind))
			candidateModel := strings.TrimSpace(candidate.Model)
			if candidateModel == "" {
				candidateModel = "*"
			}
			rank := 100
			switch {
			case candidateKind == kind && candidateModel == model:
				rank = 0
			case candidateKind == kind && candidateModel == "*":
				rank = 1
			case candidateKind == "mixed" && candidateModel == "*":
				rank = 2
			}
			if rank < bestRank {
				best = index
				bestRank = rank
			}
		}
		if best < 0 {
			return CapacityProfile{}, true, capacityProfileMissingError(kind, model)
		}
		selected := inputs[best]
		selectedModel := strings.TrimSpace(selected.Model)
		if selectedModel == "" {
			selectedModel = "*"
		}
		return CapacityProfile{
			ProfileKey:                  "acceptance:" + strings.TrimSpace(input.AcceptanceRunID) + ":" + kind + ":" + selectedModel,
			ReleaseSHA:                  releaseSHA,
			ConfigHash:                  configHash,
			AcceptanceRunID:             strings.TrimSpace(input.AcceptanceRunID),
			Kind:                        strings.ToLower(strings.TrimSpace(selected.Kind)),
			Model:                       selectedModel,
			CapacityProfile:             capacityProfile,
			StableThroughputPerSecond:   selected.StableThroughputPerSecond,
			AdmittedThroughputPerSecond: selected.AdmittedThroughputPerSecond,
			MaxRunning:                  selected.MaxRunning,
			MaxQueued:                   selected.MaxQueued,
			MaxPredictedWaitSeconds:     selected.MaxPredictedWaitSeconds,
			Metadata:                    selected.Metadata,
		}, true, nil
	}

	profile, err := scanCapacityProfile(tx.QueryRow(ctx, `
SELECT profile_key, release_sha, config_hash, COALESCE(acceptance_run_id::text, ''),
       kind, model, capacity_profile,
       stable_throughput_per_second::float8,
       admitted_throughput_per_second::float8,
       max_running, max_queued, max_predicted_wait_seconds, enabled,
       metadata, created_at, updated_at
FROM gateway_capacity_profiles
WHERE enabled = true
  AND (
    (kind = $1 AND model IN ($2, '*'))
    OR (kind = 'mixed' AND model = '*')
  )
ORDER BY CASE
  WHEN kind = $1 AND model = $2 THEN 0
  WHEN kind = $1 AND model = '*' THEN 1
  ELSE 2
END
LIMIT 1`, kind, model))
	if err != nil {
		if IsUndefinedDatabaseObject(err) {
			return CapacityProfile{}, false, nil
		}
		if errors.Is(err, pgx.ErrNoRows) {
			var enabledProfiles int
			if countErr := tx.QueryRow(ctx, `
SELECT count(*)::int
FROM gateway_capacity_profiles
WHERE enabled = true
  AND kind IN ($1, 'mixed')`, kind).Scan(&enabledProfiles); countErr != nil {
				return CapacityProfile{}, false, countErr
			}
			if enabledProfiles == 0 {
				return CapacityProfile{}, false, nil
			}
			return CapacityProfile{}, true, capacityProfileMissingError(kind, model)
		}
		return CapacityProfile{}, false, err
	}
	return profile, true, nil
}

func capacityProfileMissingError(kind string, model string) error {
	return &RateLimitExceededError{
		ScopeType: "capacity_profile",
		ScopeKey:  kind + ":" + model,
		ScopeName: "uncertified_model",
		ScopeMetadata: map[string]any{
			"kind": kind, "model": model,
		},
		Metric:     "certified_capacity",
		Limit:      0,
		Amount:     1,
		Current:    0,
		Projected:  1,
		Message:    "model has no certified production capacity profile",
		RetryAfter: 60 * time.Second,
		Retryable:  true,
		Reason:     "capacity_profile_missing",
	}
}

func capacityExceededError(
	profile CapacityProfile,
	reason string,
	metric string,
	current float64,
	limit float64,
	queued int,
	running int,
) error {
	predictedWait := float64(queued) / profile.AdmittedThroughputPerSecond
	retrySeconds := int(math.Ceil(predictedWait))
	if retrySeconds < 1 {
		retrySeconds = 1
	}
	return &RateLimitExceededError{
		ScopeType: "capacity_profile",
		ScopeKey:  profile.ProfileKey,
		ScopeName: profile.CapacityProfile,
		ScopeMetadata: map[string]any{
			"kind":                 profile.Kind,
			"model":                profile.Model,
			"releaseSha":           profile.ReleaseSHA,
			"configHash":           profile.ConfigHash,
			"maxRunning":           profile.MaxRunning,
			"maxQueued":            profile.MaxQueued,
			"running":              running,
			"predictedWaitSeconds": predictedWait,
		},
		Metric:     metric,
		Limit:      limit,
		Amount:     1,
		Current:    current,
		Projected:  current + 1,
		Message:    "certified production capacity has been reached",
		RetryAfter: time.Duration(retrySeconds) * time.Second,
		Retryable:  true,
		Reason:     reason,
		QueueDepth: queued,
		QueueLimit: profile.MaxQueued,
	}
}

func capacityKindForTask(kind string) string {
	switch strings.TrimSpace(kind) {
	case "images.generations", "images.edits", "images.vectorize":
		return "image"
	case "videos.generations", "videos.upscales":
		return "video"
	default:
		return ""
	}
}

func scanCapacityProfile(scanner taskScanner) (CapacityProfile, error) {
	var item CapacityProfile
	var metadata []byte
	if err := scanner.Scan(
		&item.ProfileKey, &item.ReleaseSHA, &item.ConfigHash, &item.AcceptanceRunID,
		&item.Kind, &item.Model, &item.CapacityProfile,
		&item.StableThroughputPerSecond, &item.AdmittedThroughputPerSecond,
		&item.MaxRunning, &item.MaxQueued, &item.MaxPredictedWaitSeconds,
		&item.Enabled, &metadata, &item.CreatedAt, &item.UpdatedAt,
	); err != nil {
		return CapacityProfile{}, err
	}
	_ = json.Unmarshal(metadata, &item.Metadata)
	if item.Metadata == nil {
		item.Metadata = map[string]any{}
	}
	return item, nil
}