package graph

// This file will be automatically regenerated based on the schema, any resolver implementations
// will be copied through when generating and any unknown code will be moved to the end.
// Code generated by github.com/99designs/gqlgen version v0.17.80

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/clidey/whodb/core/graph/model"
	"github.com/clidey/whodb/core/src"
	"github.com/clidey/whodb/core/src/auth"
	"github.com/clidey/whodb/core/src/common"
	"github.com/clidey/whodb/core/src/engine"
	"github.com/clidey/whodb/core/src/env"
	"github.com/clidey/whodb/core/src/llm"
	"github.com/clidey/whodb/core/src/log"
	gorm_plugin "github.com/clidey/whodb/core/src/plugins/gorm"
	"github.com/clidey/whodb/core/src/settings"
	"gorm.io/gorm"
)

// Login is the resolver for the Login field.
func (r *mutationResolver) Login(ctx context.Context, credentials model.LoginCredentials) (*model.StatusResponse, error) {
	advanced := []engine.Record{}
	for _, recordInput := range credentials.Advanced {
		advanced = append(advanced, engine.Record{
			Key:   recordInput.Key,
			Value: recordInput.Value,
		})
	}
	if !src.MainEngine.Choose(engine.DatabaseType(credentials.Type)).IsAvailable(&engine.PluginConfig{
		Credentials: &engine.Credentials{
			Type:     credentials.Type,
			Hostname: credentials.Hostname,
			Username: credentials.Username,
			Password: credentials.Password,
			Database: credentials.Database,
			Advanced: advanced,
		},
	}) {
		log.LogFields(log.Fields{
			"type":     credentials.Type,
			"hostname": credentials.Hostname,
			"username": credentials.Username,
			"database": credentials.Database,
		}).Error("Database connection failed during login - credentials unauthorized")
		return nil, errors.New("unauthorized")
	}
	return auth.Login(ctx, &credentials)
}

// LoginWithProfile is the resolver for the LoginWithProfile field.
func (r *mutationResolver) LoginWithProfile(ctx context.Context, profile model.LoginProfileInput) (*model.StatusResponse, error) {
    profiles := src.GetLoginProfiles()
    for i, loginProfile := range profiles {
        profileId := src.GetLoginProfileId(i, loginProfile)
        if profile.ID == profileId {
            if !src.MainEngine.Choose(engine.DatabaseType(loginProfile.Type)).IsAvailable(&engine.PluginConfig{
                Credentials: src.GetLoginCredentials(loginProfile),
            }) {
				log.LogFields(log.Fields{
					"profile_id": profile.ID,
					"type":       loginProfile.Type,
				}).Error("Database connection failed for login profile - credentials unauthorized")
				return nil, errors.New("unauthorized")
			}
            // Build full credentials so the auth layer can persist them to keyring
            resolved := src.GetLoginCredentials(loginProfile)
            credentials := &model.LoginCredentials{
                ID:       &profile.ID,
                Type:     resolved.Type,
                Hostname: resolved.Hostname,
                Username: resolved.Username,
                Password: resolved.Password,
                Database: resolved.Database,
                Advanced: func() []*model.RecordInput {
                    var out []*model.RecordInput
                    for _, rec := range resolved.Advanced {
                        out = append(out, &model.RecordInput{Key: rec.Key, Value: rec.Value})
                    }
                    return out
                }(),
            }
            if profile.Database != nil && *profile.Database != "" {
                credentials.Database = *profile.Database
            }
            return auth.Login(ctx, credentials)
        }
    }
	log.LogFields(log.Fields{
		"profile_id": profile.ID,
	}).Error("Login profile not found or not authorized")
	return nil, errors.New("login profile does not exist or is not authorized")
}

// Logout is the resolver for the Logout field.
func (r *mutationResolver) Logout(ctx context.Context) (*model.StatusResponse, error) {
	return auth.Logout(ctx)
}

