network/pkg/client/database_client.go

package client

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"github.com/rqlite/gorqlite"
)

// DatabaseClientImpl implements DatabaseClient
type DatabaseClientImpl struct {
	client     *Client
	connection *gorqlite.Connection
	mu         sync.RWMutex
}

// checkConnection verifies the client is connected
func (d *DatabaseClientImpl) checkConnection() error {
	if !d.client.isConnected() {
		return fmt.Errorf("client not connected")
	}
	return nil
}

// withRetry executes an operation with retry logic
func (d *DatabaseClientImpl) withRetry(operation func(*gorqlite.Connection) error) error {
	maxRetries := 3
	var lastErr error

	for attempt := 1; attempt <= maxRetries; attempt++ {
		conn, err := d.getRQLiteConnection()
		if err != nil {
			lastErr = err
			d.clearConnection()
			continue
		}

		if err := operation(conn); err != nil {
			lastErr = err
			d.clearConnection()
			continue
		}

		return nil
	}

	return fmt.Errorf("operation failed after %d attempts. Last error: %w", maxRetries, lastErr)
}

// Query executes a SQL query
func (d *DatabaseClientImpl) Query(ctx context.Context, sql string, args ...interface{}) (*QueryResult, error) {
	if err := d.checkConnection(); err != nil {
		return nil, err
	}

	if err := d.client.requireAccess(ctx); err != nil {
		return nil, fmt.Errorf("authentication required: %w - run CLI commands to authenticate automatically", err)
	}

	// Determine if this is a read or write operation
	isWriteOperation := d.isWriteOperation(sql)

	// Retry logic for resilient querying
	maxRetries := 3
	var lastErr error

	for attempt := 1; attempt <= maxRetries; attempt++ {
		// Get RQLite connection (tries multiple nodes)
		conn, err := d.getRQLiteConnection()
		if err != nil {
			lastErr = err

			// Clear any cached connection and try again
			d.clearConnection()
			continue
		}

		if isWriteOperation {
			// Execute write operation with parameters
			_, err := conn.WriteOneParameterized(gorqlite.ParameterizedStatement{
				Query:     sql,
				Arguments: args,
			})
			if err != nil {
				lastErr = err
				d.clearConnection()
				continue
			}

			// For write operations, return empty result set
			return &QueryResult{
				Columns: []string{"affected"},
				Rows:    [][]interface{}{{"success"}},
				Count:   1,
			}, nil
		} else {
			// Execute read operation with parameters
			result, err := conn.QueryOneParameterized(gorqlite.ParameterizedStatement{
				Query:     sql,
				Arguments: args,
			})
			if err != nil {
				lastErr = err
				d.clearConnection()
				continue
			}

			// Convert gorqlite.QueryResult to our QueryResult
			columns := result.Columns()
			numRows := int(result.NumRows())

			queryResult := &QueryResult{
				Columns: columns,
				Rows:    make([][]interface{}, 0, numRows),
				Count:   result.NumRows(),
			}

			// Iterate through rows
			for result.Next() {
				row, err := result.Slice()
				if err != nil {
					continue
				}
				queryResult.Rows = append(queryResult.Rows, row)
			}

			return queryResult, nil
		}
	}

	return nil, fmt.Errorf("query failed after %d attempts. Last error: %w", maxRetries, lastErr)
}

// isWriteOperation determines if a SQL statement is a write operation
func (d *DatabaseClientImpl) isWriteOperation(sql string) bool {
	// Convert to uppercase for comparison
	sqlUpper := strings.ToUpper(strings.TrimSpace(sql))

	// List of write operation keywords
	writeKeywords := []string{
		"INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER",
		"TRUNCATE", "REPLACE", "MERGE", "PRAGMA",
	}

	for _, keyword := range writeKeywords {
		if strings.HasPrefix(sqlUpper, keyword) {
			return true
		}
	}

	return false
}

// clearConnection clears the cached connection to force reconnection
func (d *DatabaseClientImpl) clearConnection() {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.connection != nil {
		d.connection.Close()
		d.connection = nil
	}
}

// getRQLiteConnection returns a connection to RQLite, creating one if needed
func (d *DatabaseClientImpl) getRQLiteConnection() (*gorqlite.Connection, error) {
	d.mu.RLock()
	conn := d.connection
	d.mu.RUnlock()

	if conn != nil {
		return conn, nil
	}

	newConn, err := d.connectToAvailableNode()
	if err != nil {
		return nil, err
	}

	d.mu.Lock()
	d.connection = newConn
	d.mu.Unlock()
	return newConn, nil
}

// getRQLiteNodes returns a list of RQLite node URLs with precedence:
// 1) client config DatabaseEndpoints
// 2) RQLITE_NODES env (comma/space separated)
// 3) library defaults via DefaultDatabaseEndpoints()
func (d *DatabaseClientImpl) getRQLiteNodes() []string {
	// 1) Prefer explicit configuration on the client
	if d.client != nil && d.client.config != nil && len(d.client.config.DatabaseEndpoints) > 0 {
		return dedupeStrings(normalizeEndpoints(d.client.config.DatabaseEndpoints))
	}

	// 3) Fallback to library defaults derived from bootstrap peers
	return DefaultDatabaseEndpoints()
}

