package executor

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

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/cqlerrors"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/parser"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
)

func (e *dmlExecutor) executeCollectionMutationUpdate(ctx context.Context, exec sqlExecutor, user, password string, stmt *parser.UpdateStatement, params []interface{}) (*QueryResult, bool, error) {
	if stmt == nil {
		return nil, false, nil
	}
	assignments, hasMutation, err := e.collectAssignmentInfoForUpdate(ctx, user, password, stmt)
	if err != nil {
		return nil, true, cqlerrors.Invalid(err.Error())
	}
	if !hasMutation {
		return nil, false, nil
	}
	if stmt.TTL != nil {
		return nil, true, cqlerrors.Invalid("USING TTL is not supported")
	}
	if stmt.IfExists || stmt.IfCondition != nil {
		return nil, true, cqlerrors.Invalid("IF conditions are not supported for collection mutations")
	}
	markerIndex := buildBindMarkerIndex(stmt)
	if len(params) > 0 && hasUnset(params) {
		if err := rejectUnsetInConditions(stmt.Where, stmt.IfCondition, markerIndex, params); err != nil {
			return nil, true, err
		}
	}

	pkCols, err := e.primaryKeyColumns(ctx, user, password, stmt.Table)
	if err != nil {
		return nil, true, err
	}
	whereConditions, err := extractEqualityConditions(stmt.Where)
	if err != nil {
		return nil, true, cqlerrors.Invalid(err.Error())
	}
	if err := ensurePrimaryKeyConditions(pkCols, whereConditions); err != nil {
		return nil, true, cqlerrors.Invalid(err.Error())
	}
	for _, assignment := range assignments {
		if isColumnInList(assignment.column, pkCols) {
			return nil, true, cqlerrors.Invalid(fmt.Sprintf("UPDATE cannot modify primary key column %q", assignment.column))
		}
	}

	columnTypes, err := e.lookupColumnTypesIfAvailable(ctx, user, password, stmt.Table)
	if err != nil {
		return nil, true, err
	}

	selectCols := uniqueMutationColumns(assignments)
	selectSQL, whereOrder := e.buildSelectForUpdateSQL(stmt.Table, selectCols, whereConditions, markerIndex, columnTypes)
	whereArgs, err := mapParams(params, whereOrder)
	if err != nil {
		return nil, true, err
	}

	tx, owns, err := beginMutationTx(ctx, exec)
	if err != nil {
		return nil, true, err
	}
	committed := false
	if owns {
		defer func() {
			if !committed {
				_ = tx.Rollback()
			}
		}()
	}

	existing, rowExists, err := fetchExistingCollectionValues(ctx, tx, selectSQL, whereArgs, selectCols)
	if err != nil {
		return nil, true, e.mapMySQLErrorWithSQL(ctx, err, mysqlOpWrite, selectSQL)
	}
	updates, err := applyCollectionAssignments(assignments, existing, params, markerIndex)
	if err != nil {
		return nil, true, cqlerrors.Invalid(err.Error())
	}
	updateCols := uniqueAssignmentColumns(assignments)

	var result sql.Result
	if rowExists {
		updateSQL, updateArgs := e.buildUpdateWithValuesSQL(stmt.Table, updateCols, whereConditions, markerIndex, updates, columnTypes)
		updateArgs = append(updateArgs, whereArgs...)
		metrics.RecordRequestSQL(ctx, updateSQL, updateArgs...)
		start := time.Now()
		result, err = tx.ExecContext(ctx, updateSQL, updateArgs...)
		d := time.Since(start)
		metrics.ObserveStageFromContext(ctx, metrics.StageMySQLQuery, d)
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, err == nil, d)
		if err != nil {
			return nil, true, e.mapMySQLErrorWithSQL(ctx, fmt.Errorf("failed to execute update: %w", err), mysqlOpWrite, updateSQL)
		}
	} else {
		pkValues, err := pkValuesFromConditions(pkCols, whereConditions, params, markerIndex)
		if err != nil {
			return nil, true, cqlerrors.Invalid(err.Error())
		}
		insertSQL, insertArgs := e.buildInsertWithValuesSQL(stmt.Table, pkCols, pkValues, updateCols, updates)
		metrics.RecordRequestSQL(ctx, insertSQL, insertArgs...)
		start := time.Now()
		result, err = tx.ExecContext(ctx, insertSQL, insertArgs...)
		d := time.Since(start)
		metrics.ObserveStageFromContext(ctx, metrics.StageMySQLQuery, d)
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, err == nil, d)
		if err != nil {
			return nil, true, e.mapMySQLErrorWithSQL(ctx, fmt.Errorf("failed to execute insert: %w", err), mysqlOpWrite, insertSQL)
		}
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return nil, true, fmt.Errorf("failed to get rows affected: %w", err)
	}

	if owns {
		if err := tx.Commit(); err != nil {
			return nil, true, fmt.Errorf("failed to commit transaction: %w", err)
		}
		committed = true
	}

	return &QueryResult{RowsAffected: rowsAffected, CQLWriteRows: rowsAffected}, true, nil
}

