AI-Support/backend/internal/database/database.go

package database

import (
	"fmt"
	"log"
	"os"
	"time"

	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"customer-support-system/internal/models"
)

var DB *gorm.DB

// Connect initializes the database connection
func Connect() error {
	dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s sslmode=%s TimeZone=%s",
		getEnv("DB_HOST", "localhost"),
		getEnv("DB_USER", "postgres"),
		getEnv("DB_PASSWORD", "postgres"),
		getEnv("DB_NAME", "support"),
		getEnv("DB_PORT", "5432"),
		getEnv("DB_SSLMODE", "disable"),
		getEnv("DB_TIMEZONE", "UTC"),
	)

	newLogger := logger.New(
		log.New(os.Stdout, "\n", log.LstdFlags),
		logger.Config{
			SlowThreshold: time.Second,
			LogLevel:      logger.Info,
			Colorful:      false,
		},
	)

	var err error
	DB, err = gorm.Open(postgres.Open(dsn), &gorm.Config{
		Logger: newLogger,
	})

	if err != nil {
		return fmt.Errorf("failed to connect to database: %w", err)
	}

	log.Println("Database connected successfully")
	return nil
}

// AutoMigrate runs the database migrations
func AutoMigrate() error {
	err := DB.AutoMigrate(
		&models.User{},
		&models.Conversation{},
		&models.Message{},
		&models.KnowledgeBase{},
		&models.KnowledgeBaseFeedback{},
		&models.AIModel{},
		&models.AIInteraction{},
		&models.AIFallback{},
	)

	if err != nil {
		return fmt.Errorf("failed to run database migrations: %w", err)
	}

	log.Println("Database migrations completed successfully")
	return nil
}

// Close closes the database connection
func Close() error {
	sqlDB, err := DB.DB()
	if err != nil {
		return fmt.Errorf("failed to get database instance: %w", err)
	}

	err = sqlDB.Close()
	if err != nil {
		return fmt.Errorf("failed to close database connection: %w", err)
	}

	log.Println("Database connection closed successfully")
	return nil
}

// GetDB returns the database instance
func GetDB() *gorm.DB {
	return DB
}

// getEnv gets an environment variable with a default value
func getEnv(key, defaultValue string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return defaultValue
}

// SeedDatabase seeds the database with initial data
func SeedDatabase() error {
	// Create default admin user if not exists
	var userCount int64
	if err := DB.Model(&models.User{}).Where("role = ?", "admin").Count(&userCount).Error; err != nil {
		return fmt.Errorf("failed to count admin users: %w", err)
	}

	if userCount == 0 {
		adminUser := models.User{
			Username:  "admin",
			Email:     "admin@example.com",
			Password:  "admin123",
			FirstName: "Admin",
			LastName:  "User",
			Role:      "admin",
			Active:    true,
		}

		if err := DB.Create(&adminUser).Error; err != nil {
			return fmt.Errorf("failed to create admin user: %w", err)
		}

		log.Println("Default admin user created")
	}

	// Create default AI models if not exists
	var aiModelCount int64
	if err := DB.Model(&models.AIModel{}).Count(&aiModelCount).Error; err != nil {
		return fmt.Errorf("failed to count AI models: %w", err)
	}

	if aiModelCount == 0 {
		// Create OpenAI GPT-4 model
		openAIModel := models.AIModel{
			Name:        "OpenAI GPT-4",
			Type:        "openai",
			Model:       "gpt-4",
			MaxTokens:   4000,
			Temperature: 0.7,
			TopP:        1.0,
			Active:      true,
			Priority:    2,
			Description: "OpenAI's GPT-4 model for complex queries",
		}

		if err := DB.Create(&openAIModel).Error; err != nil {
			return fmt.Errorf("failed to create OpenAI model: %w", err)
		}

		// Create Local LLaMA model
		localModel := models.AIModel{
			Name:        "Local LLaMA",
			Type:        "local",
			Model:       "llama2",
			Endpoint:    "http://localhost:11434",
			MaxTokens:   2000,
			Temperature: 0.7,
			TopP:        1.0,
			Active:      true,
			Priority:    1,
			Description: "Local LLaMA model for simple queries",
		}

		if err := DB.Create(&localModel).Error; err != nil {
			return fmt.Errorf("failed to create local model: %w", err)
		}

		log.Println("Default AI models created")
	}

	// Create sample knowledge base entries if not exists
	var kbCount int64
	if err := DB.Model(&models.KnowledgeBase{}).Count(&kbCount).Error; err != nil {
		return fmt.Errorf("failed to count knowledge base entries: %w", err)
	}

	if kbCount == 0 {
		sampleEntries := []models.KnowledgeBase{
			{
				Question: "How do I reset my password?",
				Answer:   "To reset your password, click on the 'Forgot Password' link on the login page. Enter your email address and follow the instructions sent to your inbox.",
				Category: "Account",
				Tags:     "password,reset,account",
				Priority: 5,
				Active:   true,
			},
			{
				Question: "What payment methods do you accept?",
				Answer:   "We accept all major credit cards including Visa, Mastercard, and American Express. We also support payments through PayPal and bank transfers.",
				Category: "Billing",
				Tags:     "payment,billing,methods",
				Priority: 4,
				Active:   true,
			},
			{
				Question: "How can I contact customer support?",
				Answer:   "You can contact our customer support team through the chat feature on our website, by emailing support@example.com, or by calling our toll-free number at 1-800-123-4567.",
				Category: "Support",
				Tags:     "contact,support,help",
				Priority: 3,
				Active:   true,
			},
		}

		for _, entry := range sampleEntries {
			if err := DB.Create(&entry).Error; err != nil {
				return fmt.Errorf("failed to create knowledge base entry: %w", err)
			}
		}

		log.Println("Sample knowledge base entries created")
	}

	return nil
}