// UpdateSettings is the resolver for the UpdateSettings field.
func (r *mutationResolver) UpdateSettings(ctx context.Context, newSettings model.SettingsConfigInput) (*model.StatusResponse, error) {
	var fields []settings.ISettingsField

	if newSettings.MetricsEnabled != nil {
		fields = append(fields, settings.MetricsEnabledField(common.StrPtrToBool(newSettings.MetricsEnabled)))
	}

	updated := settings.UpdateSettings(fields...)
	return &model.StatusResponse{
		Status: updated,
	}, nil
}

// AddStorageUnit is the resolver for the AddStorageUnit field.
func (r *mutationResolver) AddStorageUnit(ctx context.Context, schema string, storageUnit string, fields []*model.RecordInput) (*model.StatusResponse, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	fieldsMap := []engine.Record{}
	for _, field := range fields {
		extraFields := map[string]string{}
		for _, extraField := range field.Extra {
			extraFields[extraField.Key] = extraField.Value
		}
		fieldsMap = append(fieldsMap, engine.Record{
			Key:   field.Key,
			Value: field.Value,
			Extra: extraFields,
		})
	}
	status, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).AddStorageUnit(config, schema, storageUnit, fieldsMap)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "AddStorageUnit",
			"schema":        schema,
			"storage_unit":  storageUnit,
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	return &model.StatusResponse{
		Status: status,
	}, nil
}

// UpdateStorageUnit is the resolver for the UpdateStorageUnit field.
func (r *mutationResolver) UpdateStorageUnit(ctx context.Context, schema string, storageUnit string, values []*model.RecordInput, updatedColumns []string) (*model.StatusResponse, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	valuesMap := map[string]string{}
	for _, value := range values {
		valuesMap[value.Key] = value.Value
	}
	status, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).UpdateStorageUnit(config, schema, storageUnit, valuesMap, updatedColumns)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":       "UpdateStorageUnit",
			"schema":          schema,
			"storage_unit":    storageUnit,
			"database_type":   typeArg,
			"updated_columns": len(updatedColumns),
			"error":           err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	return &model.StatusResponse{
		Status: status,
	}, nil
}

// AddRow is the resolver for the AddRow field.
func (r *mutationResolver) AddRow(ctx context.Context, schema string, storageUnit string, values []*model.RecordInput) (*model.StatusResponse, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type

	// Debug logging
	log.LogFields(log.Fields{
		"operation":     "AddRow-Resolver",
		"schema":        schema,
		"storage_unit":  storageUnit,
		"database_type": typeArg,
		"values_count":  len(values),
	}).Debug("AddRow resolver called")

	valuesRecords := []engine.Record{}
	for _, field := range values {
		extraFields := map[string]string{}
		for _, extraField := range field.Extra {
			extraFields[extraField.Key] = extraField.Value
		}
		valuesRecords = append(valuesRecords, engine.Record{
			Key:   field.Key,
			Value: field.Value,
			Extra: extraFields,
		})
	}
	status, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).AddRow(config, schema, storageUnit, valuesRecords)

	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "AddRow",
			"schema":        schema,
			"storage_unit":  storageUnit,
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	return &model.StatusResponse{
		Status: status,
	}, nil
}

// DeleteRow is the resolver for the DeleteRow field.
func (r *mutationResolver) DeleteRow(ctx context.Context, schema string, storageUnit string, values []*model.RecordInput) (*model.StatusResponse, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	valuesMap := map[string]string{}
	for _, value := range values {
		valuesMap[value.Key] = value.Value
	}
	status, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).DeleteRow(config, schema, storageUnit, valuesMap)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "DeleteRow",
			"schema":        schema,
			"storage_unit":  storageUnit,
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	return &model.StatusResponse{
		Status: status,
	}, nil
}

