package kb

import (
	"context"
	"fmt"
	"strings"
	"time"
)

const listLatestTerminalSourceJobsForSourcesSQL = `
	SELECT
		id,
		workspace_id,
		knowledge_base_id,
		source_id,
		job_type,
		status,
		progress_json,
		error,
		attempt_count,
		lease_owner,
		lease_until,
		started_at,
		finished_at,
		created_at,
		updated_at
	FROM (
		SELECT knowledge_base_jobs.*,
			ROW_NUMBER() OVER (PARTITION BY source_id ORDER BY updated_at DESC, created_at DESC, id DESC) AS source_job_rank
		FROM knowledge_base_jobs
		WHERE workspace_id = ?
			AND knowledge_base_id = ?
			AND source_id IN ?
			AND job_type IN ?
			AND status IN ?
	) ranked_jobs
	WHERE source_job_rank = 1
	ORDER BY updated_at DESC, created_at DESC, id DESC`

func (s *Store) ListActiveSourceJobsForSources(ctx context.Context, workspaceID, kbID string, sourceIDs []string, jobTypes []string) ([]Job, error) {
	if err := s.ensureDB(); err != nil {
		return nil, err
	}
	sourceIDs = trimmedUniqueNonEmpty(sourceIDs)
	jobTypes = trimmedUniqueNonEmpty(jobTypes)
	if len(sourceIDs) == 0 || len(jobTypes) == 0 {
		return nil, nil
	}
	var jobs []Job
	err := s.db.WithContext(ctx).
		Where("workspace_id = ? AND knowledge_base_id = ? AND source_id IN ? AND job_type IN ? AND status IN ?", strings.TrimSpace(workspaceID), strings.TrimSpace(kbID), sourceIDs, jobTypes, []string{"queued", "running"}).
		Order("updated_at DESC").
		Order("created_at DESC").
		Order("id DESC").
		Find(&jobs).Error
	if err != nil {
		return nil, fmt.Errorf("list active kb source jobs: %w", err)
	}
	return jobs, nil
}

func (s *Store) ListLatestTerminalSourceJobsForSources(ctx context.Context, workspaceID, kbID string, sourceIDs []string, jobTypes []string) ([]Job, error) {
	if err := s.ensureDB(); err != nil {
		return nil, err
	}
	sourceIDs = trimmedUniqueNonEmpty(sourceIDs)
	jobTypes = trimmedUniqueNonEmpty(jobTypes)
	if len(sourceIDs) == 0 || len(jobTypes) == 0 {
		return nil, nil
	}
	var jobs []Job
	if err := s.db.WithContext(ctx).Raw(listLatestTerminalSourceJobsForSourcesSQL, strings.TrimSpace(workspaceID), strings.TrimSpace(kbID), sourceIDs, jobTypes, []string{"succeeded", "failed"}).Scan(&jobs).Error; err != nil {
		return nil, fmt.Errorf("list terminal kb source jobs: %w", err)
	}
	return jobs, nil
}

func (s *Store) ListRecentJobsForKnowledgeBase(ctx context.Context, workspaceID, kbID string, jobTypes []string, limit int) ([]Job, error) {
	if err := s.ensureDB(); err != nil {
		return nil, err
	}
	jobTypes = trimmedUniqueNonEmpty(jobTypes)
	if limit <= 0 {
		limit = 10
	}
	if limit > 100 {
		limit = 100
	}
	var jobs []Job
	query := s.db.WithContext(ctx).
		Where("workspace_id = ? AND knowledge_base_id = ?", strings.TrimSpace(workspaceID), strings.TrimSpace(kbID))
	if len(jobTypes) > 0 {
		query = query.Where("job_type IN ?", jobTypes)
	}
	err := query.
		Order("updated_at DESC").
		Order("created_at DESC").
		Order("id DESC").
		Limit(limit).
		Find(&jobs).Error
	if err != nil {
		return nil, fmt.Errorf("list recent kb jobs: %w", err)
	}
	return jobs, nil
}

// CancelSourceJobs marks all queued or running jobs for a source and job type
// as cancelled, clearing lease ownership and recording the finish time. It
// returns the number of jobs cancelled. Running jobs whose lease is still
// valid will be cancelled in the database; the worker that currently holds the
// lease will detect the cancellation on its next leased update and stop.
func (s *Store) CancelSourceJobs(ctx context.Context, workspaceID, kbID, sourceID, jobType string, now time.Time) (int64, error) {
	if err := s.ensureDB(); err != nil {
		return 0, err
	}
	sourceID = strings.TrimSpace(sourceID)
	jobType = strings.TrimSpace(jobType)
	if sourceID == "" || jobType == "" {
		return 0, nil
	}
	result := s.db.WithContext(ctx).
		Model(&Job{}).
		Where(
			"workspace_id = ? AND knowledge_base_id = ? AND source_id = ? AND job_type = ? AND status IN ?",
			strings.TrimSpace(workspaceID),
			strings.TrimSpace(kbID),
			sourceID,
			jobType,
			[]string{"queued", "running"},
		).
		Updates(map[string]any{
			"status":      "cancelled",
			"lease_owner": "",
			"lease_until": nil,
			"finished_at": now,
			"updated_at":  now,
		})
	if result.Error != nil {
		return 0, fmt.Errorf("cancel kb source jobs: %w", result.Error)
	}
	return result.RowsAffected, nil
}

func trimmedUniqueNonEmpty(values []string) []string {
	if len(values) == 0 {
		return nil
	}
	seen := map[string]struct{}{}
	out := make([]string, 0, len(values))
	for _, value := range values {
		trimmed := strings.TrimSpace(value)
		if trimmed == "" {
			continue
		}
		if _, ok := seen[trimmed]; ok {
			continue
		}
		seen[trimmed] = struct{}{}
		out = append(out, trimmed)
	}
	return out
}