func beginMutationTx(ctx context.Context, exec sqlExecutor) (*sql.Tx, bool, error) {
	if tx, ok := exec.(*sql.Tx); ok {
		return tx, false, nil
	}
	conn, ok := exec.(*sql.Conn)
	if !ok {
		return nil, false, fmt.Errorf("executor does not support transactions")
	}
	start := time.Now()
	tx, err := conn.BeginTx(ctx, nil)
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpBeginTx, err == nil, time.Since(start))
	return tx, true, err
}

func (e *dmlExecutor) buildSelectForUpdateSQL(table *parser.TableName, columns []string, conditions []equalityCondition, markerIndex map[parser.Expression]int, columnTypes map[string]string) (string, []int) {
	var sql strings.Builder
	sql.WriteString("SELECT ")
	for i, col := range columns {
		if i > 0 {
			sql.WriteString(", ")
		}
		sql.WriteString(e.translator.QuoteIdentifier(col))
	}
	sql.WriteString(" FROM ")
	sql.WriteString(e.translator.QuoteIdentifier(table.Schema))
	sql.WriteString(".")
	sql.WriteString(e.translator.QuoteIdentifier(table.Name))
	whereSQL, order := e.buildWhereClause(conditions, markerIndex, columnTypes)
	sql.WriteString(" WHERE ")
	sql.WriteString(whereSQL)
	sql.WriteString(" FOR UPDATE")
	return sql.String(), order
}

func fetchExistingCollectionValues(ctx context.Context, tx *sql.Tx, sqlQuery string, args []interface{}, columns []string) (map[string]interface{}, bool, error) {
	metrics.RecordRequestSQL(ctx, sqlQuery, args...)
	start := time.Now()
	rows, err := tx.QueryContext(ctx, sqlQuery, args...)
	d := time.Since(start)
	metrics.ObserveStageFromContext(ctx, metrics.StageMySQLQuery, d)
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryWrite, err == nil, d)
	if err != nil {
		return nil, false, err
	}
	defer rows.Close()
	if !rows.Next() {
		if err := rows.Err(); err != nil {
			return nil, false, err
		}
		return map[string]interface{}{}, false, nil
	}
	values := make([]interface{}, len(columns))
	scanArgs := make([]interface{}, len(columns))
	for i := range values {
		scanArgs[i] = &values[i]
	}
	if err := rows.Scan(scanArgs...); err != nil {
		return nil, false, err
	}
	if err := rows.Err(); err != nil {
		return nil, false, err
	}
	existing := make(map[string]interface{}, len(columns))
	for i, col := range columns {
		decoded, err := decodeExistingCollectionValue(values[i])
		if err != nil {
			return nil, true, err
		}
		existing[col] = decoded
	}
	return existing, true, nil
}