// GenerateMockData is the resolver for the GenerateMockData field.
func (r *mutationResolver) GenerateMockData(ctx context.Context, input model.MockDataGenerationInput) (*model.MockDataGenerationStatus, error) {
	log.Logger.WithField("schema", input.Schema).WithField("table", input.StorageUnit).WithField("rowCount", input.RowCount).WithField("overwrite", input.OverwriteExisting).Info("Starting mock data generation")

	// Enforce maximum row limit
	maxRowLimit := env.GetMockDataGenerationMaxRowCount()
	if input.RowCount > maxRowLimit {
		log.Logger.WithField("requested", input.RowCount).WithField("max", maxRowLimit).Error("Row count exceeds maximum limit")
		return nil, fmt.Errorf("row count exceeds maximum limit of %d", maxRowLimit)
	}

	// Check if mock data generation is allowed for this table
	if !env.IsMockDataGenerationAllowed(input.StorageUnit) {
		log.Logger.WithField("table", input.StorageUnit).Error("Mock data generation not allowed for table")
		return nil, errors.New("mock data generation is not allowed for this table")
	}

	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	plugin := src.MainEngine.Choose(engine.DatabaseType(typeArg))

	// Get table columns to understand the schema
	log.Logger.Debug("Fetching table schema")
	rowsResult, err := plugin.GetRows(config, input.Schema, input.StorageUnit, nil, []*model.SortCondition{}, 1, 0)
	if err != nil {
		log.Logger.WithError(err).Error("Failed to get table schema")
		return nil, fmt.Errorf("failed to get table schema: %w", err)
	}
	log.Logger.WithField("columnCount", len(rowsResult.Columns)).Debug("Got table columns")

	// Get column constraints from the database
	log.Logger.Debug("Fetching column constraints")
	constraints, err := plugin.GetColumnConstraints(config, input.Schema, input.StorageUnit)
	if err != nil {
		log.Logger.WithError(err).Warn("Failed to get column constraints, using empty constraints")
		// Use empty constraints on error
		constraints = make(map[string]map[string]any)
	}
	log.Logger.WithField("constraintCount", len(constraints)).Debug("Got column constraints")

	// Pre-generate rows with brute force approach to ensure we get exactly the requested count
	generator := src.NewMockDataGenerator()
	// Generate extra rows as buffer for potential constraint violations
	// We'll generate 50% more rows as buffer, capped at requested + 100
	bufferSize := min(input.RowCount/2, 100)
	maxGenerationAttempts := input.RowCount + bufferSize

	generatedRowsBuffer := make([][]engine.Record, 0, maxGenerationAttempts)
	attemptsCount := 0
	maxTotalAttempts := maxGenerationAttempts * 10 // Overall safety limit to prevent infinite loops

	// Keep generating until we have enough rows or hit the safety limit
	log.Logger.WithField("targetRows", maxGenerationAttempts).Info("Starting row generation")
	for len(generatedRowsBuffer) < maxGenerationAttempts && attemptsCount < maxTotalAttempts {
		attemptsCount++
		rowData, genErr := generator.GenerateRowDataWithConstraints(rowsResult.Columns, constraints)
		if genErr != nil {
			log.Logger.WithError(genErr).WithField("attempt", attemptsCount).Debug("Failed to generate row")
			continue
		}
		generatedRowsBuffer = append(generatedRowsBuffer, rowData)
	}

	log.Logger.WithField("generatedRows", len(generatedRowsBuffer)).WithField("attempts", attemptsCount).Info("Completed row generation buffer")

	if len(generatedRowsBuffer) == 0 {
		log.Logger.Error("Failed to generate any valid rows")
		return nil, errors.New("failed to generate any valid rows after multiple attempts")
	}

	// Check if this is a GORM-based plugin (SQL databases)
	gormPlugin, isGormPlugin := plugin.PluginFunctions.(*gorm_plugin.GormPlugin)

	generatedRows := 0
	if isGormPlugin {
		// Use transaction for SQL databases
		err = gormPlugin.ExecuteInTransaction(config, func(tx *gorm.DB) error {
			// Clear existing data if requested
			if input.OverwriteExisting {
				if err := gormPlugin.ClearTableDataInTx(tx, input.Schema, input.StorageUnit); err != nil {
					return fmt.Errorf("failed to clear existing data: %w", err)
				}
			}

			// Insert rows until we reach exactly the requested count
			// Use brute force approach - keep trying from our buffer until we get the exact count
			successfulRows := 0
			bufferIndex := 0
			retryWithDefaults := false

			for successfulRows < input.RowCount && bufferIndex < len(generatedRowsBuffer) {
				rowData := generatedRowsBuffer[bufferIndex]
				bufferIndex++

				if err := gormPlugin.AddRowInTx(tx, input.Schema, input.StorageUnit, rowData); err != nil {
					errStr := err.Error()
					log.Logger.WithError(err).WithField("rowIndex", bufferIndex).WithField("successfulSoFar", successfulRows).Error("Failed to insert row")
					// If it's a constraint error, try with default values
					if strings.Contains(errStr, "constraint") || strings.Contains(errStr, "CHECK") || strings.Contains(errStr, "check") {
						log.Logger.Error("Constraint violation detected, trying with default values")
						// Try once with defaults
						if !retryWithDefaults {
							defaultRow := generator.GenerateRowWithDefaults(rowsResult.Columns)
							if retryErr := gormPlugin.AddRowInTx(tx, input.Schema, input.StorageUnit, defaultRow); retryErr == nil {
								log.Logger.Info("Successfully inserted row with default values")
								successfulRows++
								retryWithDefaults = false
								continue
							} else {
								log.Logger.WithError(retryErr).Error("Failed to insert row with default values")
							}
							retryWithDefaults = true // Mark that we've tried defaults
						}
						// Skip this row and try next from buffer
						continue
					}
					// For non-constraint errors in overwrite mode, fail the transaction
					if input.OverwriteExisting {
						log.Logger.WithError(err).Error("Non-constraint error in overwrite mode, failing transaction")
						return fmt.Errorf("failed to insert row: %w", err)
					}
					log.Logger.WithError(err).Warn("Non-constraint error in append mode, skipping row")
					// For append mode, skip and continue
					continue
				}
				successfulRows++
				retryWithDefaults = false
				if successfulRows%10 == 0 {
					log.Logger.WithField("progress", successfulRows).WithField("target", input.RowCount).Debug("Mock data insertion progress")
				}
			}

			// If we couldn't generate exactly the requested amount, try generating more on the fly
			log.Logger.WithField("successfulRows", successfulRows).WithField("target", input.RowCount).Debug("Checking if additional generation needed")
			additionalAttempts := 0
			maxAdditionalAttempts := input.RowCount * 5 // Safety limit

			for successfulRows < input.RowCount && additionalAttempts < maxAdditionalAttempts {
				additionalAttempts++

				// Generate a new row
				newRow, genErr := generator.GenerateRowDataWithConstraints(rowsResult.Columns, constraints)
				if genErr != nil {
					continue
				}

				if err := gormPlugin.AddRowInTx(tx, input.Schema, input.StorageUnit, newRow); err != nil {
					// Try with defaults
					defaultRow := generator.GenerateRowWithDefaults(rowsResult.Columns)
					if retryErr := gormPlugin.AddRowInTx(tx, input.Schema, input.StorageUnit, defaultRow); retryErr == nil {
						successfulRows++
					}
					continue
				}
				successfulRows++
			}

			generatedRows = successfulRows
			log.Logger.WithField("generatedRows", generatedRows).WithField("target", input.RowCount).Info("Completed mock data insertion in transaction")

			// If overwriting and we couldn't generate enough rows, fail the transaction
			if input.OverwriteExisting && generatedRows == 0 {
				log.Logger.Error("Failed to generate any valid rows in overwrite mode")
				return errors.New("failed to generate any valid rows - existing data preserved")
			}

			return nil
		})
	} else {
		// For NoSQL databases, operations are not transactional
		if input.OverwriteExisting {
			if _, err := plugin.ClearTableData(config, input.Schema, input.StorageUnit); err != nil {
				return nil, fmt.Errorf("failed to clear existing data: %w", err)
			}
		}

		// Insert rows until we reach exactly the requested count
		bufferIndex := 0
		for generatedRows < input.RowCount && bufferIndex < len(generatedRowsBuffer) {
			rowData := generatedRowsBuffer[bufferIndex]
			bufferIndex++

			success, addErr := plugin.AddRow(config, input.Schema, input.StorageUnit, rowData)
			if addErr != nil || !success {
				// Try with defaults
				defaultRow := generator.GenerateRowWithDefaults(rowsResult.Columns)
				success2, err2 := plugin.AddRow(config, input.Schema, input.StorageUnit, defaultRow)
				if err2 == nil && success2 {
					generatedRows++
				}
				continue
			}
			generatedRows++
		}

		// If we need more rows, generate them on the fly
		additionalAttempts := 0
		maxAdditionalAttempts := input.RowCount * 5

		for generatedRows < input.RowCount && additionalAttempts < maxAdditionalAttempts {
			additionalAttempts++

			newRow, genErr := generator.GenerateRowDataWithConstraints(rowsResult.Columns, constraints)
			if genErr != nil {
				continue
			}

			success, addErr := plugin.AddRow(config, input.Schema, input.StorageUnit, newRow)
			if addErr != nil || !success {
				// Try with defaults
				defaultRow := generator.GenerateRowWithDefaults(rowsResult.Columns)
				success2, err2 := plugin.AddRow(config, input.Schema, input.StorageUnit, defaultRow)
				if err2 == nil && success2 {
					generatedRows++
				}
				continue
			}
			generatedRows++
		}
	}

	if err != nil {
		log.Logger.WithError(err).WithField("generatedRows", generatedRows).Error("Mock data generation failed")
		return nil, fmt.Errorf("mock data generation failed: %w", err)
	}

	// Return success if any rows were generated
	log.Logger.WithField("generatedRows", generatedRows).Info("Mock data generation completed successfully")
	return &model.MockDataGenerationStatus{
		AmountGenerated: generatedRows,
	}, nil
}