// hasPort checks if a hostport string has a port suffix
func hasPort(hostport string) bool {
	// cheap check for :port suffix (IPv6 with brackets handled by url.Parse earlier)
	if i := strings.LastIndex(hostport, ":"); i > -1 && i < len(hostport)-1 {
		// ensure the segment after ':' is numeric-ish
		for _, c := range hostport[i+1:] {
			if c < '0' || c > '9' {
				return false
			}
		}
		return true
	}
	return false
}

// connectToAvailableNode tries to connect to any available RQLite node
func (d *DatabaseClientImpl) connectToAvailableNode() (*gorqlite.Connection, error) {
	// Get RQLite nodes from environment or use defaults
	rqliteNodes := d.getRQLiteNodes()

	var lastErr error

	for _, rqliteURL := range rqliteNodes {
		var conn *gorqlite.Connection
		var err error

		conn, err = gorqlite.Open(rqliteURL)
		if err != nil {
			lastErr = err
			continue
		}

		// Test the connection with a simple query to ensure it's working
		// and the node has leadership or can serve reads
		if err := d.testConnection(conn); err != nil {
			lastErr = err
			continue
		}

		return conn, nil
	}

	return nil, fmt.Errorf("failed to connect to any RQLite instance. Last error: %w", lastErr)
}

// testConnection performs a health check on the RQLite connection
func (d *DatabaseClientImpl) testConnection(conn *gorqlite.Connection) error {
	// Try a simple read query first (works even without leadership)
	result, err := conn.QueryOne("SELECT 1")
	if err != nil {
		return fmt.Errorf("read test failed: %w", err)
	}

	// Check if we got a valid result
	if result.NumRows() == 0 {
		return fmt.Errorf("read test returned no rows")
	}

	return nil
}

// Transaction executes multiple queries in a transaction
func (d *DatabaseClientImpl) Transaction(ctx context.Context, queries []string) error {
	if !d.client.isConnected() {
		return fmt.Errorf("client not connected")
	}

	if err := d.client.requireAccess(ctx); err != nil {
		return fmt.Errorf("authentication required: %w - run CLI commands to authenticate automatically", err)
	}

	maxRetries := 3
	var lastErr error

	for attempt := 1; attempt <= maxRetries; attempt++ {
		// Get RQLite connection
		conn, err := d.getRQLiteConnection()
		if err != nil {
			lastErr = err
			d.clearConnection()
			continue
		}

		// Execute all queries in the transaction
		success := true
		for _, query := range queries {
			_, err := conn.WriteOne(query)
			if err != nil {
				lastErr = err
				success = false
				d.clearConnection()
				break
			}
		}

		if success {
			return nil
		}
	}

	return fmt.Errorf("transaction failed after %d attempts. Last error: %w", maxRetries, lastErr)
}

// CreateTable creates a new table
func (d *DatabaseClientImpl) CreateTable(ctx context.Context, schema string) error {
	if err := d.checkConnection(); err != nil {
		return err
	}

	if err := d.client.requireAccess(ctx); err != nil {
		return fmt.Errorf("authentication required: %w - run CLI commands to authenticate automatically", err)
	}

	return d.withRetry(func(conn *gorqlite.Connection) error {
		_, err := conn.WriteOne(schema)
		return err
	})
}

// DropTable drops a table
func (d *DatabaseClientImpl) DropTable(ctx context.Context, tableName string) error {
	if err := d.checkConnection(); err != nil {
		return err
	}

	return d.withRetry(func(conn *gorqlite.Connection) error {
		dropSQL := fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)
		_, err := conn.WriteOne(dropSQL)
		return err
	})
}

// GetSchema returns schema information
func (d *DatabaseClientImpl) GetSchema(ctx context.Context) (*SchemaInfo, error) {
	if !d.client.isConnected() {
		return nil, fmt.Errorf("client not connected")
	}

	if err := d.client.requireAccess(ctx); err != nil {
		return nil, fmt.Errorf("authentication required: %w - run CLI commands to authenticate automatically", err)
	}

	// Get RQLite connection
	conn, err := d.getRQLiteConnection()
	if err != nil {
		return nil, fmt.Errorf("failed to get RQLite connection: %w", err)
	}

	// Query for all tables
	result, err := conn.QueryOne("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
	if err != nil {
		return nil, fmt.Errorf("failed to query table list: %w", err)
	}

	schema := &SchemaInfo{
		Tables: make([]TableInfo, 0),
	}

	// Iterate through tables
	for result.Next() {
		row, err := result.Slice()
		if err != nil {
			return nil, fmt.Errorf("failed to get table row: %w", err)
		}

		if len(row) > 0 {
			tableName := fmt.Sprintf("%v", row[0])

			// Get column information for this table
			columnResult, err := conn.QueryOne(fmt.Sprintf("PRAGMA table_info(%s)", tableName))
			if err != nil {
				continue // Skip this table if we can't get column info
			}

			tableInfo := TableInfo{
				Name:    tableName,
				Columns: make([]ColumnInfo, 0),
			}

			// Parse column information
			for columnResult.Next() {
				colRow, err := columnResult.Slice()
				if err != nil {
					continue
				}

				if len(colRow) >= 6 {
					columnInfo := ColumnInfo{
						Name:     fmt.Sprintf("%v", colRow[1]),        // name
						Type:     fmt.Sprintf("%v", colRow[2]),        // type
						Nullable: fmt.Sprintf("%v", colRow[3]) == "0", // notnull (0 = nullable, 1 = not null)
					}
					tableInfo.Columns = append(tableInfo.Columns, columnInfo)
				}
			}

			schema.Tables = append(schema.Tables, tableInfo)
		}
	}

	return schema, nil
}