pizzasql/pkg/storage/schema.go

package storage

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/goccy/go-json"

	"github.com/danfragoso/pizzasql-next/pkg/analyzer"
)

// Schema represents a table schema.
type Schema struct {
	Name          string    `json:"name"`
	Columns       []Column  `json:"columns"`
	PrimaryKey    string    `json:"primary_key"`
	CreatedAt     time.Time `json:"created_at"`
	NextRowID     int64     `json:"next_rowid"`
	AutoIncrement bool      `json:"autoincrement"`
}

// Column represents a column definition.
type Column struct {
	Name       string      `json:"name"`
	Type       string      `json:"type"`
	Nullable   bool        `json:"nullable"`
	Default    interface{} `json:"default,omitempty"`
	PrimaryKey bool        `json:"primary_key"`
}

// Index represents an index definition.
type Index struct {
	Name      string        `json:"name"`
	Table     string        `json:"table"`
	Columns   []IndexColumn `json:"columns"`
	Unique    bool          `json:"unique"`
	CreatedAt time.Time     `json:"created_at"`
}

// IndexColumn represents a column in an index.
type IndexColumn struct {
	Name string `json:"name"`
	Desc bool   `json:"desc"`
}

// SchemaManager manages table schemas.
type SchemaManager struct {
	pool             *KVPool
	database         string
	cache            map[string]*Schema
	rowIDInitialized map[string]bool
	version          uint64
	mu               sync.RWMutex
}

// NewSchemaManager creates a new schema manager.
func NewSchemaManager(pool *KVPool, database string) *SchemaManager {
	return &SchemaManager{
		pool:             pool,
		database:         database,
		cache:            make(map[string]*Schema),
		rowIDInitialized: make(map[string]bool),
	}
}

// GetDatabaseName returns the database name.
func (m *SchemaManager) GetDatabaseName() string {
	return m.database
}

// GetPool returns the KV pool.
func (m *SchemaManager) GetPool() *KVPool {
	return m.pool
}

// Version returns the in-process schema catalog version. It is incremented for
// schema/index definition changes so cached executors can resync their analyzer
// catalogs without scanning storage on every query.
func (m *SchemaManager) Version() uint64 {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.version
}

func (m *SchemaManager) bumpVersionLocked() {
	m.version++
}

// schemaKey returns the key for a table schema.
func (m *SchemaManager) schemaKey(table string) string {
	return fmt.Sprintf("%s:_schema:%s", m.database, strings.ToLower(table))
}

// catalogKey returns the key for the table catalog.
func (m *SchemaManager) catalogKey() string {
	return fmt.Sprintf("%s:_sys:tables", m.database)
}

// rowIDKey returns the key for a table's next ROWID counter.
func (m *SchemaManager) rowIDKey(table string) string {
	return fmt.Sprintf("%s:_sys:rowid:%s", m.database, strings.ToLower(table))
}

// CreateTable creates a new table.
func (m *SchemaManager) CreateTable(schema *Schema) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Check if table already exists
	key := m.schemaKey(schema.Name)
	err := m.pool.WithClient(func(c *KVClient) error {
		_, err := c.Read(key)
		return err
	})
	if err == nil {
		return fmt.Errorf("table already exists: %s", schema.Name)
	}

	// Set creation time
	schema.CreatedAt = time.Now()

	// Determine primary key if not set
	if schema.PrimaryKey == "" {
		for _, col := range schema.Columns {
			if col.PrimaryKey {
				schema.PrimaryKey = col.Name
				break
			}
		}
		// No explicit primary key declared — use synthetic _rowid_ so user
		// columns remain unconstrained and can hold duplicate or NULL values.
		if schema.PrimaryKey == "" {
			schema.PrimaryKey = "_rowid_"
		}
	}

	// Serialize schema
	data, err := json.Marshal(schema)
	if err != nil {
		return fmt.Errorf("failed to serialize schema: %w", err)
	}

	// Write schema
	err = m.pool.WithClient(func(c *KVClient) error {
		return c.Write(key, string(data))
	})
	if err != nil {
		return fmt.Errorf("failed to write schema: %w", err)
	}

	// Update catalog
	if err := m.addToCatalog(schema.Name); err != nil {
		// Rollback schema write
		m.pool.WithClient(func(c *KVClient) error {
			return c.Delete(key)
		})
		return err
	}

	// Update cache
	m.cache[strings.ToLower(schema.Name)] = schema
	m.bumpVersionLocked()

	return nil
}

