internal/db/userbot_manager.go

package db

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.uber.org/zap"
)

// UserBotManager handles user bot operations
type UserBotManager struct {
	collection   *mongo.Collection
	log          *zap.Logger
	encryptionKey []byte // 32 bytes for AES-256
}

// NewUserBotManager creates a new UserBotManager instance
func NewUserBotManager(database *mongo.Database, log *zap.Logger, encryptionKey string) *UserBotManager {
	// Ensure encryption key is 32 bytes for AES-256
	key := make([]byte, 32)
	copy(key, []byte(encryptionKey))
	
	return &UserBotManager{
		collection:    database.Collection("user_bots"),
		log:           log.Named("UserBotManager"),
		encryptionKey: key,
	}
}

// encryptToken encrypts a bot token using AES-256-GCM
func (ubm *UserBotManager) encryptToken(token string) (string, error) {
	block, err := aes.NewCipher(ubm.encryptionKey)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	ciphertext := gcm.Seal(nonce, nonce, []byte(token), nil)
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// decryptToken decrypts a bot token
func (ubm *UserBotManager) decryptToken(encryptedToken string) (string, error) {
	data, err := base64.StdEncoding.DecodeString(encryptedToken)
	if err != nil {
		return "", err
	}

	block, err := aes.NewCipher(ubm.encryptionKey)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return "", fmt.Errorf("ciphertext too short")
	}

	nonce, ciphertext := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

// AddUserBot adds or updates a user's bot configuration
func (ubm *UserBotManager) AddUserBot(ctx context.Context, userID int64, botToken, botUsername string, botID int64) error {
	// Encrypt the bot token
	encryptedToken, err := ubm.encryptToken(botToken)
	if err != nil {
		return fmt.Errorf("failed to encrypt bot token: %w", err)
	}

	userBot := UserBot{
		UserID:      userID,
		BotToken:    encryptedToken,
		BotUsername: botUsername,
		BotID:       botID,
		IsActive:    true,
		IsVerified:  true, // Assume verified if we got here
		LastUsed:    time.Now(),
		ErrorCount:  0,
		CreatedAt:   time.Now(),
		UpdatedAt:   time.Now(),
	}

	// Upsert the user bot (replace if exists)
	opts := options.Replace().SetUpsert(true)
	_, err = ubm.collection.ReplaceOne(
		ctx,
		bson.M{"user_id": userID},
		userBot,
		opts,
	)

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

	ubm.log.Info("Added user bot", 
		zap.Int64("userID", userID), 
		zap.String("botUsername", botUsername),
		zap.Int64("botID", botID))

	return nil
}

// GetUserBot retrieves a user's bot configuration
func (ubm *UserBotManager) GetUserBot(ctx context.Context, userID int64) (*UserBot, error) {
	var userBot UserBot
	
	err := ubm.collection.FindOne(ctx, bson.M{"user_id": userID}).Decode(&userBot)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, nil // No bot configured
		}
		return nil, fmt.Errorf("failed to get user bot: %w", err)
	}

	return &userBot, nil
}

// GetDecryptedBotToken retrieves and decrypts a user's bot token
func (ubm *UserBotManager) GetDecryptedBotToken(ctx context.Context, userID int64) (string, error) {
	userBot, err := ubm.GetUserBot(ctx, userID)
	if err != nil {
		return "", err
	}
	
	if userBot == nil {
		return "", fmt.Errorf("no bot configured for user")
	}

	if !userBot.IsActive {
		return "", fmt.Errorf("user bot is disabled")
	}

	token, err := ubm.decryptToken(userBot.BotToken)
	if err != nil {
		return "", fmt.Errorf("failed to decrypt bot token: %w", err)
	}

	return token, nil
}

// UpdateBotUsage updates the last used time and resets error count on successful use
func (ubm *UserBotManager) UpdateBotUsage(ctx context.Context, userID int64) error {
	update := bson.M{
		"$set": bson.M{
			"last_used":   time.Now(),
			"updated_at":  time.Now(),
			"error_count": 0,
			"last_error":  "",
		},
	}

	_, err := ubm.collection.UpdateOne(
		ctx,
		bson.M{"user_id": userID},
		update,
	)

	return err
}

// RecordBotError records an error for a user's bot
func (ubm *UserBotManager) RecordBotError(ctx context.Context, userID int64, errorMsg string) error {
	update := bson.M{
		"$inc": bson.M{"error_count": 1},
		"$set": bson.M{
			"last_error":  errorMsg,
			"updated_at":  time.Now(),
		},
	}

	// Disable bot if too many consecutive errors
	result, err := ubm.collection.UpdateOne(
		ctx,
		bson.M{"user_id": userID},
		update,
	)

	if err != nil {
		return err
	}

	// Check if we need to disable the bot due to too many errors
	if result.ModifiedCount > 0 {
		userBot, err := ubm.GetUserBot(ctx, userID)
		if err == nil && userBot != nil && userBot.ErrorCount >= 5 {
			// Disable bot after 5 consecutive errors
			ubm.DisableUserBot(ctx, userID)
			ubm.log.Warn("Disabled user bot due to too many errors",
				zap.Int64("userID", userID),
				zap.String("lastError", errorMsg))
		}
	}

	return nil
}

// DisableUserBot disables a user's bot
func (ubm *UserBotManager) DisableUserBot(ctx context.Context, userID int64) error {
	update := bson.M{
		"$set": bson.M{
			"is_active":  false,
			"updated_at": time.Now(),
		},
	}

	_, err := ubm.collection.UpdateOne(
		ctx,
		bson.M{"user_id": userID},
		update,
	)

	return err
}

// EnableUserBot enables a user's bot
func (ubm *UserBotManager) EnableUserBot(ctx context.Context, userID int64) error {
	update := bson.M{
		"$set": bson.M{
			"is_active":   true,
			"error_count": 0,
			"last_error":  "",
			"updated_at":  time.Now(),
		},
	}

	_, err := ubm.collection.UpdateOne(
		ctx,
		bson.M{"user_id": userID},
		update,
	)

	return err
}

// RemoveUserBot removes a user's bot configuration
func (ubm *UserBotManager) RemoveUserBot(ctx context.Context, userID int64) error {
	_, err := ubm.collection.DeleteOne(ctx, bson.M{"user_id": userID})
	if err != nil {
		return fmt.Errorf("failed to remove user bot: %w", err)
	}

	ubm.log.Info("Removed user bot", zap.Int64("userID", userID))
	return nil
}

// ListActiveBots returns all active user bots (for admin purposes)
func (ubm *UserBotManager) ListActiveBots(ctx context.Context) ([]UserBot, error) {
	cursor, err := ubm.collection.Find(ctx, bson.M{"is_active": true})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var bots []UserBot
	if err := cursor.All(ctx, &bots); err != nil {
		return nil, err
	}

	return bots, nil
}