package query

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/pingcap/log"
	"go.uber.org/zap"

	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/utils/metacache"
)

const (
	statusPollInterval    = 200 * time.Millisecond
	defaultResultRows     = int64(1000)
	maxResultRows         = int64(1000) // Athena GetQueryResults MaxResults hard limit
	defaultDownloadExpire = time.Hour
	maxDownloadExpire     = 12 * time.Hour
)

type clusterHeader struct {
	Provider   string `header:"x-provider" binding:"required"`
	Region     string `header:"x-region" binding:"required"`
	TenantID   string `header:"x-org-id"`
	ProjectID  string `header:"x-project-id"`
	ClusterID  string `header:"x-cluster-id" binding:"required"`
	DeployType string `header:"x-deploy-type"`
	Env        string `header:"x-env"`
	BizType    string `header:"x-biz-type"`
}

// SubmitRequest is the generic SQL analytics query body.
//
// Scoped datasets (logs, log-stats, topsql, …): scope maps dive-path partitions onto a
// WITH <alias> AS (...) CTE; user SQL queries the virtual alias.
// dataset=raw (default / omitted): user SQL is passed through unchanged.
type SubmitRequest struct {
	SQL        string     `json:"sql"`
	Database   string     `json:"database"`
	Dataset    string     `json:"dataset"`
	OutputMode string     `json:"output_mode"`
	Scope      Scope      `json:"scope"`
	TimeRange  *TimeRange `json:"time_range"`
	// WaitSeconds blocks the submit request until the query finishes or the timeout elapses.
	// nil => DefaultWaitSeconds (env OLAP_DEFAULT_WAIT_SECONDS, default 30);
	// 0 => async only (legacy behavior); clamped to MaxWaitSeconds (env OLAP_MAX_WAIT_SECONDS, default 60).
	WaitSeconds *int `json:"wait_seconds"`
	// MaxRows controls how many rows are returned in wait-mode inline result (Athena max 1000).
	MaxRows *int64 `json:"max_rows"`
}

type errorResponse struct {
	Error    bool   `json:"error"`
	Message  string `json:"message"`
	FullText string `json:"full_text,omitempty"`
	Code     int    `json:"code,omitempty"`
}

type okEnvelope struct {
	Status string `json:"status"`
	Data   any    `json:"data"`
}

// Service hosts the /olap HTTP API (OLAP SQL analytics).
type Service struct {
	cfg     Config
	engine  Engine
	jobs    *jobStore
	catalog *Catalog
	meta    metacache.MetaCache // optional; used to map applicationID → o11y cluster UUID for slowlogs
}

// NewService constructs the query service for the local vendor/region.
// meta may be nil (cluster id → o11y id resolution for slowlogs is then skipped).
func NewService(cfg Config, meta metacache.MetaCache) (*Service, error) {
	cfg.normalizeWaitSeconds()
	eng, err := NewEngine(cfg)
	if err != nil {
		return nil, err
	}
	cat, err := NewCatalog(cfg)
	if err != nil {
		return nil, err
	}
	s := &Service{cfg: cfg, engine: eng, jobs: newJobStore(), catalog: cat, meta: meta}
	go s.ensureOnStartup()
	return s, nil
}

func (s *Service) ensureOnStartup() {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	if ae, ok := s.engine.(*athenaEngine); ok {
		if err := ae.EnsureWorkGroup(ctx); err != nil {
			log.Warn("olap athena workgroup ensure failed", zap.Error(err),
				zap.String("workgroup", s.cfg.ResolvedAthenaWorkGroup()),
				zap.String("output_location", s.cfg.AthenaOutputLocation))
		}
	}

	if s.catalog == nil {
		return
	}
	if err := s.catalog.EnsureAll(ctx); err != nil {
		log.Warn("olap catalog startup ensure failed", zap.Error(err))
		return
	}
	log.Info("olap catalog tables ensured",
		zap.String("region", s.cfg.Region),
		zap.String("logical_region", s.cfg.LogsTableRegionKey()),
		zap.String("workgroup", s.cfg.ResolvedAthenaWorkGroup()),
		zap.String("output_location", s.cfg.AthenaOutputLocation))
}