// DropTable drops a table.
func (m *SchemaManager) DropTable(name string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := m.schemaKey(name)

	// Check if table exists
	err := m.pool.WithClient(func(c *KVClient) error {
		_, err := c.Read(key)
		return err
	})
	if err != nil {
		return fmt.Errorf("table not found: %s", name)
	}

	// Delete all rows
	dataPrefix := fmt.Sprintf("%s:_data:%s:", m.database, strings.ToLower(name))
	err = m.pool.WithClient(func(c *KVClient) error {
		// Get all keys with this prefix and delete them
		// Note: This is a simplified version - in production you'd want batch delete
		values, err := c.Reads(dataPrefix)
		if err != nil {
			return err
		}
		// The Reads command returns values, not keys, so we can't delete them directly
		// In a real implementation, we'd need a keys scan command
		_ = values
		return nil
	})

	// Delete schema
	err = m.pool.WithClient(func(c *KVClient) error {
		return c.Delete(key)
	})
	if err != nil {
		return fmt.Errorf("failed to delete schema: %w", err)
	}

	// Delete ROWID state.
	m.pool.WithClient(func(c *KVClient) error {
		return c.Delete(m.rowIDKey(name))
	})

	// Update catalog
	if err := m.removeFromCatalog(name); err != nil {
		return err
	}

	// Update cache
	tableLower := strings.ToLower(name)
	delete(m.cache, tableLower)
	delete(m.rowIDInitialized, tableLower)
	m.bumpVersionLocked()

	return nil
}

// GetSchema retrieves a table schema.
func (m *SchemaManager) GetSchema(name string) (*Schema, error) {
	m.mu.RLock()
	if schema, ok := m.cache[strings.ToLower(name)]; ok {
		m.mu.RUnlock()
		return schema, nil
	}
	m.mu.RUnlock()

	m.mu.Lock()
	defer m.mu.Unlock()

	// Double-check after acquiring write lock
	if schema, ok := m.cache[strings.ToLower(name)]; ok {
		return schema, nil
	}

	key := m.schemaKey(name)
	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err != nil {
		if err == ErrKeyNotFound {
			return nil, fmt.Errorf("table not found: %s", name)
		}
		return nil, err
	}

	var schema Schema
	if err := json.Unmarshal([]byte(data), &schema); err != nil {
		return nil, fmt.Errorf("failed to parse schema: %w", err)
	}

	m.cache[strings.ToLower(name)] = &schema
	return &schema, nil
}

// TableExists checks if a table exists.
func (m *SchemaManager) TableExists(name string) bool {
	_, err := m.GetSchema(name)
	return err == nil
}

// ListTables returns all table names.
func (m *SchemaManager) ListTables() ([]string, error) {
	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(m.catalogKey())
		return err
	})
	if err != nil {
		if err == ErrKeyNotFound {
			return nil, nil
		}
		return nil, err
	}

	var tables []string
	if err := json.Unmarshal([]byte(data), &tables); err != nil {
		return nil, fmt.Errorf("failed to parse catalog: %w", err)
	}

	return tables, nil
}

// addToCatalog adds a table to the catalog.
func (m *SchemaManager) addToCatalog(name string) error {
	tables, err := m.ListTables()
	if err != nil && err != ErrKeyNotFound {
		return err
	}

	// Check if already exists
	lowerName := strings.ToLower(name)
	for _, t := range tables {
		if strings.ToLower(t) == lowerName {
			return nil
		}
	}

	tables = append(tables, name)
	data, err := json.Marshal(tables)
	if err != nil {
		return err
	}

	return m.pool.WithClient(func(c *KVClient) error {
		return c.Write(m.catalogKey(), string(data))
	})
}

