package ai

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"io"
	"time"

	"gorm.io/gorm/clause"
	"senlinai-agent/backend/internal/models"
)

type Gateway struct {
	systemKey        string
	encryptionSecret string
}

type SelectedKey struct {
	Provider string
	APIKey   string
	KeyType  string
}

func NewGateway(systemKey string) *Gateway {
	return NewGatewayWithSecret(systemKey, "development-ai-key-secret-change-me")
}

func NewGatewayWithSecret(systemKey string, encryptionSecret string) *Gateway {
	return &Gateway{systemKey: systemKey, encryptionSecret: encryptionSecret}
}

func (g *Gateway) SaveUserKey(userID uint, provider string, apiKey string) error {
	encrypted, err := encryptAPIKey(apiKey, g.encryptionSecret)
	if err != nil {
		return err
	}
	key := models.SenlinAgentAIKey{UserID: userID, Provider: provider, EncryptedAPIKey: encrypted}
	return models.DBService.Clauses(clause.OnConflict{
		Columns:   []clause.Column{{Name: "user_id"}},
		DoUpdates: clause.AssignmentColumns([]string{"provider", "encrypted_api_key", "updated_at"}),
	}).Create(&key).Error
}

func (g *Gateway) SelectKey(userID uint) (SelectedKey, error) {
	var userKey models.SenlinAgentAIKey
	if err := models.DBService.Where("user_id = ?", userID).First(&userKey).Error; err == nil {
		apiKey, err := decryptAPIKey(userKey.EncryptedAPIKey, g.encryptionSecret)
		if err != nil {
			return SelectedKey{}, err
		}
		return SelectedKey{Provider: userKey.Provider, APIKey: apiKey, KeyType: "user"}, nil
	}
	if g.systemKey == "" {
		return SelectedKey{}, errors.New("no ai key available")
	}
	return SelectedKey{Provider: "openai", APIKey: g.systemKey, KeyType: "system"}, nil
}

func (g *Gateway) RecordCall(userID uint, provider string, usedKeyType string, action string, status string, errText string) error {
	return models.DBService.Create(&models.SenlinAgentAICallLog{
		UserID:      userID,
		Provider:    provider,
		UsedKeyType: usedKeyType,
		Action:      action,
		Status:      status,
		Error:       errText,
	}).Error
}

func (g *Gateway) CheckRateLimit(userID uint, action string, limit int, window time.Duration) error {
	if limit <= 0 {
		return nil
	}
	var count int64
	if err := models.DBService.Model(&models.SenlinAgentAICallLog{}).
		Where("user_id = ? AND action = ? AND created_at >= ?", userID, action, time.Now().Add(-window)).
		Count(&count).Error; err != nil {
		return err
	}
	if count >= int64(limit) {
		return errors.New("ai rate limit exceeded")
	}
	return nil
}

func encryptAPIKey(apiKey string, secret string) (string, error) {
	if apiKey == "" {
		return "", errors.New("api key is required")
	}
	block, err := aes.NewCipher(encryptionKey(secret))
	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(nil, nonce, []byte(apiKey), nil)
	return "v1:" + base64.RawURLEncoding.EncodeToString(append(nonce, ciphertext...)), nil
}

func decryptAPIKey(encrypted string, secret string) (string, error) {
	if len(encrypted) < 3 || encrypted[:3] != "v1:" {
		return encrypted, nil
	}
	raw, err := base64.RawURLEncoding.DecodeString(encrypted[3:])
	if err != nil {
		return "", errors.New("invalid encrypted api key")
	}
	block, err := aes.NewCipher(encryptionKey(secret))
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	if len(raw) < gcm.NonceSize() {
		return "", errors.New("invalid encrypted api key")
	}
	nonce := raw[:gcm.NonceSize()]
	ciphertext := raw[gcm.NonceSize():]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", errors.New("invalid encrypted api key")
	}
	return string(plaintext), nil
}

func encryptionKey(secret string) []byte {
	sum := sha256.Sum256([]byte(secret))
	return sum[:]
}