// Version is the resolver for the Version field.
func (r *queryResolver) Version(ctx context.Context) (string, error) {
	return env.GetClideyQuickContainerImage(), nil
}

// Profiles is the resolver for the Profiles field.
func (r *queryResolver) Profiles(ctx context.Context) ([]*model.LoginProfile, error) {
	profiles := []*model.LoginProfile{}
	for i, profile := range src.GetLoginProfiles() {
		profileName := src.GetLoginProfileId(i, profile)
		loginProfile := &model.LoginProfile{
			ID:                   profileName,
			Type:                 model.DatabaseType(profile.Type),
			Database:             &profile.Database,
			IsEnvironmentDefined: true,
			Source:               profile.Source,
		}
		if len(profile.Alias) > 0 {
			loginProfile.Alias = &profile.Alias
		}
		if len(profile.CustomId) > 0 {
			loginProfile.ID = profile.CustomId
		}
		profiles = append(profiles, loginProfile)
	}
	return profiles, nil
}

// Database is the resolver for the Database field.
func (r *queryResolver) Database(ctx context.Context, typeArg string) ([]string, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	databases, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).GetDatabases(config)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "GetDatabases",
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	return databases, nil
}

// Schema is the resolver for the Schema field.
func (r *queryResolver) Schema(ctx context.Context) ([]string, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	schemas, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).GetAllSchemas(config)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "GetAllSchemas",
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	return schemas, nil
}

