package postgres

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
)

const jobColumns = `id, owner_id, external_client_id, capability, provider, req_key, status, prompt,
input_asset_ids, input_urls, output_asset_ids, provider_task_id, provider_dispatch_started_at, request_payload, response_payload,
error, retry_of, idempotency_key, idempotency_fingerprint, priority, attempts, max_attempts,
scheduled_at, locked_at, locked_by, started_at, completed_at, webhook_url, webhook_attempts,
webhook_last_status, usage_context, billing, dispatch_ready_at, finalized_at, created_at, updated_at`

const ClaimJobsSQL = `SELECT ` + jobColumns + ` FROM public.claim_generation_jobs($1::text, $2::integer, $3::integer)`
const WriteJobBillingSQL = `UPDATE public.generation_jobs SET billing = $2::jsonb, updated_at = now() WHERE id = $1::text RETURNING id`
const ActivateJobCreationSQL = `UPDATE public.generation_jobs
SET billing = $2::jsonb, dispatch_ready_at = COALESCE(dispatch_ready_at, now()), updated_at = now()
WHERE id = $1::text AND (dispatch_ready_at IS NULL OR billing = $2::jsonb) RETURNING id`
const WriteJobOutputAssetIDsSQL = `UPDATE public.generation_jobs SET output_asset_ids = $2::text[], updated_at = now() WHERE id = $1::text RETURNING id`
const WriteJobBillingFencedSQL = `UPDATE public.generation_jobs SET billing = $2::jsonb, updated_at = now() WHERE id = $1::text AND status = $3::text AND locked_by = $4::text RETURNING id`
const WriteJobOutputAssetIDsFencedSQL = `UPDATE public.generation_jobs SET output_asset_ids = $2::text[], updated_at = now() WHERE id = $1::text AND status = $3::text AND locked_by = $4::text RETURNING id`
const FailJobCreationSQL = `UPDATE public.generation_jobs
SET status = $2::text, error = $3::jsonb, billing = $4::jsonb,
    completed_at = COALESCE(completed_at, now()), locked_at = NULL, locked_by = NULL, updated_at = now()
WHERE id = $1::text RETURNING id`

func (db *Database) ListJobs(ctx context.Context, filter jobs.ListFilter) ([]jobs.Job, error) {
	if err := db.requirePostgres("list generation jobs"); err != nil {
		return nil, err
	}
	clauses := []string{"owner_id = $1::text"}
	args := []any{filter.Scope.OwnerID}
	if filter.Scope.ExternalClientID != "" {
		args = append(args, filter.Scope.ExternalClientID)
		clauses = append(clauses, fmt.Sprintf("external_client_id = $%d::text", len(args)))
	}
	if filter.Status != "" {
		args = append(args, string(filter.Status))
		clauses = append(clauses, fmt.Sprintf("status = $%d::text", len(args)))
	}
	if filter.Capability != "" {
		args = append(args, filter.Capability)
		clauses = append(clauses, fmt.Sprintf("capability = $%d::text", len(args)))
	}
	if filter.Before != nil {
		args = append(args, *filter.Before)
		clauses = append(clauses, fmt.Sprintf("created_at < $%d::timestamptz", len(args)))
	}
	args = append(args, filter.Limit)
	query := `SELECT ` + jobColumns + ` FROM public.generation_jobs WHERE ` + strings.Join(clauses, " AND ") +
		fmt.Sprintf(" ORDER BY created_at DESC LIMIT $%d::integer", len(args))
	return db.queryJobs(ctx, query, args...)
}

func (db *Database) FindJob(ctx context.Context, scope jobs.Scope, id string) (jobs.Job, bool, error) {
	if err := db.requirePostgres("find generation job"); err != nil {
		return jobs.Job{}, false, err
	}
	query := `SELECT ` + jobColumns + ` FROM public.generation_jobs WHERE id = $1::text AND owner_id = $2::text`
	args := []any{id, scope.OwnerID}
	if scope.ExternalClientID != "" {
		query += ` AND external_client_id = $3::text`
		args = append(args, scope.ExternalClientID)
	}
	query += ` LIMIT 1`
	return db.queryOneJob(ctx, query, args...)
}

func (db *Database) FindIdempotentJob(ctx context.Context, scope jobs.Scope, key string) (jobs.Job, bool, error) {
	if err := db.requirePostgres("find idempotent generation job"); err != nil {
		return jobs.Job{}, false, err
	}
	query := `SELECT ` + jobColumns + ` FROM public.generation_jobs
WHERE owner_id = $1::text AND external_client_id = $2::text AND idempotency_key = $3::text LIMIT 1`
	return db.queryOneJob(ctx, query, scope.OwnerID, scope.ExternalClientID, key)
}