func decodeExistingCollectionValue(value interface{}) (interface{}, error) {
	if value == nil {
		return nil, nil
	}
	switch v := value.(type) {
	case []byte:
		return decodeJSONValue(v)
	case string:
		return decodeJSONValue([]byte(v))
	default:
		return value, nil
	}
}

func pkValuesFromConditions(pkCols []string, conditions []equalityCondition, params []interface{}, markerIndex map[parser.Expression]int) ([]interface{}, error) {
	condMap := make(map[string]parser.Expression, len(conditions))
	for _, cond := range conditions {
		condMap[strings.ToLower(cond.column)] = cond.value
	}
	values := make([]interface{}, len(pkCols))
	for i, pk := range pkCols {
		expr, ok := condMap[strings.ToLower(pk)]
		if !ok {
			return nil, fmt.Errorf("missing primary key column %q in WHERE clause", pk)
		}
		val, err := resolveScalarExpr(expr, params, markerIndex)
		if err != nil {
			return nil, err
		}
		values[i] = val
	}
	return values, nil
}

func uniqueMutationColumns(assignments []assignmentInfo) []string {
	seen := make(map[string]struct{}, len(assignments))
	out := make([]string, 0, len(assignments))
	for _, info := range assignments {
		if info.mutation == nil {
			continue
		}
		key := strings.ToLower(info.column)
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		out = append(out, info.column)
	}
	return out
}

func uniqueAssignmentColumns(assignments []assignmentInfo) []string {
	seen := make(map[string]struct{}, len(assignments))
	out := make([]string, 0, len(assignments))
	for _, info := range assignments {
		key := strings.ToLower(info.column)
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		out = append(out, info.column)
	}
	return out
}

func (e *dmlExecutor) buildUpdateWithValuesSQL(table *parser.TableName, columns []string, conditions []equalityCondition, markerIndex map[parser.Expression]int, values map[string]interface{}, columnTypes map[string]string) (string, []interface{}) {
	var sql strings.Builder
	sql.WriteString("UPDATE ")
	sql.WriteString(e.translator.QuoteIdentifier(table.Schema))
	sql.WriteString(".")
	sql.WriteString(e.translator.QuoteIdentifier(table.Name))
	sql.WriteString(" SET ")
	args := make([]interface{}, 0, len(columns))
	for i, col := range columns {
		if i > 0 {
			sql.WriteString(", ")
		}
		sql.WriteString(e.translator.QuoteIdentifier(col))
		sql.WriteString(" = ?")
		args = append(args, values[col])
	}
	whereSQL, _ := e.buildWhereClause(conditions, markerIndex, columnTypes)
	sql.WriteString(" WHERE ")
	sql.WriteString(whereSQL)
	return sql.String(), args
}

func (e *dmlExecutor) buildInsertWithValuesSQL(table *parser.TableName, pkCols []string, pkValues []interface{}, columns []string, values map[string]interface{}) (string, []interface{}) {
	var sql strings.Builder
	sql.WriteString("INSERT INTO ")
	sql.WriteString(e.translator.QuoteIdentifier(table.Schema))
	sql.WriteString(".")
	sql.WriteString(e.translator.QuoteIdentifier(table.Name))

	insertCols := make([]string, 0, len(pkCols)+len(columns))
	args := make([]interface{}, 0, len(pkCols)+len(columns))
	for i, pk := range pkCols {
		insertCols = append(insertCols, pk)
		if i < len(pkValues) {
			args = append(args, pkValues[i])
		} else {
			args = append(args, nil)
		}
	}
	for _, col := range columns {
		if isColumnInList(col, pkCols) {
			continue
		}
		insertCols = append(insertCols, col)
		args = append(args, values[col])
	}

	sql.WriteString(" (")
	for i, col := range insertCols {
		if i > 0 {
			sql.WriteString(", ")
		}
		sql.WriteString(e.translator.QuoteIdentifier(col))
	}
	sql.WriteString(") VALUES (")
	for i := range insertCols {
		if i > 0 {
			sql.WriteString(", ")
		}
		sql.WriteString("?")
	}
	sql.WriteString(")")
	return sql.String(), args
}