repository.go

package contentdata

import (
	"context"
	"database/sql"
	"encoding/base64"
	"fmt"
	"reflect"
	"strings"

	"github.com/goccy/go-json"
	"github.com/goccy/go-zetasqlite"
	"go.uber.org/zap"
	bigqueryv2 "google.golang.org/api/bigquery/v2"

	"github.com/goccy/bigquery-emulator/internal/connection"
	"github.com/goccy/bigquery-emulator/internal/logger"
	"github.com/goccy/bigquery-emulator/internal/metadata"
	internaltypes "github.com/goccy/bigquery-emulator/internal/types"
	"github.com/goccy/bigquery-emulator/types"
)

const ViewQueryEndCutset = ";\n \t"

type Repository struct {
	db *sql.DB
}

func NewRepository(db *sql.DB) *Repository {
	return &Repository{
		db: db,
	}
}

func (r *Repository) getConnection(ctx context.Context, projectID, datasetID string) (*sql.Conn, error) {
	if projectID == "" {
		return nil, fmt.Errorf("invalid projectID. projectID is empty")
	}
	conn, err := r.db.Conn(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get connection: %w", err)
	}
	if err := conn.Raw(func(c interface{}) error {
		zetasqliteConn, ok := c.(*zetasqlite.ZetaSQLiteConn)
		if !ok {
			return fmt.Errorf("failed to get ZetaSQLiteConn from %T", c)
		}
		if datasetID == "" {
			_ = zetasqliteConn.SetNamePath([]string{projectID})
		} else {
			_ = zetasqliteConn.SetNamePath([]string{projectID, datasetID})
		}
		const maxNamePath = 3 // projectID and datasetID and tableID
		zetasqliteConn.SetMaxNamePath(maxNamePath)
		return nil
	}); err != nil {
		return nil, fmt.Errorf("failed to setup connection: %w", err)
	}
	return conn, nil
}

func (r *Repository) tablePath(projectID, datasetID, tableID string) string {
	var tablePath []string
	if projectID != "" {
		tablePath = append(tablePath, projectID)
	}
	if datasetID != "" {
		tablePath = append(tablePath, datasetID)
	}
	tablePath = append(tablePath, tableID)
	return strings.Join(tablePath, ".")
}

func (r *Repository) routinePath(projectID, datasetID, routineID string) string {
	var routinePath []string
	if projectID != "" {
		routinePath = append(routinePath, projectID)
	}
	if datasetID != "" {
		routinePath = append(routinePath, datasetID)
	}
	routinePath = append(routinePath, routineID)
	return strings.Join(routinePath, ".")
}