// InitOlapApi registers OLAP SQL analytics routes alongside slowquery/sqls.
func InitOlapApi(g *gin.RouterGroup, s *Service) {
	InitMetrics()
	g.GET("/olap/tables", s.handleListTables)
	g.POST("/olap/scope/list", s.handleListScope)
	g.POST("/olap", s.handleSubmit)
	g.GET("/olap/:query_id", s.handleStatus)
	g.GET("/olap/:query_id/results", s.handleResults)
	g.GET("/olap/:query_id/download", s.handleDownload)
	g.DELETE("/olap/:query_id", s.handleCancel)
}

func (s *Service) resolveWaitSeconds(v *int) int {
	s.cfg.normalizeWaitSeconds()
	if v == nil {
		return s.cfg.DefaultWaitSeconds
	}
	if *v <= 0 {
		return 0
	}
	if *v > s.cfg.MaxWaitSeconds {
		return s.cfg.MaxWaitSeconds
	}
	return *v
}

func resolveMaxRows(v *int64) (int64, error) {
	if v == nil {
		return defaultResultRows, nil
	}
	if *v <= 0 || *v > maxResultRows {
		return 0, fmt.Errorf("max_rows must be 1..%d", maxResultRows)
	}
	return *v, nil
}

func resolveDownloadExpire(raw string) (time.Duration, error) {
	if raw == "" {
		return defaultDownloadExpire, nil
	}
	sec, err := strconv.ParseInt(raw, 10, 64)
	if err != nil || sec <= 0 {
		return 0, fmt.Errorf("expires_seconds must be a positive integer")
	}
	d := time.Duration(sec) * time.Second
	if d > maxDownloadExpire {
		return maxDownloadExpire, nil
	}
	return d, nil
}

func isTerminalStatus(status string) bool {
	switch status {
	case StatusSucceeded, StatusFailed, StatusCancelled:
		return true
	default:
		return false
	}
}