func (db *Database) CreateJob(ctx context.Context, job jobs.Job) (jobs.Job, error) {
	if err := db.requirePostgres("create generation job"); err != nil {
		return jobs.Job{}, err
	}
	const columns = `id, owner_id, external_client_id, capability, provider, req_key, status, prompt,
input_asset_ids, input_urls, output_asset_ids, provider_task_id, provider_dispatch_started_at, request_payload, response_payload,
error, retry_of, idempotency_key, idempotency_fingerprint, priority, attempts, max_attempts,
scheduled_at, locked_at, locked_by, started_at, completed_at, webhook_url, webhook_attempts,
webhook_last_status, usage_context, billing, dispatch_ready_at, finalized_at, created_at, updated_at`
	placeholders := make([]string, 36)
	for index := range placeholders {
		placeholders[index] = fmt.Sprintf("$%d", index+1)
	}
	query := `INSERT INTO public.generation_jobs (` + columns + `) VALUES (` + strings.Join(placeholders, ", ") + `) RETURNING ` + jobColumns
	row, found, err := db.queryOneJob(ctx, query, jobArguments(job)...)
	if err != nil {
		if sqlState(err) == "23505" && job.ExternalClientID != "" && job.IdempotencyKey != "" {
			return jobs.Job{}, jobs.ErrUniqueIdempotency
		}
		return jobs.Job{}, fmt.Errorf("create generation job: %w", err)
	}
	if !found {
		return jobs.Job{}, fmt.Errorf("create generation job returned no row")
	}
	return row, nil
}

func (db *Database) UpdateJob(ctx context.Context, id string, patch jobs.Patch) (jobs.Job, error) {
	if err := db.requirePostgres("update generation job"); err != nil {
		return jobs.Job{}, err
	}
	sets := make([]string, 0, 12)
	args := []any{id}
	add := func(column, cast string, value any) {
		args = append(args, value)
		sets = append(sets, fmt.Sprintf("%s = $%d%s", column, len(args), cast))
	}
	if patch.Status != nil {
		add("status", "::text", string(*patch.Status))
	}
	if patch.Error != nil {
		raw, err := json.Marshal(patch.Error)
		if err != nil {
			return jobs.Job{}, fmt.Errorf("encode generation job error: %w", err)
		}
		add("error", "::jsonb", raw)
	}
	if patch.ClearError {
		sets = append(sets, "error = NULL")
	}
	if patch.Attempts != nil {
		add("attempts", "::integer", *patch.Attempts)
	}
	if patch.ScheduledAt != nil {
		add("scheduled_at", "::timestamptz", *patch.ScheduledAt)
	}
	if patch.CompletedAt != nil {
		add("completed_at", "::timestamptz", *patch.CompletedAt)
	}
	if patch.DispatchReadyAt != nil {
		add("dispatch_ready_at", "::timestamptz", *patch.DispatchReadyAt)
	}
	if patch.FinalizedAt != nil {
		add("finalized_at", "::timestamptz", *patch.FinalizedAt)
	}
	if patch.ProviderTaskID != nil {
		add("provider_task_id", "::text", *patch.ProviderTaskID)
	}
	if patch.ProviderDispatchStartedAt != nil {
		add("provider_dispatch_started_at", "::timestamptz", *patch.ProviderDispatchStartedAt)
	}
	if patch.SetResponsePayload {
		add("response_payload", "::jsonb", optionalJSON(patch.ResponsePayload))
	}
	if patch.ClearProviderTaskID {
		sets = append(sets, "provider_task_id = NULL")
	}
	if patch.ClearProviderDispatch {
		sets = append(sets, "provider_dispatch_started_at = NULL")
	}
	if patch.ClearLease {
		sets = append(sets, "locked_at = NULL", "locked_by = NULL")
	}
	if patch.WebhookAttempts != nil {
		add("webhook_attempts", "::integer", *patch.WebhookAttempts)
	}
	if patch.SetWebhookStatus {
		add("webhook_last_status", "::jsonb", patch.WebhookLastStatus)
	}
	if len(sets) == 0 {
		return db.mustFindJobByID(ctx, id)
	}
	sets = append(sets, "updated_at = now()")
	where := []string{"id = $1::text"}
	if len(patch.ExpectedStatuses) != 0 {
		values := make([]string, len(patch.ExpectedStatuses))
		for index, status := range patch.ExpectedStatuses {
			values[index] = string(status)
		}
		args = append(args, values)
		where = append(where, fmt.Sprintf("status = ANY($%d::text[])", len(args)))
	}
	if patch.ExpectedLockedBy != nil {
		args = append(args, *patch.ExpectedLockedBy)
		where = append(where, fmt.Sprintf("locked_by = $%d::text", len(args)))
	}
	query := `UPDATE public.generation_jobs SET ` + strings.Join(sets, ", ") + ` WHERE ` + strings.Join(where, " AND ") + ` RETURNING ` + jobColumns
	job, found, err := db.queryOneJob(ctx, query, args...)
	if err != nil {
		return jobs.Job{}, fmt.Errorf("update generation job: %w", err)
	}
	if !found {
		if len(patch.ExpectedStatuses) != 0 || patch.ExpectedLockedBy != nil {
			return jobs.Job{}, jobs.ErrTransitionConflict
		}
		return jobs.Job{}, fmt.Errorf("generation job not found: %s", id)
	}
	return job, nil
}