// StorageUnit is the resolver for the StorageUnit field.
func (r *queryResolver) StorageUnit(ctx context.Context, schema string) ([]*model.StorageUnit, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	units, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).GetStorageUnits(config, schema)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "GetStorageUnits",
			"schema":        schema,
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	storageUnits := []*model.StorageUnit{}
	for _, unit := range units {
		storageUnit := engine.GetStorageUnitModel(unit)
		storageUnit.IsMockDataGenerationAllowed = env.IsMockDataGenerationAllowed(unit.Name)
		storageUnits = append(storageUnits, storageUnit)
	}
	return storageUnits, nil
}

// Row is the resolver for the Row field.
func (r *queryResolver) Row(ctx context.Context, schema string, storageUnit string, where *model.WhereCondition, sort []*model.SortCondition, pageSize int, pageOffset int) (*model.RowsResult, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	rowsResult, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).GetRows(config, schema, storageUnit, where, sort, pageSize, pageOffset)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "GetRows",
			"schema":        schema,
			"storage_unit":  storageUnit,
			"database_type": typeArg,
			"page_size":     pageSize,
			"page_offset":   pageOffset,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	columns := []*model.Column{}
	for _, column := range rowsResult.Columns {
		columns = append(columns, &model.Column{
			Type: column.Type,
			Name: column.Name,
		})
	}
	return &model.RowsResult{
		Columns:       columns,
		Rows:          rowsResult.Rows,
		DisableUpdate: rowsResult.DisableUpdate,
	}, nil
}