// removeFromCatalog removes a table from the catalog.
func (m *SchemaManager) removeFromCatalog(name string) error {
	tables, err := m.ListTables()
	if err != nil {
		return err
	}

	lowerName := strings.ToLower(name)
	newTables := make([]string, 0, len(tables))
	for _, t := range tables {
		if strings.ToLower(t) != lowerName {
			newTables = append(newTables, t)
		}
	}

	data, err := json.Marshal(newTables)
	if err != nil {
		return err
	}

	return m.pool.WithClient(func(c *KVClient) error {
		return c.Write(m.catalogKey(), string(data))
	})
}

// InvalidateCache clears the cache for a table.
func (m *SchemaManager) InvalidateCache(name string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	tableLower := strings.ToLower(name)
	delete(m.cache, tableLower)
	delete(m.rowIDInitialized, tableLower)
}

// ToAnalyzerTableInfo converts a Schema to analyzer.TableInfo.
func (s *Schema) ToAnalyzerTableInfo() *analyzer.TableInfo {
	info := &analyzer.TableInfo{
		Name: s.Name,
	}

	for _, col := range s.Columns {
		info.Columns = append(info.Columns, analyzer.ColumnInfo{
			Name:       col.Name,
			Type:       analyzer.TypeFromName(col.Type),
			Nullable:   col.Nullable,
			PrimaryKey: col.PrimaryKey,
			TableName:  s.Name,
		})
	}

	return info
}

// GetColumn returns a column by name.
func (s *Schema) GetColumn(name string) (*Column, bool) {
	lowerName := strings.ToLower(name)
	for i := range s.Columns {
		if strings.ToLower(s.Columns[i].Name) == lowerName {
			return &s.Columns[i], true
		}
	}
	return nil, false
}

// GetNextRowID gets and increments the next ROWID for a table.
func (m *SchemaManager) GetNextRowID(table string) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	schema, err := m.getSchemaLocked(table)
	if err != nil {
		return 0, err
	}

	nextRowID, err := m.getNextRowIDLocked(schema)
	if err != nil {
		return 0, err
	}

	schema.NextRowID = nextRowID + 1

	return nextRowID, nil
}

// UpdateMaxRowID updates the next ROWID if the provided value is higher.
func (m *SchemaManager) UpdateMaxRowID(table string, rowid int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	schema, err := m.getSchemaLocked(table)
	if err != nil {
		return err
	}

	nextRowID, err := m.getNextRowIDLocked(schema)
	if err != nil {
		return err
	}

	if rowid >= nextRowID {
		schema.NextRowID = rowid + 1
	}

	return nil
}

// getNextRowIDLocked returns a table's in-memory ROWID counter (must hold lock).
// On first use after startup, the counter is derived from durable row data so
// ROWID movement does not add a separate WAL entry.
func (m *SchemaManager) getNextRowIDLocked(schema *Schema) (int64, error) {
	tableLower := strings.ToLower(schema.Name)
	if m.rowIDInitialized[tableLower] {
		if schema.NextRowID < 1 {
			schema.NextRowID = 1
		}
		return schema.NextRowID, nil
	}

	nextRowID, err := m.deriveNextRowIDLocked(schema)
	if err != nil {
		return 0, err
	}
	if schema.NextRowID > nextRowID {
		nextRowID = schema.NextRowID
	}
	if nextRowID < 1 {
		nextRowID = 1
	}

	schema.NextRowID = nextRowID
	m.rowIDInitialized[tableLower] = true
	return schema.NextRowID, nil
}