func (s *Service) handleSubmit(c *gin.Context) {
	var hdr clusterHeader
	if err := c.ShouldBindHeader(&hdr); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}
	if err := s.checkRegion(hdr.Region); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid region", err.Error())
		return
	}
	if normalizeProvider(hdr.Provider) != normalizeProvider(s.cfg.Vendor) {
		writeErr(c, http.StatusBadRequest, "invalid provider",
			"x-provider "+hdr.Provider+" does not match deployment vendor "+s.cfg.Vendor)
		return
	}

	var req SubmitRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}
	if strings.TrimSpace(req.SQL) == "" {
		writeErr(c, http.StatusBadRequest, "invalid request", "sql is required")
		return
	}
	if err := ValidateSelectSQL(req.SQL); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid sql", err.Error())
		return
	}
	if err := RejectDangerousSQL(req.SQL); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid sql", err.Error())
		return
	}
	dataset := normalizeDataset(req.Dataset)
	queryDB := firstNonEmpty(req.Database, s.cfg.AthenaDatabase)
	if isScopedDataset(dataset) {
		queryDB = firstNonEmpty(req.Database, logsDatabaseName)
	}
	if err := s.ensureCatalogForSubmit(c.Request.Context(), req.SQL, queryDB, dataset); err != nil {
		var nr *NotReadyError
		if errors.As(err, &nr) {
			writeErr(c, http.StatusNotImplemented, "engine not ready", err.Error())
			return
		}
		log.Warn("olap catalog ensure failed", zap.Error(err))
		writeErr(c, http.StatusBadGateway, "catalog ensure failed", err.Error())
		return
	}
	outputMode := req.OutputMode
	if outputMode == "" {
		outputMode = OutputModeInline
	}
	if outputMode != OutputModeInline && outputMode != OutputModeExport {
		writeErr(c, http.StatusBadRequest, "invalid request", "output_mode must be inline or export")
		return
	}
	maxRows, err := resolveMaxRows(req.MaxRows)
	if err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}

	tenantID := hdr.TenantID
	if org, ok := req.Scope.GetString("org"); ok && org != "" {
		tenantID = org
	}

	if normalizeDataset(dataset) == DatasetSlowlogs {
		if err := s.resolveSlowlogsClusterID(c.Request.Context(), hdr, req.Scope); err != nil {
			writeErr(c, http.StatusBadRequest, "invalid scope", err.Error())
			return
		}
	}

	effectiveSQL, err := BuildEffectiveSQL(req.SQL, RewriteOptions{
		Dataset:         dataset,
		Region:          s.cfg.LogsTableRegionKey(),
		LogsTableSuffix: s.cfg.LogsTableSuffix,
		Scope:           req.Scope,
		TimeRange:       req.TimeRange,
	})
	if err != nil {
		writeErr(c, http.StatusBadRequest, "invalid sql", err.Error())
		return
	}

	handle, err := s.engine.Submit(c.Request.Context(), &EngineSubmitRequest{
		SQL:      effectiveSQL,
		Database: queryDB,
	})
	if err != nil {
		var nr *NotReadyError
		if errors.As(err, &nr) {
			writeErr(c, http.StatusNotImplemented, "engine not ready", err.Error())
			return
		}
		log.Error("query submit failed", zap.Error(err), zap.String("cluster", hdr.ClusterID))
		writeErr(c, http.StatusBadGateway, "query submit failed", err.Error())
		return
	}

	now := time.Now().UTC()
	job := s.jobs.create(&Job{
		Provider:       normalizeProvider(hdr.Provider),
		Region:         s.cfg.Region,
		EngineName:     handle.EngineName,
		ExecutionID:    handle.ExecutionID,
		SQL:            req.SQL,
		EffectiveSQL:   effectiveSQL,
		Database:       queryDB,
		Dataset:        dataset,
		OutputMode:     outputMode,
		ClusterID:      hdr.ClusterID,
		TenantID:       tenantID,
		GitHubUsername: githubUsernameFromContext(c),
		Status:         StatusRunning,
		SubmittedAt:    now,
	})
	logOlapQuerySubmitted(job, dataset, s.cfg.ResolvedAthenaWorkGroup())

	data := gin.H{
		"query_id":      job.ID,
		"engine":        job.EngineName,
		"region":        job.Region,
		"database":      job.Database,
		"dataset":       dataset,
		"status":        job.Status,
		"submitted_at":  job.SubmittedAt.Format(time.RFC3339),
		"output_mode":   job.OutputMode,
		"effective_sql": job.EffectiveSQL,
		"wait_seconds":  s.resolveWaitSeconds(req.WaitSeconds),
	}
	httpStatus := http.StatusAccepted

	waitSeconds := s.resolveWaitSeconds(req.WaitSeconds)
	if waitSeconds > 0 {
		st, waitErr := s.waitForTerminal(c.Request.Context(), handle, time.Duration(waitSeconds)*time.Second)
		if waitErr != nil {
			log.Warn("query wait failed", zap.Error(waitErr), zap.String("query_id", job.ID))
		} else if st != nil {
			updated := s.applyExecutionStatus(job.ID, st)
			if updated == nil {
				updated = job
			}
			data["status"] = updated.Status
			data["scanned_bytes"] = updated.ScannedBytes
			data["engine_execution_time_ms"] = updated.EngineExecutionTimeMillis
			if updated.CompletedAt != nil {
				data["completed_at"] = updated.CompletedAt.Format(time.RFC3339)
			}
			if isTerminalStatus(updated.Status) {
				httpStatus = http.StatusOK
			}
			if updated.Error != "" && (updated.Status == StatusFailed || updated.Status == StatusCancelled) {
				data["error_message"] = updated.Error
			}
			if updated.Status == StatusSucceeded {
				switch outputMode {
				case OutputModeInline:
					result, fetchErr := s.fetchResultPayload(c.Request.Context(), handle, job.ID, "", maxRows)
					if fetchErr != nil {
						log.Warn("query wait fetch results failed", zap.Error(fetchErr), zap.String("query_id", job.ID))
					} else {
						data["result"] = result
					}
				case OutputModeExport:
					result, downloadErr := s.fetchDownloadPayload(c.Request.Context(), handle, job.ID, st.ScannedBytes, defaultDownloadExpire)
					if downloadErr != nil {
						log.Warn("query wait presign download failed", zap.Error(downloadErr), zap.String("query_id", job.ID))
						data["download_path"] = "/api/v1/olap/" + job.ID + "/download"
					} else {
						data["result"] = result
					}
				}
			}
		}
	}

	c.JSON(httpStatus, okEnvelope{Status: "ok", Data: data})
}