// Columns is the resolver for the Columns field.
func (r *queryResolver) Columns(ctx context.Context, schema string, storageUnit string) ([]*model.Column, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	columnsResult, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).GetColumnsForTable(config, schema, storageUnit)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "GetColumnsForTable",
			"schema":        schema,
			"storage_unit":  storageUnit,
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	columns := []*model.Column{}
	for _, column := range columnsResult {
		columns = append(columns, &model.Column{
			Type: column.Type,
			Name: column.Name,
		})
	}
	return columns, nil
}

// RawExecute is the resolver for the RawExecute field.
func (r *queryResolver) RawExecute(ctx context.Context, query string) (*model.RowsResult, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	rowsResult, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).RawExecute(config, query)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "RawExecute",
			"database_type": typeArg,
			"query":         query,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	columns := []*model.Column{}
	for _, column := range rowsResult.Columns {
		columns = append(columns, &model.Column{
			Type: column.Type,
			Name: column.Name,
		})
	}
	return &model.RowsResult{
		Columns: columns,
		Rows:    rowsResult.Rows,
	}, nil
}

// Graph is the resolver for the Graph field.
func (r *queryResolver) Graph(ctx context.Context, schema string) ([]*model.GraphUnit, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	graphUnits, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).GetGraph(config, schema)
	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "GetGraph",
			"schema":        schema,
			"database_type": typeArg,
			"error":         err.Error(),
		}).Error("Database operation failed")
		return nil, err
	}
	graphUnitsModel := []*model.GraphUnit{}
	for _, graphUnit := range graphUnits {
		relations := []*model.GraphUnitRelationship{}
		for _, relation := range graphUnit.Relations {
			relations = append(relations, &model.GraphUnitRelationship{
				Name:         relation.Name,
				Relationship: model.GraphUnitRelationshipType(relation.RelationshipType),
			})
		}
		graphUnitsModel = append(graphUnitsModel, &model.GraphUnit{
			Unit:      engine.GetStorageUnitModel(graphUnit.Unit),
			Relations: relations,
		})
	}
	return graphUnitsModel, nil
}

// AIProviders is the resolver for the AIProviders field.
func (r *queryResolver) AIProviders(ctx context.Context) ([]*model.AIProvider, error) {
	providers := env.GetConfiguredChatProviders()
	aiProviders := []*model.AIProvider{}
	for _, provider := range providers {
		aiProviders = append(aiProviders, &model.AIProvider{
			Type:                 provider.Type,
			ProviderID:           provider.ProviderId,
			IsEnvironmentDefined: true,
		})
	}
	return aiProviders, nil
}