// deriveNextRowIDLocked scans durable row values to recover max(rowid)+1.
func (m *SchemaManager) deriveNextRowIDLocked(schema *Schema) (int64, error) {
	prefix := fmt.Sprintf("%s:_data:%s:", m.database, strings.ToLower(schema.Name))
	var values []string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		values, err = c.Reads(prefix)
		return err
	})
	if err != nil {
		return 0, err
	}

	var maxRowID int64
	for _, value := range values {
		var row Row
		if err := json.Unmarshal([]byte(value), &row); err != nil {
			return 0, fmt.Errorf("failed to parse row while deriving ROWID: %w", err)
		}

		if rowid, ok := valueAsInt64(row["_rowid_"]); ok && rowid > maxRowID {
			maxRowID = rowid
		}
	}

	return maxRowID + 1, nil
}

func valueAsInt64(value interface{}) (int64, bool) {
	switch v := value.(type) {
	case int64:
		return v, true
	case int:
		return int64(v), true
	case float64:
		return int64(v), true
	default:
		return 0, false
	}
}

// getSchemaLocked retrieves schema (must hold lock).
func (m *SchemaManager) getSchemaLocked(name string) (*Schema, error) {
	if schema, ok := m.cache[strings.ToLower(name)]; ok {
		return schema, nil
	}

	key := m.schemaKey(name)
	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err != nil {
		if err == ErrKeyNotFound {
			return nil, fmt.Errorf("table not found: %s", name)
		}
		return nil, err
	}

	var schema Schema
	if err := json.Unmarshal([]byte(data), &schema); err != nil {
		return nil, fmt.Errorf("failed to parse schema: %w", err)
	}

	m.cache[strings.ToLower(name)] = &schema
	return &schema, nil
}

// saveSchemaLocked saves schema (must hold lock).
func (m *SchemaManager) saveSchemaLocked(schema *Schema) error {
	data, err := json.Marshal(schema)
	if err != nil {
		return fmt.Errorf("failed to serialize schema: %w", err)
	}

	key := m.schemaKey(schema.Name)
	err = m.pool.WithClient(func(c *KVClient) error {
		return c.Write(key, string(data))
	})
	if err != nil {
		return fmt.Errorf("failed to write schema: %w", err)
	}

	m.cache[strings.ToLower(schema.Name)] = schema
	return nil
}

// Index management methods

// indexKey returns the key for an index.
func (m *SchemaManager) indexKey(name string) string {
	return fmt.Sprintf("%s:index:%s", m.database, strings.ToLower(name))
}

// indexListKey returns the key for the index list.
func (m *SchemaManager) indexListKey() string {
	return fmt.Sprintf("%s:indexes", m.database)
}

// CreateIndex creates a new index.
func (m *SchemaManager) CreateIndex(index *Index) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Check if index already exists
	key := m.indexKey(index.Name)
	err := m.pool.WithClient(func(c *KVClient) error {
		_, err := c.Read(key)
		return err
	})
	if err == nil {
		return fmt.Errorf("index already exists: %s", index.Name)
	}

	// Verify table exists
	if _, err := m.getSchemaLocked(index.Table); err != nil {
		return fmt.Errorf("table not found: %s", index.Table)
	}

	// Save index
	index.CreatedAt = time.Now()
	data, err := json.Marshal(index)
	if err != nil {
		return fmt.Errorf("failed to serialize index: %w", err)
	}

	err = m.pool.WithClient(func(c *KVClient) error {
		return c.Write(key, string(data))
	})
	if err != nil {
		return fmt.Errorf("failed to write index: %w", err)
	}

	// Add to index list
	if err := m.addToIndexList(index.Name); err != nil {
		return err
	}
	m.bumpVersionLocked()
	return nil
}

// DropIndex drops an index.
func (m *SchemaManager) DropIndex(name string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	key := m.indexKey(name)
	err := m.pool.WithClient(func(c *KVClient) error {
		return c.Delete(key)
	})
	if err != nil {
		return fmt.Errorf("failed to delete index: %w", err)
	}

	if err := m.removeFromIndexList(name); err != nil {
		return err
	}
	m.bumpVersionLocked()
	return nil
}

