package executor

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

	"go.uber.org/zap"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/mysqlbackend"
)

type cachedStmtExecutor struct {
	pool   *mysqlbackend.UserPool
	logger *zap.Logger
}

func newCachedStmtExecutor(pool *mysqlbackend.UserPool, logger *zap.Logger) *cachedStmtExecutor {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &cachedStmtExecutor{pool: pool, logger: logger}
}

func (e *cachedStmtExecutor) ExecContext(ctx context.Context, sqlText string, args ...interface{}) (sql.Result, error) {
	if e == nil || e.pool == nil {
		return nil, fmt.Errorf("prepared statement cache not configured")
	}
	if sqlText == "" {
		return nil, fmt.Errorf("sql is required")
	}
	start := time.Now()
	stmt, err := e.pool.PrepareCached(ctx, sqlText)
	if err != nil {
		e.pool.RecordPreparedError()
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, false, time.Since(start))
		return nil, err
	}
	result, err := stmt.ExecContext(ctx, args...)
	if err == nil {
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, true, time.Since(start))
		return result, nil
	}
	e.pool.RecordPreparedError()
	if mysqlbackend.IsPreparedStmtStale(err) {
		e.pool.EvictPrepared(sqlText)
		stmt, prepErr := e.pool.PrepareCached(ctx, sqlText)
		if prepErr != nil {
			e.pool.RecordPreparedError()
			metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, false, time.Since(start))
			return result, err
		}
		result, retryErr := stmt.ExecContext(ctx, args...)
		if retryErr != nil {
			e.pool.RecordPreparedError()
			metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, false, time.Since(start))
			return result, retryErr
		}
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, true, time.Since(start))
		return result, nil
	}
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpExecWrite, false, time.Since(start))
	return result, err
}

func (e *cachedStmtExecutor) QueryContext(ctx context.Context, sqlText string, args ...interface{}) (*sql.Rows, error) {
	if e == nil || e.pool == nil {
		return nil, fmt.Errorf("prepared statement cache not configured")
	}
	if sqlText == "" {
		return nil, fmt.Errorf("sql is required")
	}
	start := time.Now()
	stmt, err := e.pool.PrepareCached(ctx, sqlText)
	if err != nil {
		e.pool.RecordPreparedError()
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryRead, false, time.Since(start))
		return nil, err
	}
	rows, err := stmt.QueryContext(ctx, args...)
	if err == nil {
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryRead, true, time.Since(start))
		return rows, nil
	}
	e.pool.RecordPreparedError()
	if mysqlbackend.IsPreparedStmtStale(err) {
		e.pool.EvictPrepared(sqlText)
		stmt, prepErr := e.pool.PrepareCached(ctx, sqlText)
		if prepErr != nil {
			e.pool.RecordPreparedError()
			metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryRead, false, time.Since(start))
			return nil, err
		}
		rows, retryErr := stmt.QueryContext(ctx, args...)
		if retryErr != nil {
			e.pool.RecordPreparedError()
			metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryRead, false, time.Since(start))
			return nil, retryErr
		}
		metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryRead, true, time.Since(start))
		return rows, nil
	}
	metrics.ObserveMySQLOperation(ctx, metrics.MySQLOpQueryRead, false, time.Since(start))
	return nil, err
}