func (db *Database) DeleteJob(ctx context.Context, id string) error {
	if err := db.requirePostgres("delete generation job"); err != nil {
		return err
	}
	rows, err := db.querier.Query(ctx, `DELETE FROM public.generation_jobs WHERE id = $1::text RETURNING id`, id)
	if err != nil {
		return fmt.Errorf("delete generation job: %w", err)
	}
	defer rows.Close()
	if !rows.Next() {
		return fmt.Errorf("generation job not found: %s", id)
	}
	var deleted string
	if err := rows.Scan(&deleted); err != nil {
		return fmt.Errorf("scan deleted generation job: %w", err)
	}
	return nil
}

// WriteBilling and WriteOutputAssetIDs implement the deliberately narrow
// orchestration state seam without widening jobs.Patch into a generic update
// bag. Both fail closed when the target disappeared during processing.
func (db *Database) WriteBilling(ctx context.Context, id string, value json.RawMessage) error {
	if len(value) == 0 || !json.Valid(value) {
		return fmt.Errorf("write generation billing: invalid JSON")
	}
	return db.writeJobState(ctx, "write generation billing", WriteJobBillingSQL, id, value)
}

func (db *Database) ActivateCreation(ctx context.Context, id string, value json.RawMessage) error {
	if len(value) == 0 || !json.Valid(value) {
		return fmt.Errorf("activate generation creation: invalid JSON")
	}
	return db.writeJobState(ctx, "activate generation creation", ActivateJobCreationSQL, id, value)
}

func (db *Database) WriteOutputAssetIDs(ctx context.Context, id string, values []string) error {
	if values == nil {
		values = []string{}
	}
	return db.writeJobState(ctx, "write generation output assets", WriteJobOutputAssetIDsSQL, id, values)
}

func (db *Database) WriteBillingFenced(ctx context.Context, id string, value json.RawMessage, status jobs.Status, lockedBy string) error {
	if len(value) == 0 || !json.Valid(value) || lockedBy == "" {
		return fmt.Errorf("write generation billing: invalid fenced state")
	}
	return db.writeJobStateConflict(ctx, "write generation billing", WriteJobBillingFencedSQL, id, value, string(status), lockedBy)
}

func (db *Database) WriteOutputAssetIDsFenced(ctx context.Context, id string, values []string, status jobs.Status, lockedBy string) error {
	if lockedBy == "" {
		return fmt.Errorf("write generation output assets: invalid fenced state")
	}
	if values == nil {
		values = []string{}
	}
	return db.writeJobStateConflict(ctx, "write generation output assets", WriteJobOutputAssetIDsFencedSQL, id, values, string(status), lockedBy)
}

// FailCreation commits the externally observable result of a failed charge in
// one statement. A job must never remain queued with a wallet-side failure.
func (db *Database) FailCreation(ctx context.Context, job jobs.Job) error {
	if job.Status != jobs.StatusFailed || job.Error == nil || len(job.Billing) == 0 || !json.Valid(job.Billing) {
		return fmt.Errorf("fail generation creation: invalid terminal state")
	}
	errorJSON, err := json.Marshal(job.Error)
	if err != nil {
		return fmt.Errorf("fail generation creation: encode error: %w", err)
	}
	return db.writeJobStateArgs(ctx, "fail generation creation", FailJobCreationSQL, job.ID, string(job.Status), json.RawMessage(errorJSON), job.Billing)
}