func (s *Service) handleListTables(c *gin.Context) {
	var hdr clusterHeader
	if err := c.ShouldBindHeader(&hdr); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}
	if err := s.checkRegion(hdr.Region); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid region", err.Error())
		return
	}
	if normalizeProvider(hdr.Provider) != normalizeProvider(s.cfg.Vendor) {
		writeErr(c, http.StatusBadRequest, "invalid provider",
			"x-provider "+hdr.Provider+" does not match deployment vendor "+s.cfg.Vendor)
		return
	}
	tables, err := s.catalog.ListTables(c.Request.Context())
	if err != nil {
		writeErr(c, http.StatusBadGateway, "list tables failed", err.Error())
		return
	}
	c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: gin.H{
		"region":   s.cfg.Region,
		"database": logsDatabaseName,
		"datasets": listDatasetInfos(s.cfg),
		"tables":   tables,
	}})
}

func (s *Service) handleListScope(c *gin.Context) {
	var hdr clusterHeader
	if err := c.ShouldBindHeader(&hdr); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}
	if err := s.checkRegion(hdr.Region); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid region", err.Error())
		return
	}
	if normalizeProvider(hdr.Provider) != normalizeProvider(s.cfg.Vendor) {
		writeErr(c, http.StatusBadRequest, "invalid provider",
			"x-provider "+hdr.Provider+" does not match deployment vendor "+s.cfg.Vendor)
		return
	}
	var req ListScopeRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}
	if strings.TrimSpace(req.Dataset) == "" || strings.TrimSpace(req.Field) == "" {
		writeErr(c, http.StatusBadRequest, "invalid request", "dataset and field are required")
		return
	}
	if s.catalog == nil {
		writeErr(c, http.StatusNotImplemented, "engine not ready", "catalog not configured")
		return
	}
	result, err := s.catalog.ListScopeValues(c.Request.Context(), req)
	if err != nil {
		msg := err.Error()
		if strings.Contains(msg, "not available") {
			writeErr(c, http.StatusNotImplemented, "engine not ready", msg)
			return
		}
		if strings.Contains(msg, "requires") || strings.Contains(msg, "unknown") ||
			strings.Contains(msg, "max_keys") || strings.Contains(msg, "does not support") ||
			strings.Contains(msg, "must be") {
			writeErr(c, http.StatusBadRequest, "invalid request", msg)
			return
		}
		log.Warn("olap scope list failed", zap.Error(err))
		writeErr(c, http.StatusBadGateway, "scope list failed", msg)
		return
	}
	c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: result})
}

func (s *Service) ensureCatalogForSubmit(ctx context.Context, sql, defaultDatabase, dataset string) error {
	if s.catalog == nil {
		return nil
	}
	// Always ensure o11y-logs DB + registry tables for scoped datasets,
	// and when raw SQL may reference them.
	if isScopedDataset(dataset) {
		return s.catalog.EnsureAll(ctx)
	}
	return s.ensureCatalogForSQL(ctx, sql, defaultDatabase)
}