// AIModel is the resolver for the AIModel field.
func (r *queryResolver) AIModel(ctx context.Context, providerID *string, modelType string, token *string) ([]string, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))

	// Initialize ExternalModel to prevent nil pointer dereference
	config.ExternalModel = &engine.ExternalModel{
		Type: modelType,
	}

	if providerID != nil {
		providers := env.GetConfiguredChatProviders()
		for _, provider := range providers {
			if provider.ProviderId == *providerID {
				config.ExternalModel.Token = provider.APIKey
				break
			}
		}
	} else if token != nil {
		config.ExternalModel.Token = *token
	}
	models, err := llm.Instance(config).GetSupportedModels()
	if err != nil {
		log.LogFields(log.Fields{
			"operation":   "GetSupportedModels",
			"model_type":  modelType,
			"provider_id": providerID,
			"error":       err.Error(),
		}).Error("AI operation failed")
		return nil, err
	}
	return models, nil
}

// AIChat is the resolver for the AIChat field.
func (r *queryResolver) AIChat(ctx context.Context, providerID *string, modelType string, token *string, schema string, input model.ChatInput) ([]*model.AIChatMessage, error) {
	config := engine.NewPluginConfig(auth.GetCredentials(ctx))
	typeArg := config.Credentials.Type
	if providerID != nil {
		providers := env.GetConfiguredChatProviders()
		for _, provider := range providers {
			if provider.ProviderId == *providerID {
				config.ExternalModel = &engine.ExternalModel{
					Type:  modelType,
					Token: provider.APIKey,
				}
			}
		}
	} else {
		config.ExternalModel = &engine.ExternalModel{
			Type: modelType,
		}
		if token != nil {
			config.ExternalModel.Token = *token
		}
	}
	messages, err := src.MainEngine.Choose(engine.DatabaseType(typeArg)).Chat(config, schema, input.Model, input.PreviousConversation, input.Query)

	if err != nil {
		log.LogFields(log.Fields{
			"operation":     "Chat",
			"schema":        schema,
			"database_type": typeArg,
			"model":         input.Model,
			"model_type":    modelType,
			"provider_id":   providerID,
			"query":         input.Query,
			"error":         err.Error(),
		}).Error("AI chat operation failed")
		return nil, err
	}

	chatResponse := []*model.AIChatMessage{}

	for _, message := range messages {
		var result *model.RowsResult
		if strings.HasPrefix(message.Type, "sql") {
			columns := []*model.Column{}
			for _, column := range message.Result.Columns {
				columns = append(columns, &model.Column{
					Type: column.Type,
					Name: column.Name,
				})
			}
			result = &model.RowsResult{
				Columns: columns,
				Rows:    message.Result.Rows,
			}
		}
		chatResponse = append(chatResponse, &model.AIChatMessage{
			Type:   message.Type,
			Result: result,
			Text:   message.Text,
		})
	}

	return chatResponse, nil
}

// SettingsConfig is the resolver for the SettingsConfig field.
func (r *queryResolver) SettingsConfig(ctx context.Context) (*model.SettingsConfig, error) {
	currentSettings := settings.Get()
	return &model.SettingsConfig{MetricsEnabled: &currentSettings.MetricsEnabled}, nil
}

// MockDataMaxRowCount is the resolver for the MockDataMaxRowCount field.
func (r *queryResolver) MockDataMaxRowCount(ctx context.Context) (int, error) {
	return env.GetMockDataGenerationMaxRowCount(), nil
}

// Mutation returns MutationResolver implementation.
func (r *Resolver) Mutation() MutationResolver { return &mutationResolver{r} }

// Query returns QueryResolver implementation.
func (r *Resolver) Query() QueryResolver { return &queryResolver{r} }

type mutationResolver struct{ *Resolver }
type queryResolver struct{ *Resolver }