func (db *Database) writeJobState(ctx context.Context, operation, query, id string, value any) error {
	return db.writeJobStateArgs(ctx, operation, query, id, value)
}

func (db *Database) writeJobStateArgs(ctx context.Context, operation, query, id string, values ...any) error {
	if err := db.requirePostgres(operation); err != nil {
		return err
	}
	args := make([]any, 0, len(values)+1)
	args = append(args, id)
	args = append(args, values...)
	rows, err := db.querier.Query(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("%s: %w", operation, err)
	}
	defer rows.Close()
	if !rows.Next() {
		if err := rows.Err(); err != nil {
			return fmt.Errorf("%s: %w", operation, err)
		}
		return fmt.Errorf("%s: generation job not found", operation)
	}
	var returnedID string
	if err := rows.Scan(&returnedID); err != nil {
		return fmt.Errorf("%s: scan result: %w", operation, err)
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("%s: %w", operation, err)
	}
	return nil
}

func (db *Database) writeJobStateConflict(ctx context.Context, operation, query, id string, values ...any) error {
	err := db.writeJobStateArgs(ctx, operation, query, id, values...)
	if err != nil && strings.Contains(err.Error(), "generation job not found") {
		return jobs.ErrTransitionConflict
	}
	return err
}

func (db *Database) ClaimJobs(ctx context.Context, workerID string, limit, lockTimeoutSeconds int) ([]jobs.Job, error) {
	if err := db.requirePostgres("claim generation jobs"); err != nil {
		return nil, err
	}
	limit = max(1, min(limit, 20))
	if lockTimeoutSeconds <= 0 {
		lockTimeoutSeconds = 300
	}
	return db.queryJobs(ctx, ClaimJobsSQL, workerID, limit, lockTimeoutSeconds)
}

func (db *Database) mustFindJobByID(ctx context.Context, id string) (jobs.Job, error) {
	job, found, err := db.queryOneJob(ctx, `SELECT `+jobColumns+` FROM public.generation_jobs WHERE id = $1::text LIMIT 1`, id)
	if err != nil {
		return jobs.Job{}, err
	}
	if !found {
		return jobs.Job{}, fmt.Errorf("generation job not found: %s", id)
	}
	return job, nil
}