func (s *Service) ensureCatalogForSQL(ctx context.Context, sql, defaultDatabase string) error {
	if s.catalog == nil {
		return nil
	}
	refs, err := ExtractTableReferences(sql, defaultDatabase)
	if err != nil {
		// Athena dialect may not parse with the TiDB parser; still ensure registry tables.
		log.Warn("olap extract table refs failed; ensuring all registry tables", zap.Error(err))
		return s.catalog.EnsureAll(ctx)
	}
	if len(refs) == 0 {
		return s.catalog.EnsureAll(ctx)
	}
	return s.catalog.EnsureReferenced(ctx, refs)
}

func (s *Service) waitForTerminal(ctx context.Context, handle *Handle, timeout time.Duration) (*ExecutionStatus, error) {
	deadline := time.Now().Add(timeout)
	var last *ExecutionStatus
	for {
		st, err := s.engine.Status(ctx, handle)
		if err != nil {
			return last, err
		}
		last = st
		if isTerminalStatus(st.Status) {
			return st, nil
		}
		if !time.Now().Before(deadline) {
			return st, nil
		}
		remaining := time.Until(deadline)
		sleep := statusPollInterval
		if sleep > remaining {
			sleep = remaining
		}
		timer := time.NewTimer(sleep)
		select {
		case <-ctx.Done():
			timer.Stop()
			return last, ctx.Err()
		case <-timer.C:
		}
	}
}

func (s *Service) fetchResultPayload(ctx context.Context, handle *Handle, queryID, pageToken string, maxRows int64) (gin.H, error) {
	page, err := s.engine.FetchPage(ctx, handle, pageToken, maxRows)
	if err != nil {
		return nil, err
	}
	logOlapQueryResults(queryID, page)
	return formatResultPayload(queryID, page), nil
}

func (s *Service) fetchDownloadPayload(ctx context.Context, handle *Handle, queryID string, scannedBytes int64, expires time.Duration) (gin.H, error) {
	presigned, err := s.engine.PresignResult(ctx, handle, expires)
	if err != nil {
		return nil, err
	}
	return formatDownloadPayload(queryID, scannedBytes, presigned), nil
}

func formatResultPayload(queryID string, page *Page) gin.H {
	rows := make([]map[string]any, 0, len(page.Rows))
	for _, r := range page.Rows {
		m := map[string]any{}
		for i, col := range page.Columns {
			var val any
			if i < len(r) {
				val = r[i]
			}
			m[col.Name] = val
		}
		rows = append(rows, m)
	}
	return gin.H{
		"query_id":        queryID,
		"columns":         page.Columns,
		"rows":            rows,
		"next_page_token": page.NextToken,
		"has_more":        page.NextToken != "",
	}
}

func formatDownloadPayload(queryID string, scannedBytes int64, presigned *PresignedResult) gin.H {
	item := gin.H{
		"url":        presigned.URL,
		"expires_at": presigned.ExpiresAt.Format(time.RFC3339),
	}
	if presigned.Bytes != nil {
		item["bytes"] = *presigned.Bytes
	}
	return gin.H{
		"query_id":      queryID,
		"status":        StatusSucceeded,
		"format":        firstNonEmpty(presigned.Format, "csv"),
		"s3_uri":        presigned.S3URI,
		"file_count":    1,
		"download_urls": []gin.H{item},
		"expires_at":    presigned.ExpiresAt.Format(time.RFC3339),
		"scanned_bytes": scannedBytes,
	}
}