// IndexExists checks if an index exists.
func (m *SchemaManager) IndexExists(name string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()

	key := m.indexKey(name)
	err := m.pool.WithClient(func(c *KVClient) error {
		_, err := c.Read(key)
		return err
	})
	return err == nil
}

// GetIndex retrieves an index by name.
func (m *SchemaManager) GetIndex(name string) (*Index, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	key := m.indexKey(name)
	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err != nil {
		return nil, fmt.Errorf("index not found: %s", name)
	}

	var index Index
	if err := json.Unmarshal([]byte(data), &index); err != nil {
		return nil, fmt.Errorf("failed to parse index: %w", err)
	}

	return &index, nil
}

// ListIndexes returns all index names.
func (m *SchemaManager) ListIndexes() ([]string, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	key := m.indexListKey()
	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err != nil {
		return []string{}, nil
	}

	var indexes []string
	if err := json.Unmarshal([]byte(data), &indexes); err != nil {
		return []string{}, nil
	}

	return indexes, nil
}

// ListTableIndexes returns all indexes for a table.
func (m *SchemaManager) ListTableIndexes(table string) ([]*Index, error) {
	indexes, err := m.ListIndexes()
	if err != nil {
		return nil, err
	}

	var result []*Index
	for _, name := range indexes {
		idx, err := m.GetIndex(name)
		if err != nil {
			continue
		}
		if strings.EqualFold(idx.Table, table) {
			result = append(result, idx)
		}
	}

	return result, nil
}

// addToIndexList adds an index name to the list.
func (m *SchemaManager) addToIndexList(name string) error {
	key := m.indexListKey()
	var indexes []string

	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err == nil {
		json.Unmarshal([]byte(data), &indexes)
	}

	indexes = append(indexes, name)
	newData, _ := json.Marshal(indexes)

	return m.pool.WithClient(func(c *KVClient) error {
		return c.Write(key, string(newData))
	})
}

// removeFromIndexList removes an index name from the list.
func (m *SchemaManager) removeFromIndexList(name string) error {
	key := m.indexListKey()
	var indexes []string

	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err != nil {
		return nil
	}
	json.Unmarshal([]byte(data), &indexes)

	var newIndexes []string
	for _, idx := range indexes {
		if !strings.EqualFold(idx, name) {
			newIndexes = append(newIndexes, idx)
		}
	}

	newData, _ := json.Marshal(newIndexes)
	return m.pool.WithClient(func(c *KVClient) error {
		return c.Write(key, string(newData))
	})
}

// AddColumn adds a new column to a table.
func (m *SchemaManager) AddColumn(table string, column Column) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	schema, err := m.getSchemaUnsafe(table)
	if err != nil {
		return err
	}

	// Check if column already exists
	for _, col := range schema.Columns {
		if strings.EqualFold(col.Name, column.Name) {
			return fmt.Errorf("column already exists: %s", column.Name)
		}
	}

	// Add column
	schema.Columns = append(schema.Columns, column)

	// Update schema
	return m.updateSchemaUnsafe(schema)
}

// DropColumn removes a column from a table.
func (m *SchemaManager) DropColumn(table, columnName string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	schema, err := m.getSchemaUnsafe(table)
	if err != nil {
		return err
	}

	// Cannot drop primary key column
	if strings.EqualFold(schema.PrimaryKey, columnName) {
		return fmt.Errorf("cannot drop primary key column: %s", columnName)
	}

	// Find and remove column
	newColumns := make([]Column, 0, len(schema.Columns)-1)
	found := false
	for _, col := range schema.Columns {
		if strings.EqualFold(col.Name, columnName) {
			found = true
			continue
		}
		newColumns = append(newColumns, col)
	}

	if !found {
		return fmt.Errorf("column not found: %s", columnName)
	}

	schema.Columns = newColumns

	// Update schema
	return m.updateSchemaUnsafe(schema)
}

