package connector_workqueue

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

// ReconcileLaneReadyStates shares the single-lane transition: choose heads,
// demote other due ready rows, then promote only when no owner is processing.
// A caller using WithTx keeps every phase in its existing admitted transaction.
func (r *Repository) ReconcileLaneReadyStates(ctx context.Context, inputs []CreateWorkItemInput) error {
	type scope struct {
		pack       string
		generation int64
		now        time.Time
	}
	groups := map[scope]map[string]struct{}{}
	for _, input := range inputs {
		in := r.defaultCreateWorkItemInput(input)
		if strings.TrimSpace(in.InstalledPackID) == "" || in.Generation < 0 {
			return ErrInvalidWorkItem
		}
		lane := workItemLaneKey(in)
		if lane == "" {
			continue
		}
		key := scope{in.InstalledPackID, in.Generation, in.Now.UTC()}
		if groups[key] == nil {
			groups[key] = map[string]struct{}{}
		}
		groups[key][lane] = struct{}{}
	}
	for key, lanes := range groups {
		keys := make([]string, 0, len(lanes))
		for lane := range lanes {
			keys = append(keys, lane)
		}
		sort.Strings(keys)
		in := CreateWorkItemInput{InstalledPackID: key.pack, Generation: key.generation, Now: key.now}
		for offset := 0; offset < len(keys); offset += 128 {
			chunk := keys[offset:min(offset+128, len(keys))]
			if len(chunk) == 1 {
				in.SourceLaneKey = chunk[0]
				if err := r.reconcileLaneReadyState(ctx, in); err != nil {
					return err
				}
				continue
			}
			heads, err := r.nextLaneReadyHeads(ctx, in, chunk, key.now)
			if err != nil {
				return err
			}
			if err := r.demoteNonHeadLaneReadyBatch(ctx, in, chunk, heads, key.now); err != nil {
				return err
			}
			if len(heads) > 0 {
				if err := r.promoteLaneReadyBatch(ctx, in, chunk, heads, key.now); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func inSQL(values []string) string {
	return "(" + strings.TrimRight(strings.Repeat("?,", len(values)), ",") + ")"
}

func equalsOrIn(column string, values []string) string {
	if len(values) == 1 {
		return column + " = ?"
	}
	return column + " IN " + inSQL(values)
}
func appendStrings(args []any, values []string) []any {
	for _, value := range values {
		args = append(args, value)
	}
	return args
}

func laneReadyBatchQuery(in CreateWorkItemInput, lanes []string, now time.Time) (string, []any) {
	args := appendStrings([]any{in.InstalledPackID, in.Generation}, lanes)
	args = append(args, WorkStatePending, WorkStateRetry, WorkStateBlocked, WorkStateBlockedRetry, now)
	args = append(args, args...)
	selectSQL := "w.WorkId"
	if len(lanes) > 1 {
		selectSQL += `, ROW_NUMBER() OVER (
 PARTITION BY COALESCE(NULLIF(w.TargetMutationLaneKey,''),w.SourceLaneKey,'')
 ORDER BY ` + laneReadyHeadOrderBy("w", "e") + `) AS lane_position`
	}
	query := `SELECT ` + selectSQL + `
 FROM (
 SELECT * FROM _System_ConnectorWorkItem USE INDEX (idx_connector_work_item_target_lane)
 WHERE InstalledPackId = ? AND Generation = ? AND ` + equalsOrIn("TargetMutationLaneKey", lanes) + `
 AND State IN (?, ?, ?, ?) AND RunAfter <= ?
 UNION ALL
 SELECT * FROM _System_ConnectorWorkItem USE INDEX (idx_connector_work_item_source_lane_gen)
 WHERE InstalledPackId = ? AND Generation = ? AND ` + equalsOrIn("SourceLaneKey", lanes) + `
 AND (TargetMutationLaneKey IS NULL OR TargetMutationLaneKey = '')
 AND State IN (?, ?, ?, ?) AND RunAfter <= ?
 ) w LEFT JOIN _System_ConnectorEvent e ON e.Id = w.EventId
 WHERE ` + workItemExecutionInputPredicate("w")
	if len(lanes) == 1 {
		return query + " ORDER BY " + laneReadyHeadOrderBy("w", "e") + " LIMIT 1", args
	}
	return "SELECT WorkId FROM (" + query + ") ranked WHERE lane_position=1", args
}

func (r *Repository) nextLaneReadyHeads(ctx context.Context, in CreateWorkItemInput, lanes []string, now time.Time) ([]string, error) {
	query, args := laneReadyBatchQuery(in, lanes, now)
	rows, err := r.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var heads []string
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		heads = append(heads, id)
	}
	sort.Strings(heads)
	return heads, rows.Err()
}

func (r *Repository) demoteNonHeadLaneReadyBatch(ctx context.Context, in CreateWorkItemInput, lanes, heads []string, now time.Time) error {
	args := []any{WorkStatePending, WorkStateBlocked, WorkStateRetry, WorkStateBlockedRetry, now, in.InstalledPackID, in.Generation}
	args = appendStrings(args, lanes)
	args = append(args, WorkStatePending, WorkStateRetry, now)
	filter := ""
	if len(heads) > 0 {
		filter = " AND WorkId NOT IN " + inSQL(heads)
		if len(heads) == 1 {
			filter = " AND WorkId <> ?"
		}
		args = appendStrings(args, heads)
	}
	for _, branch := range []struct{ index, predicate string }{
		{"idx_connector_work_item_target_lane", equalsOrIn("TargetMutationLaneKey", lanes)},
		{"idx_connector_work_item_source_lane_gen", equalsOrIn("SourceLaneKey", lanes) + " AND (TargetMutationLaneKey IS NULL OR TargetMutationLaneKey = '')"},
	} {
		_, err := r.db.ExecContext(ctx, `UPDATE _System_ConnectorWorkItem USE INDEX (`+branch.index+`)
 SET State = CASE WHEN State = ? THEN ? WHEN State = ? THEN ? ELSE State END, LastModifiedDate = ?
 WHERE InstalledPackId = ? AND Generation = ? AND `+branch.predicate+`
 AND State IN (?, ?) AND RunAfter <= ?`+filter, args...)
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *Repository) promoteLaneReadyBatch(ctx context.Context, in CreateWorkItemInput, lanes, heads []string, now time.Time) error {
	args := appendStrings([]any{WorkStateBlocked, WorkStatePending, WorkStateBlockedRetry, WorkStateRetry, now}, heads)
	args = append(args, WorkStateBlocked, WorkStateBlockedRetry)
	// Only outer expressions normalize identity. Inner indexed columns remain
	// bare equalities, exactly as in single-lane processing guards.
	pack, generation, lane := "?", "?", "?"
	if len(lanes) > 1 {
		pack, generation = "w.InstalledPackId", "w.Generation"
		lane = "COALESCE(NULLIF(w.TargetMutationLaneKey,''),w.SourceLaneKey,'')"
	}
	for range 2 {
		if len(lanes) == 1 {
			args = append(args, in.InstalledPackID, in.Generation, lanes[0])
		}
		args = append(args, WorkStateProcessing)
	}
	query := fmt.Sprintf(`UPDATE _System_ConnectorWorkItem AS w
 SET State = CASE WHEN State = ? THEN ? WHEN State = ? THEN ? ELSE State END, LastModifiedDate = ?
 WHERE %s AND State IN (?, ?) AND %s
 AND NOT EXISTS (%s) AND NOT EXISTS (%s)`, equalsOrIn("WorkId", heads), workItemExecutionInputPredicate("w"), fmt.Sprintf(laneProcessingTargetSQL, pack, generation, lane), fmt.Sprintf(laneProcessingSourceSQL, pack, generation, lane))
	_, err := r.db.ExecContext(ctx, query, args...)
	return err
}