func (s *Service) handleStatus(c *gin.Context) {
	job, hdr, ok := s.loadJobAuthorized(c)
	if !ok {
		return
	}
	st, err := s.engine.Status(c.Request.Context(), &Handle{EngineName: job.EngineName, ExecutionID: job.ExecutionID})
	if err != nil {
		var nr *NotReadyError
		if errors.As(err, &nr) {
			writeErr(c, http.StatusNotImplemented, "engine not ready", err.Error())
			return
		}
		writeErr(c, http.StatusBadGateway, "query status failed", err.Error())
		return
	}
	updated := s.applyExecutionStatus(job.ID, st)
	if updated == nil {
		updated = job
	}
	_ = hdr
	data := gin.H{
		"query_id":                 updated.ID,
		"engine":                   updated.EngineName,
		"region":                   updated.Region,
		"status":                   updated.Status,
		"submitted_at":             updated.SubmittedAt.Format(time.RFC3339),
		"scanned_bytes":            updated.ScannedBytes,
		"engine_execution_time_ms": updated.EngineExecutionTimeMillis,
		"total_execution_time_ms":  updated.TotalExecutionTimeMillis,
		"output_mode":              updated.OutputMode,
	}
	if updated.CompletedAt != nil {
		data["completed_at"] = updated.CompletedAt.Format(time.RFC3339)
	}
	if updated.Error != "" && (updated.Status == StatusFailed || updated.Status == StatusCancelled) {
		data["error_message"] = updated.Error
	}
	c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: data})
}

func (s *Service) handleResults(c *gin.Context) {
	job, _, ok := s.loadJobAuthorized(c)
	if !ok {
		return
	}
	st, err := s.engine.Status(c.Request.Context(), &Handle{EngineName: job.EngineName, ExecutionID: job.ExecutionID})
	if err != nil {
		writeErr(c, http.StatusBadGateway, "query status failed", err.Error())
		return
	}
	if st.Status != StatusSucceeded {
		writeErr(c, http.StatusConflict, "query not ready", "status="+st.Status)
		return
	}
	maxRows := defaultResultRows
	if v := c.Query("max_rows"); v != "" {
		n, err := strconv.ParseInt(v, 10, 64)
		if err != nil {
			writeErr(c, http.StatusBadRequest, "invalid request", "max_rows must be 1..1000")
			return
		}
		resolved, err := resolveMaxRows(&n)
		if err != nil {
			writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
			return
		}
		maxRows = resolved
	}
	result, err := s.fetchResultPayload(c.Request.Context(), &Handle{EngineName: job.EngineName, ExecutionID: job.ExecutionID}, job.ID, c.Query("page_token"), maxRows)
	if err != nil {
		writeErr(c, http.StatusBadGateway, "fetch results failed", err.Error())
		return
	}
	c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: result})
}

func (s *Service) handleDownload(c *gin.Context) {
	job, _, ok := s.loadJobAuthorized(c)
	if !ok {
		return
	}
	st, err := s.engine.Status(c.Request.Context(), &Handle{EngineName: job.EngineName, ExecutionID: job.ExecutionID})
	if err != nil {
		var nr *NotReadyError
		if errors.As(err, &nr) {
			writeErr(c, http.StatusNotImplemented, "engine not ready", err.Error())
			return
		}
		writeErr(c, http.StatusBadGateway, "query status failed", err.Error())
		return
	}
	if st.Status != StatusSucceeded {
		writeErr(c, http.StatusConflict, "query not ready", "status="+st.Status)
		return
	}
	expires, err := resolveDownloadExpire(c.Query("expires_seconds"))
	if err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return
	}
	presigned, err := s.engine.PresignResult(c.Request.Context(), &Handle{EngineName: job.EngineName, ExecutionID: job.ExecutionID}, expires)
	if err != nil {
		var nr *NotReadyError
		if errors.As(err, &nr) {
			writeErr(c, http.StatusNotImplemented, "engine not ready", err.Error())
			return
		}
		writeErr(c, http.StatusBadGateway, "presign result failed", err.Error())
		return
	}
	c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: formatDownloadPayload(job.ID, st.ScannedBytes, presigned)})
}