// RenameTable renames a table.
func (m *SchemaManager) RenameTable(oldName, newName string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Check if old table exists
	schema, err := m.getSchemaUnsafe(oldName)
	if err != nil {
		return err
	}

	// Check if new table name already exists
	_, err = m.getSchemaUnsafe(newName)
	if err == nil {
		return fmt.Errorf("table already exists: %s", newName)
	}

	// Update schema name
	schema.Name = newName

	rowIDKey := m.rowIDKey(oldName)

	// Delete old schema
	oldKey := m.schemaKey(oldName)
	err = m.pool.WithClient(func(c *KVClient) error {
		return c.Delete(oldKey)
	})
	if err != nil {
		return err
	}

	// Remove from catalog
	m.removeFromCatalog(oldName)

	// Move ROWID state.
	m.pool.WithClient(func(c *KVClient) error {
		return c.Delete(rowIDKey)
	})

	// Update cache
	oldLower := strings.ToLower(oldName)
	newLower := strings.ToLower(newName)
	wasInitialized := m.rowIDInitialized[oldLower]
	delete(m.cache, oldLower)
	delete(m.rowIDInitialized, oldLower)

	// Write new schema
	newKey := m.schemaKey(newName)
	data, _ := json.Marshal(schema)
	err = m.pool.WithClient(func(c *KVClient) error {
		return c.Write(newKey, string(data))
	})
	if err != nil {
		return err
	}

	// Add to catalog
	m.addToCatalog(newName)

	// Update cache
	m.cache[newLower] = schema
	if wasInitialized {
		m.rowIDInitialized[newLower] = true
	}
	m.bumpVersionLocked()

	return nil
}

// RenameColumn renames a column in a table.
func (m *SchemaManager) RenameColumn(table, oldName, newName string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	schema, err := m.getSchemaUnsafe(table)
	if err != nil {
		return err
	}

	// Check if new column name already exists
	for _, col := range schema.Columns {
		if strings.EqualFold(col.Name, newName) {
			return fmt.Errorf("column already exists: %s", newName)
		}
	}

	// Find and rename column
	found := false
	for i, col := range schema.Columns {
		if strings.EqualFold(col.Name, oldName) {
			schema.Columns[i].Name = newName
			found = true

			// Update primary key reference if needed
			if strings.EqualFold(schema.PrimaryKey, oldName) {
				schema.PrimaryKey = newName
			}
			break
		}
	}

	if !found {
		return fmt.Errorf("column not found: %s", oldName)
	}

	// Update schema
	return m.updateSchemaUnsafe(schema)
}

// getSchemaUnsafe gets a schema without locking (internal use).
func (m *SchemaManager) getSchemaUnsafe(table string) (*Schema, error) {
	tableLower := strings.ToLower(table)

	// Check cache
	if schema, ok := m.cache[tableLower]; ok {
		return schema, nil
	}

	// Read from storage
	key := m.schemaKey(table)
	var data string
	err := m.pool.WithClient(func(c *KVClient) error {
		var err error
		data, err = c.Read(key)
		return err
	})
	if err != nil {
		return nil, fmt.Errorf("table not found: %s", table)
	}

	var schema Schema
	if err := json.Unmarshal([]byte(data), &schema); err != nil {
		return nil, err
	}

	m.cache[tableLower] = &schema
	return &schema, nil
}

// updateSchemaUnsafe updates a schema without locking (internal use).
func (m *SchemaManager) updateSchemaUnsafe(schema *Schema) error {
	key := m.schemaKey(schema.Name)
	data, _ := json.Marshal(schema)

	err := m.pool.WithClient(func(c *KVClient) error {
		return c.Write(key, string(data))
	})
	if err != nil {
		return err
	}

	// Update cache
	m.cache[strings.ToLower(schema.Name)] = schema
	m.bumpVersionLocked()
	return nil
}