func (r *Repository) CreateTable(ctx context.Context, tx *connection.Tx, table *bigqueryv2.Table) error {
	if err := tx.ContentRepoMode(); err != nil {
		return err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()
	ref := table.TableReference
	if ref == nil {
		return fmt.Errorf("TableReference is nil")
	}
	fields := make([]string, 0, len(table.Schema.Fields))
	for _, field := range table.Schema.Fields {
		fields = append(fields, fmt.Sprintf("`%s` %s", field.Name, r.encodeSchemaField(field)))
	}
	tablePath := r.tablePath(ref.ProjectId, ref.DatasetId, ref.TableId)
	query := fmt.Sprintf("CREATE TABLE `%s` (%s)", tablePath, strings.Join(fields, ","))
	if _, err := tx.Tx().ExecContext(ctx, query); err != nil {
		return fmt.Errorf("failed to create table %s: %w", query, err)
	}
	return nil
}

func getSchemaFromResult(result sql.Result) (*bigqueryv2.TableSchema, error) {
	changedCatalog, err := zetasqlite.ChangedCatalogFromResult(result)
	if err != nil {
		return nil, fmt.Errorf("failed to get changed catalog: %w", err)
	}
	if len(changedCatalog.Table.Added) != 1 {
		return nil, fmt.Errorf("catalog detected %d tables added; but expected one", len(changedCatalog.Table.Added))
	}
	createdTable := changedCatalog.Table.Added[0]
	fields := make([]*bigqueryv2.TableFieldSchema, 0, len(createdTable.Columns))
	for _, col := range createdTable.Columns {
		zetasqlType, err := col.Type.ToZetaSQLType()
		if err != nil {
			return nil, fmt.Errorf("failed to get zetasql type: %w", err)
		}
		fields = append(fields, types.TableFieldSchemaFromZetaSQLType(col.Name, zetasqlType))
	}

	return &bigqueryv2.TableSchema{Fields: fields}, nil
}

func (r *Repository) CreateView(ctx context.Context, tx *connection.Tx, table *bigqueryv2.Table) (*bigqueryv2.TableSchema, error) {
	if err := tx.ContentRepoMode(); err != nil {
		return nil, err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()
	ref := table.TableReference
	if ref == nil {
		return nil, fmt.Errorf("TableReference is nil")
	}
	viewDefinition := table.View
	if viewDefinition == nil {
		return nil, fmt.Errorf("ViewDefinition is nil")
	}
	tablePath := r.tablePath(ref.ProjectId, ref.DatasetId, ref.TableId)
	query := fmt.Sprintf("CREATE VIEW `%s` AS (%s)", tablePath, strings.TrimRight(viewDefinition.Query, ViewQueryEndCutset))
	result, err := tx.Tx().ExecContext(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("failed to create view %s: %w", query, err)
	}
	schema, err := getSchemaFromResult(result)
	return schema, err
}

func (r *Repository) encodeSchemaField(field *bigqueryv2.TableFieldSchema) string {
	var elem string
	if field.Type == "RECORD" {
		types := make([]string, 0, len(field.Fields))
		for _, f := range field.Fields {
			types = append(types, fmt.Sprintf("`%s` %s", f.Name, r.encodeSchemaField(f)))
		}
		elem = fmt.Sprintf("STRUCT<%s>", strings.Join(types, ","))
	} else {
		elem = types.Type(field.Type).ZetaSQLTypeKind().String()
	}
	if field.Mode == "REPEATED" {
		return fmt.Sprintf("ARRAY<%s>", elem)
	}
	return elem
}

func (r *Repository) Query(ctx context.Context, tx *connection.Tx, projectID, datasetID, query string, params []*bigqueryv2.QueryParameter) (*internaltypes.QueryResponse, error) {
	tx.SetProjectAndDataset(projectID, datasetID)
	if err := tx.ContentRepoMode(); err != nil {
		return nil, err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()

	values := []interface{}{}
	for _, param := range params {
		value, err := r.queryParameterValueToGoValue(param.ParameterValue)
		if err != nil {
			return nil, err
		}
		if param.Name != "" {
			values = append(values, sql.Named(param.Name, value))
		} else {
			values = append(values, value)
		}
	}
	fields := []*bigqueryv2.TableFieldSchema{}
	logger.Logger(ctx).Info(
		"",
		zap.String("query", query),
		zap.Any("values", values),
	)
	// We must pass the query parameters to zetasqlite so the analyzer uses the proper typings
	if err := tx.Conn().Raw(func(c interface{}) error {
		zetasqliteConn, ok := c.(*zetasqlite.ZetaSQLiteConn)
		if !ok {
			return fmt.Errorf("failed to get ZetaSQLiteConn from %T", c)
		}
		zetasqliteConn.SetQueryParameters(params)
		return nil
	}); err != nil {
		return nil, fmt.Errorf("failed to setup connection: %w", err)
	}
	rows, err := tx.Tx().QueryContext(ctx, query, values...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	changedCatalog, err := zetasqlite.ChangedCatalogFromRows(rows)
	if err != nil {
		return nil, fmt.Errorf("failed to get changed catalog: %w", err)
	}
	colNames, err := rows.Columns()
	if err != nil {
		return nil, fmt.Errorf("failed to get columns: %w", err)
	}
	columnTypes, err := rows.ColumnTypes()
	if err != nil {
		return nil, fmt.Errorf("failed to get column types: %w", err)
	}
	tableRows := []*internaltypes.TableRow{}
	if err != nil {
		return nil, fmt.Errorf("failed to get column types: %w", err)
	}
	for i := 0; i < len(columnTypes); i++ {
		typ, err := zetasqlite.UnmarshalDatabaseTypeName(columnTypes[i].DatabaseTypeName())
		if err != nil {
			return nil, fmt.Errorf("failed to get type from database type name: %w", err)
		}
		zetasqlType, err := typ.ToZetaSQLType()
		if err != nil {
			return nil, err
		}
		fields = append(fields, types.TableFieldSchemaFromZetaSQLType(colNames[i], zetasqlType))
	}

	var (
		totalBytes int64
		result     = [][]interface{}{}
	)
	for rows.Next() {
		values := make([]interface{}, 0, len(columnTypes))
		for i := 0; i < len(columnTypes); i++ {
			var v interface{}
			values = append(values, &v)
		}
		if err := rows.Scan(values...); err != nil {
			return nil, fmt.Errorf("failed to scan row: %w", err)
		}
		if err := rows.Err(); err != nil {
			return nil, fmt.Errorf("failed to execute query: %w", err)
		}
		cells := make([]*internaltypes.TableCell, 0, len(values))
		resultValues := make([]interface{}, 0, len(values))
		for idx, value := range values {
			v := reflect.ValueOf(value).Elem().Interface()
			if v == nil && fields[idx].Mode == string(types.RepeatedMode) {
				// GoogleSQL for BigQuery translates a NULL array into an empty array in the query result
				v = []interface{}{}
			}
			cell, err := r.convertValueToCell(v, fields[idx])
			if err != nil {
				return nil, fmt.Errorf("failed to convert value to cell: %w", err)
			}
			cell.Name = colNames[idx]
			cells = append(cells, cell)
			totalBytes += cell.Bytes
			resultValues = append(resultValues, v)
		}
		result = append(result, resultValues)
		tableRows = append(tableRows, &internaltypes.TableRow{
			F: cells,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to scan rows: %w", err)
	}
	logger.Logger(ctx).Debug("query result", zap.Any("rows", result))
	return &internaltypes.QueryResponse{
		Schema: &bigqueryv2.TableSchema{
			Fields: fields,
		},
		TotalRows:      uint64(len(tableRows)),
		JobComplete:    true,
		Rows:           tableRows,
		TotalBytes:     totalBytes,
		ChangedCatalog: changedCatalog,
	}, nil
}

func (r *Repository) queryParameterValueToGoValue(value *bigqueryv2.QueryParameterValue) (interface{}, error) {
	switch {
	case len(value.ArrayValues) != 0:
		arr := make([]interface{}, 0, len(value.ArrayValues))
		for _, v := range value.ArrayValues {
			elem, err := r.queryParameterValueToGoValue(v)
			if err != nil {
				return nil, err
			}
			arr = append(arr, elem)
		}
		return arr, nil
	case len(value.StructValues) != 0:
		st := make(map[string]interface{}, len(value.StructValues))
		for k, v := range value.StructValues {
			elem, err := r.queryParameterValueToGoValue(&v)
			if err != nil {
				return nil, err
			}
			st[k] = elem
		}
		return st, nil
	}

	// Check if the Value field is marked as null in NullFields
	// This is how Google's API client indicates a null value even though
	// Value is typed as string (not *string)
	for _, field := range value.NullFields {
		if field == "Value" {
			return nil, nil
		}
	}

	return value.Value, nil
}

// zetasqlite returns map[string]interface{} value as struct value, also returns []interface{} value as array value.
// we need to convert them to specifically TableRow and TableCell type.
// schema provides the field ordering for struct types to ensure deterministic field order.
func (r *Repository) convertValueToCell(value interface{}, schema *bigqueryv2.TableFieldSchema) (*internaltypes.TableCell, error) {
	if value == nil {
		return &internaltypes.TableCell{V: nil}, nil
	}
	rv := reflect.ValueOf(value)
	kind := rv.Type().Kind()
	if kind == reflect.Map {
		// value is struct type
		var (
			cells      []*internaltypes.TableCell
			totalBytes int64
		)

		// Build a map of field values for quick lookup
		fieldValues := make(map[string]reflect.Value)
		keys := rv.MapKeys()
		for _, key := range keys {
			fieldValues[key.Interface().(string)] = rv.MapIndex(key)
		}

		// Process fields in schema order to ensure deterministic ordering
		// (Go map iteration order is randomized)
		if schema != nil && schema.Fields != nil {
			for _, fieldSchema := range schema.Fields {
				fieldValue, exists := fieldValues[fieldSchema.Name]
				if !exists {
					// Field not present in data, skip it
					continue
				}
				cell, err := r.convertValueToCell(fieldValue.Interface(), fieldSchema)
				if err != nil {
					return nil, err
				}
				cell.Name = fieldSchema.Name
				totalBytes += cell.Bytes
				cells = append(cells, cell)
			}
		} else {
			// Fallback: no schema available, process in arbitrary order
			for _, key := range keys {
				cell, err := r.convertValueToCell(rv.MapIndex(key).Interface(), nil)
				if err != nil {
					return nil, err
				}
				cell.Name = key.Interface().(string)
				totalBytes += cell.Bytes
				cells = append(cells, cell)
			}
		}
		return &internaltypes.TableCell{V: internaltypes.TableRow{F: cells}, Bytes: totalBytes}, nil
	}
	if kind != reflect.Slice && kind != reflect.Array {
		v := fmt.Sprint(value)
		return &internaltypes.TableCell{V: v, Bytes: int64(len(v))}, nil
	}
	// array type
	var (
		cells            = []*internaltypes.TableCell{}
		totalBytes int64 = 0
	)
	// For array elements, schema.Type will be the element type (e.g., STRUCT for array of structs)
	// and schema.Fields will contain the struct fields
	var elementSchema *bigqueryv2.TableFieldSchema
	if schema != nil {
		elementSchema = &bigqueryv2.TableFieldSchema{
			Name:   schema.Name,
			Type:   schema.Type,
			Mode:   "NULLABLE", // Array elements can be nullable
			Fields: schema.Fields,
		}
	}
	for i := 0; i < rv.Len(); i++ {
		cell, err := r.convertValueToCell(rv.Index(i).Interface(), elementSchema)
		if err != nil {
			return nil, err
		}
		totalBytes += cell.Bytes
		cells = append(cells, cell)
	}
	return &internaltypes.TableCell{V: cells, Bytes: totalBytes}, nil
}

func (r *Repository) CreateOrReplaceTable(ctx context.Context, tx *connection.Tx, projectID, datasetID string, table *types.Table) error {
	tx.SetProjectAndDataset(projectID, datasetID)
	if err := tx.ContentRepoMode(); err != nil {
		return err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()

	columns := make([]string, 0, len(table.Columns))
	for _, column := range table.Columns {
		columns = append(columns,
			fmt.Sprintf("`%s` %s", column.Name, column.FormatType()),
		)
	}
	ddl := fmt.Sprintf(
		"CREATE OR REPLACE TABLE `%s` (%s)",
		r.tablePath(projectID, datasetID, table.ID), strings.Join(columns, ","),
	)

	if _, err := tx.Tx().ExecContext(ctx, ddl); err != nil {
		return fmt.Errorf("failed to execute DDL %s: %w", ddl, err)
	}
	return nil
}

func (r *Repository) AddTableData(ctx context.Context, tx *connection.Tx, projectID, datasetID string, table *types.Table, overwrite bool) error {
	if len(table.Data) == 0 {
		return nil
	}
	tx.SetProjectAndDataset(projectID, datasetID)
	if err := tx.ContentRepoMode(); err != nil {
		return err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()

	placeholders := make([]string, 0, len(table.Columns))
	columnsWithEscape := make([]string, 0, len(table.Columns))
	for _, column := range table.Columns {
		placeholders = append(placeholders, "?")
		columnsWithEscape = append(columnsWithEscape, fmt.Sprintf("`%s`", column.Name))
	}
	query := fmt.Sprintf(
		"INSERT `%s` (%s) VALUES (%s)",
		r.tablePath(projectID, datasetID, table.ID),
		strings.Join(columnsWithEscape, ","),
		strings.Join(placeholders, ","),
	)

	if overwrite {
		_, err := tx.Tx().ExecContext(ctx, fmt.Sprintf(
			"DELETE FROM `%s` WHERE true",
			r.tablePath(projectID, datasetID, table.ID),
		))
		if err != nil {
			return fmt.Errorf("failed to truncate table: %w", err)
		}
	}

	stmt, err := tx.Tx().PrepareContext(ctx, query)
	if err != nil {
		return err
	}

	for _, data := range table.Data {
		values := make([]interface{}, 0, len(table.Columns))
		for _, column := range table.Columns {
			if value, found := data[column.Name]; found {
				isTimestampColumn := column.Type == types.TIMESTAMP
				isJsonColumn := column.Type == types.JSON
				isBytesColumn := column.Type == types.BYTES
				inputString, isInputString := value.(string)

				if isInputString && isTimestampColumn {
					parsedTimestamp, err := zetasqlite.TimeFromTimestampValue(inputString)
					// If we could parse the timestamp, use it when inserting, otherwise fallback to the supplied value
					if err == nil {
						values = append(values, parsedTimestamp)
						continue
					}
				}

				if isInputString && isJsonColumn {
					var jsonValue interface{}
					if err := json.Unmarshal([]byte(inputString), &jsonValue); err != nil {
						return fmt.Errorf("failed to unmarshal json value [%s]: %w", inputString, err)
					}
					values = append(values, jsonValue)
					continue
				}

				if isInputString && isBytesColumn {
					if inputString == "" {
						values = append(values, []byte{})
						continue
					}

					decoded, err := base64.StdEncoding.DecodeString(inputString)
					if err != nil {
						return fmt.Errorf("failed to decode base64 bytes: %v - %w", inputString, err)
					}

					values = append(values, decoded)
					continue
				}

				values = append(values, value)
			} else {
				values = append(values, nil)
			}
		}

		if _, err := stmt.ExecContext(ctx, values...); err != nil {
			return err
		}
	}

	return nil
}

func (r *Repository) DeleteTables(ctx context.Context, tx *connection.Tx, projectID, datasetID string, tables []*metadata.Table) error {
	tx.SetProjectAndDataset(projectID, datasetID)
	if err := tx.ContentRepoMode(); err != nil {
		return err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()

	for _, table := range tables {
		tablePath := r.tablePath(projectID, datasetID, table.ID)
		logger.Logger(ctx).Debug("delete table", zap.String("table", tablePath))
		tableContent, err := table.Content()
		if err != nil {
			return fmt.Errorf("failed to delete table %s: %w", tablePath, err)
		}
		var query string
		switch tableContent.Type {
		case string(internaltypes.MaterializedViewTableType):
			query = fmt.Sprintf("DROP MATERIALIZED VIEW `%s`", tablePath)
		case string(internaltypes.DefaultTableType), string(internaltypes.ViewTableType):
			query = fmt.Sprintf("DROP %s `%s`", tableContent.Type, tablePath)
		default:
			return fmt.Errorf("failed to delete table with unsupported table type: %s", tableContent.Type)
		}
		if _, err := tx.Tx().ExecContext(ctx, query); err != nil {
			return fmt.Errorf("failed to delete table %s: %w", query, err)
		}
	}
	return nil
}

type RoutineType string

const (
	ScalarFunctionType      RoutineType = "SCALAR_FUNCTION"
	ProcedureType           RoutineType = "PROCEDURE"
	TableValuedFunctionType RoutineType = "TABLE_VALUED_FUNCTION"
)

type RoutineLanguageType string

const (
	LanguageTypeSQL        RoutineLanguageType = "SQL"
	LanguageTypeJavaScript RoutineLanguageType = "JavaScript"
)

func (r *Repository) AddRoutineByMetaData(ctx context.Context, tx *connection.Tx, routine *bigqueryv2.Routine) error {
	ref := routine.RoutineReference
	tx.SetProjectAndDataset(ref.ProjectId, ref.DatasetId)
	if err := tx.ContentRepoMode(); err != nil {
		return err
	}
	defer func() {
		_ = tx.MetadataRepoMode()
	}()

	var routineType string
	switch RoutineType(routine.RoutineType) {
	case ScalarFunctionType:
		routineType = "CREATE FUNCTION"
	case ProcedureType:
		routineType = "CREATE PROCEDURE"
	case TableValuedFunctionType:
		routineType = "CREATE TABLE FUNCTION"
	default:
		return fmt.Errorf("invalid routine type %s", routine.RoutineType)
	}
	switch RoutineLanguageType(routine.Language) {
	case LanguageTypeSQL:
	case LanguageTypeJavaScript:
		return fmt.Errorf("unsupported language: JavaScript")
	default:
		return fmt.Errorf("invalid language %s", routine.Language)
	}
	args := make([]string, 0, len(routine.Arguments))
	for _, arg := range routine.Arguments {
		if arg.Name == "" {
			return fmt.Errorf("invalid argument: missing name of argument")
		}
		if arg.DataType == nil {
			return fmt.Errorf("invalid argument: missing data type for %s", arg.Name)
		}
		args = append(args, fmt.Sprintf("%s %s", arg.Name, arg.DataType.TypeKind))
	}
	var retType string
	if routine.ReturnType != nil {
		retType = fmt.Sprintf(" RETURNS %s", routine.ReturnType.TypeKind)
	}
	if routine.DefinitionBody == "" {
		return fmt.Errorf("invalid body: missing function body")
	}
	query := fmt.Sprintf(
		"%s `%s`(%s)%s AS (%s)",
		routineType,
		r.routinePath(ref.ProjectId, ref.DatasetId, ref.RoutineId),
		strings.Join(args, ", "),
		retType,
		routine.DefinitionBody,
	)
	if _, err := tx.Tx().ExecContext(ctx, query); err != nil {
		return fmt.Errorf("failed to create function %s: %w", query, err)
	}
	return nil
}