func (s *Service) handleCancel(c *gin.Context) {
	job, _, ok := s.loadJobAuthorized(c)
	if !ok {
		return
	}
	handle := &Handle{EngineName: job.EngineName, ExecutionID: job.ExecutionID}
	if err := s.engine.Cancel(c.Request.Context(), handle); err != nil {
		writeErr(c, http.StatusBadGateway, "cancel failed", err.Error())
		return
	}
	// Prefer Athena's final statistics when available for chargeback.
	if st, err := s.engine.Status(c.Request.Context(), handle); err == nil && st != nil {
		if st.Status == "" || st.Status == StatusRunning || st.Status == StatusQueued {
			st.Status = StatusCancelled
		}
		if st.CompletedAt == nil {
			now := time.Now().UTC()
			st.CompletedAt = &now
		}
		updated := s.applyExecutionStatus(job.ID, st)
		if updated != nil {
			c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: gin.H{
				"query_id":                 updated.ID,
				"status":                   updated.Status,
				"scanned_bytes":            updated.ScannedBytes,
				"engine_execution_time_ms": updated.EngineExecutionTimeMillis,
			}})
			return
		}
	}
	s.jobs.update(job.ID, func(j *Job) {
		j.Status = StatusCancelled
		now := time.Now().UTC()
		j.CompletedAt = &now
		if !j.UsageLogged {
			j.UsageLogged = true
			cp := *j
			logOlapQueryFinished(&cp, &ExecutionStatus{Status: StatusCancelled}, s.cfg.ResolvedAthenaWorkGroup())
		}
	})
	c.JSON(http.StatusOK, okEnvelope{Status: "ok", Data: gin.H{"query_id": job.ID, "status": StatusCancelled}})
}

// applyExecutionStatus persists engine status onto the job and emits a one-shot
// usage log the first time the query reaches a terminal state.
func (s *Service) applyExecutionStatus(queryID string, st *ExecutionStatus) *Job {
	if st == nil {
		return nil
	}
	var finished Job
	shouldLog := false
	updated, ok := s.jobs.update(queryID, func(j *Job) {
		j.Status = st.Status
		j.Error = st.StateChangeReason
		j.ScannedBytes = st.ScannedBytes
		j.EngineExecutionTimeMillis = st.EngineExecutionTimeMillis
		j.TotalExecutionTimeMillis = st.TotalExecutionTimeMillis
		j.QueryQueueTimeMillis = st.QueryQueueTimeMillis
		j.QueryPlanningTimeMillis = st.QueryPlanningTimeMillis
		if st.CompletedAt != nil {
			j.CompletedAt = st.CompletedAt
		}
		if isTerminalStatus(st.Status) && !j.UsageLogged {
			j.UsageLogged = true
			finished = *j
			shouldLog = true
		}
	})
	if !ok {
		return nil
	}
	if shouldLog {
		logOlapQueryFinished(&finished, st, s.cfg.ResolvedAthenaWorkGroup())
	}
	return updated
}

func (s *Service) loadJobAuthorized(c *gin.Context) (*Job, clusterHeader, bool) {
	var hdr clusterHeader
	if err := c.ShouldBindHeader(&hdr); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid request", err.Error())
		return nil, hdr, false
	}
	if err := s.checkRegion(hdr.Region); err != nil {
		writeErr(c, http.StatusBadRequest, "invalid region", err.Error())
		return nil, hdr, false
	}
	id := c.Param("query_id")
	job, ok := s.jobs.get(id)
	if !ok {
		writeErr(c, http.StatusNotFound, "query not found", id)
		return nil, hdr, false
	}
	if job.ClusterID != hdr.ClusterID {
		writeErr(c, http.StatusForbidden, "forbidden", "query does not belong to this cluster")
		return nil, hdr, false
	}
	return job, hdr, true
}

func (s *Service) checkRegion(reqRegion string) error {
	if !strings.EqualFold(strings.TrimSpace(reqRegion), strings.TrimSpace(s.cfg.Region)) {
		return errors.New("x-region " + reqRegion + " does not match deployment region " + s.cfg.Region)
	}
	return nil
}

func writeErr(c *gin.Context, code int, message, full string) {
	c.JSON(code, errorResponse{Error: true, Message: message, FullText: full, Code: code})
}