func (db *Database) queryJobs(ctx context.Context, query string, args ...any) ([]jobs.Job, error) {
	rows, err := db.querier.Query(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]jobs.Job, 0)
	for rows.Next() {
		item, err := scanJob(rows)
		if err != nil {
			return nil, err
		}
		items = append(items, item)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

func (db *Database) queryOneJob(ctx context.Context, query string, args ...any) (jobs.Job, bool, error) {
	rows, err := db.querier.Query(ctx, query, args...)
	if err != nil {
		return jobs.Job{}, false, err
	}
	defer rows.Close()
	if !rows.Next() {
		return jobs.Job{}, false, rows.Err()
	}
	job, err := scanJob(rows)
	if err != nil {
		return jobs.Job{}, false, err
	}
	if err := rows.Err(); err != nil {
		return jobs.Job{}, false, err
	}
	return job, true, nil
}

func scanJob(rows Rows) (jobs.Job, error) {
	var job jobs.Job
	var externalClientID, prompt, providerTaskID, retryOf, idempotencyKey, fingerprint sql.NullString
	var lockedBy, webhookURL sql.NullString
	var lockedAt, startedAt, completedAt, dispatchReadyAt, finalizedAt, providerDispatchStartedAt sql.NullTime
	var requestPayload, responsePayload, errorPayload, webhookStatus, usageContext, billing []byte
	var status string
	err := rows.Scan(
		&job.ID, &job.OwnerID, &externalClientID, &job.Capability, &job.Provider, &job.ReqKey, &status, &prompt,
		&job.InputAssetIDs, &job.InputURLs, &job.OutputAssetIDs, &providerTaskID, &providerDispatchStartedAt, &requestPayload, &responsePayload,
		&errorPayload, &retryOf, &idempotencyKey, &fingerprint, &job.Priority, &job.Attempts, &job.MaxAttempts,
		&job.ScheduledAt, &lockedAt, &lockedBy, &startedAt, &completedAt, &webhookURL, &job.WebhookAttempts,
		&webhookStatus, &usageContext, &billing, &dispatchReadyAt, &finalizedAt, &job.CreatedAt, &job.UpdatedAt,
	)
	if err != nil {
		return jobs.Job{}, fmt.Errorf("scan generation job: %w", err)
	}
	job.Status = jobs.Status(status)
	job.ExternalClientID = nullString(externalClientID)
	job.Prompt = nullString(prompt)
	job.ProviderTaskID = nullString(providerTaskID)
	job.ProviderDispatchStartedAt = nullTimePointer(providerDispatchStartedAt)
	job.RetryOf = nullString(retryOf)
	job.IdempotencyKey = nullString(idempotencyKey)
	job.IdempotencyFingerprint = nullString(fingerprint)
	job.LockedBy = nullString(lockedBy)
	job.WebhookURL = nullString(webhookURL)
	job.LockedAt = nullTimePointer(lockedAt)
	job.StartedAt = nullTimePointer(startedAt)
	job.CompletedAt = nullTimePointer(completedAt)
	job.DispatchReadyAt = nullTimePointer(dispatchReadyAt)
	job.FinalizedAt = nullTimePointer(finalizedAt)
	job.RequestPayload = normalizedJSON(requestPayload, `{}`)
	job.ResponsePayload = normalizedOptionalJSON(responsePayload)
	job.WebhookLastStatus = normalizedOptionalJSON(webhookStatus)
	job.UsageContext = normalizedOptionalJSON(usageContext)
	job.Billing = normalizedOptionalJSON(billing)
	if len(errorPayload) > 0 && string(errorPayload) != "null" {
		if err := json.Unmarshal(errorPayload, &job.Error); err != nil {
			return jobs.Job{}, fmt.Errorf("decode generation job error: %w", err)
		}
	}
	if job.InputAssetIDs == nil {
		job.InputAssetIDs = []string{}
	}
	if job.InputURLs == nil {
		job.InputURLs = []string{}
	}
	if job.OutputAssetIDs == nil {
		job.OutputAssetIDs = []string{}
	}
	return job, nil
}

func jobArguments(job jobs.Job) []any {
	errorPayload, _ := json.Marshal(job.Error)
	if job.Error == nil {
		errorPayload = nil
	}
	return []any{
		job.ID, job.OwnerID, optionalDatabaseText(job.ExternalClientID), job.Capability, job.Provider, job.ReqKey,
		string(job.Status), optionalDatabaseText(job.Prompt), job.InputAssetIDs, job.InputURLs, job.OutputAssetIDs,
		optionalDatabaseText(job.ProviderTaskID), job.ProviderDispatchStartedAt, job.RequestPayload, optionalJSON(job.ResponsePayload), optionalJSON(errorPayload),
		optionalDatabaseText(job.RetryOf), optionalDatabaseText(job.IdempotencyKey), optionalDatabaseText(job.IdempotencyFingerprint),
		job.Priority, job.Attempts, job.MaxAttempts, job.ScheduledAt, job.LockedAt, optionalDatabaseText(job.LockedBy),
		job.StartedAt, job.CompletedAt, optionalDatabaseText(job.WebhookURL), job.WebhookAttempts,
		optionalJSON(job.WebhookLastStatus), optionalJSON(job.UsageContext), optionalJSON(job.Billing), job.DispatchReadyAt, job.FinalizedAt, job.CreatedAt, job.UpdatedAt,
	}
}

func (db *Database) requirePostgres(operation string) error {
	if db == nil || db.config.Backend != BackendPostgres || db.querier == nil {
		return fmt.Errorf("%s: PostgreSQL is unavailable", operation)
	}
	return nil
}

func nullString(value sql.NullString) string {
	if value.Valid {
		return value.String
	}
	return ""
}

func nullTimePointer(value sql.NullTime) *time.Time {
	if !value.Valid {
		return nil
	}
	copy := value.Time
	return &copy
}

func optionalJSON(value []byte) any {
	if len(value) == 0 || string(value) == "null" {
		return nil
	}
	return value
}

func normalizedJSON(value []byte, fallback string) json.RawMessage {
	if len(value) == 0 || string(value) == "null" {
		return json.RawMessage(fallback)
	}
	return append(json.RawMessage(nil), value...)
}

func normalizedOptionalJSON(value []byte) json.RawMessage {
	if len(value) == 0 || string(value) == "null" {
		return nil
	}
	return append(json.RawMessage(nil), value...)
}

type sqlStateError interface{ SQLState() string }

func sqlState(err error) string {
	for err != nil {
		if state, ok := err.(sqlStateError); ok {
			return state.SQLState()
		}
		type unwrapper interface{ Unwrap() error }
		wrapped, ok := err.(unwrapper)
		if !ok {
			break
		}
		err = wrapped.Unwrap()
	}
	return ""
}

var _ jobs.Store = (*Database)(nil)