// Package keymanager 高性能密钥管理器
// @author OpenAI Proxy Team
// @version 2.0.0
package keymanager

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"openai-multi-key-proxy/internal/config"

	"github.com/sirupsen/logrus"
)

// KeyInfo 密钥信息
type KeyInfo struct {
	Key     string `json:"key"`
	Index   int    `json:"index"`
	Preview string `json:"preview"`
}

// Stats 统计信息
type Stats struct {
	CurrentIndex     int64 `json:"currentIndex"`
	TotalKeys        int   `json:"totalKeys"`
	HealthyKeys      int   `json:"healthyKeys"`
	BlacklistedKeys  int   `json:"blacklistedKeys"`
	SuccessCount     int64 `json:"successCount"`
	FailureCount     int64 `json:"failureCount"`
	MemoryUsage      MemoryUsage `json:"memoryUsage"`
}

// MemoryUsage 内存使用情况
type MemoryUsage struct {
	FailureCountsSize int `json:"failureCountsSize"`
	BlacklistSize     int `json:"blacklistSize"`
}

// BlacklistDetail 黑名单详情
type BlacklistDetail struct {
	Index      int    `json:"index"`
	LineNumber int    `json:"lineNumber"`
	KeyPreview string `json:"keyPreview"`
	FullKey    string `json:"fullKey"`
}

// BlacklistInfo 黑名单信息
type BlacklistInfo struct {
	TotalBlacklisted int               `json:"totalBlacklisted"`
	TotalKeys        int               `json:"totalKeys"`
	HealthyKeys      int               `json:"healthyKeys"`
	BlacklistedKeys  []BlacklistDetail `json:"blacklistedKeys"`
}

// KeyManager 密钥管理器
type KeyManager struct {
	keysFilePath    string
	keys            []string
	keyPreviews     []string
	currentIndex    int64
	blacklistedKeys sync.Map
	successCount    int64
	failureCount    int64
	keyFailureCounts sync.Map
	
	// 性能优化：预编译正则表达式
	permanentErrorPatterns []*regexp.Regexp
	
	// 内存管理
	cleanupTicker *time.Ticker
	stopCleanup   chan bool
	
	// 读写锁保护密钥列表
	keysMutex sync.RWMutex
}

// NewKeyManager 创建新的密钥管理器
func NewKeyManager(keysFilePath string) *KeyManager {
	if keysFilePath == "" {
		keysFilePath = config.AppConfig.Keys.FilePath
	}
	
	km := &KeyManager{
		keysFilePath: keysFilePath,
		currentIndex: int64(config.AppConfig.Keys.StartIndex),
		stopCleanup:  make(chan bool),
		
		// 预编译正则表达式
		permanentErrorPatterns: []*regexp.Regexp{
			regexp.MustCompile(`(?i)invalid api key`),
			regexp.MustCompile(`(?i)incorrect api key`),
			regexp.MustCompile(`(?i)api key not found`),
			regexp.MustCompile(`(?i)unauthorized`),
			regexp.MustCompile(`(?i)account deactivated`),
			regexp.MustCompile(`(?i)billing`),
		},
	}
	
	// 启动内存清理
	km.setupMemoryCleanup()
	
	return km
}

// LoadKeys 加载密钥文件
func (km *KeyManager) LoadKeys() error {
	file, err := os.Open(km.keysFilePath)
	if err != nil {
		return fmt.Errorf("无法打开密钥文件: %w", err)
	}
	defer file.Close()
	
	var keys []string
	scanner := bufio.NewScanner(file)
	
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && strings.HasPrefix(line, "sk-") {
			keys = append(keys, line)
		}
	}
	
	if err := scanner.Err(); err != nil {
		return fmt.Errorf("读取密钥文件失败: %w", err)
	}
	
	if len(keys) == 0 {
		return fmt.Errorf("密钥文件中没有有效的API密钥")
	}
	
	km.keysMutex.Lock()
	km.keys = keys
	// 预生成密钥预览，避免运行时重复计算
	km.keyPreviews = make([]string, len(keys))
	for i, key := range keys {
		if len(key) > 20 {
			km.keyPreviews[i] = key[:20] + "..."
		} else {
			km.keyPreviews[i] = key
		}
	}
	km.keysMutex.Unlock()
	
	logrus.Infof("✅ 成功加载 %d 个 API 密钥", len(keys))
	return nil
}

// GetNextKey 获取下一个可用的密钥（高性能版本）
func (km *KeyManager) GetNextKey() (*KeyInfo, error) {
	km.keysMutex.RLock()
	keysLen := len(km.keys)
	km.keysMutex.RUnlock()
	
	if keysLen == 0 {
		return nil, fmt.Errorf("没有可用的 API 密钥")
	}
	
	// 检查是否所有密钥都被拉黑
	blacklistedCount := 0
	km.blacklistedKeys.Range(func(key, value interface{}) bool {
		blacklistedCount++
		return true
	})
	
	if blacklistedCount >= keysLen {
		logrus.Warn("⚠️ 所有密钥都被拉黑，重置黑名单")
		km.blacklistedKeys = sync.Map{}
		km.keyFailureCounts = sync.Map{}
	}
	
	// 使用原子操作避免竞态条件
	attempts := 0
	for attempts < keysLen {
		currentIdx := atomic.AddInt64(&km.currentIndex, 1) - 1
		keyIndex := int(currentIdx) % keysLen
		
		km.keysMutex.RLock()
		selectedKey := km.keys[keyIndex]
		keyPreview := km.keyPreviews[keyIndex]
		km.keysMutex.RUnlock()
		
		if _, blacklisted := km.blacklistedKeys.Load(selectedKey); !blacklisted {
			return &KeyInfo{
				Key:     selectedKey,
				Index:   keyIndex,
				Preview: keyPreview,
			}, nil
		}
		
		attempts++
	}
	
	// 兜底：返回第一个密钥
	km.keysMutex.RLock()
	firstKey := km.keys[0]
	firstPreview := km.keyPreviews[0]
	km.keysMutex.RUnlock()
	
	return &KeyInfo{
		Key:     firstKey,
		Index:   0,
		Preview: firstPreview,
	}, nil
}

// RecordSuccess 记录密钥使用成功
func (km *KeyManager) RecordSuccess(key string) {
	atomic.AddInt64(&km.successCount, 1)
	// 成功时重置该密钥的失败计数
	km.keyFailureCounts.Delete(key)
}

// RecordFailure 记录密钥使用失败
func (km *KeyManager) RecordFailure(key string, err error) {
	atomic.AddInt64(&km.failureCount, 1)
	
	// 检查是否是永久性错误
	if km.isPermanentError(err) {
		km.blacklistedKeys.Store(key, true)
		km.keyFailureCounts.Delete(key) // 清理计数
		logrus.Warnf("🚫 密钥已被拉黑（永久性错误）: %s (%s)", key[:20]+"...", err.Error())
		return
	}
	
	// 临时性错误：增加失败计数
	currentFailures := 0
	if val, exists := km.keyFailureCounts.Load(key); exists {
		currentFailures = val.(int)
	}
	newFailures := currentFailures + 1
	km.keyFailureCounts.Store(key, newFailures)
	
	threshold := config.AppConfig.Keys.BlacklistThreshold
	if newFailures >= threshold {
		km.blacklistedKeys.Store(key, true)
		km.keyFailureCounts.Delete(key) // 清理计数
		logrus.Warnf("🚫 密钥已被拉黑（达到阈值）: %s (失败 %d 次: %s)", key[:20]+"...", newFailures, err.Error())
	} else {
		logrus.Warnf("⚠️ 密钥失败: %s (%d/%d 次: %s)", key[:20]+"...", newFailures, threshold, err.Error())
	}
}

// isPermanentError 判断是否是永久性错误
func (km *KeyManager) isPermanentError(err error) bool {
	errorMessage := err.Error()
	for _, pattern := range km.permanentErrorPatterns {
		if pattern.MatchString(errorMessage) {
			return true
		}
	}
	return false
}

// GetStats 获取密钥统计信息
func (km *KeyManager) GetStats() *Stats {
	km.keysMutex.RLock()
	totalKeys := len(km.keys)
	km.keysMutex.RUnlock()
	
	blacklistedCount := 0
	km.blacklistedKeys.Range(func(key, value interface{}) bool {
		blacklistedCount++
		return true
	})
	
	failureCountsSize := 0
	km.keyFailureCounts.Range(func(key, value interface{}) bool {
		failureCountsSize++
		return true
	})
	
	return &Stats{
		CurrentIndex:    atomic.LoadInt64(&km.currentIndex),
		TotalKeys:       totalKeys,
		HealthyKeys:     totalKeys - blacklistedCount,
		BlacklistedKeys: blacklistedCount,
		SuccessCount:    atomic.LoadInt64(&km.successCount),
		FailureCount:    atomic.LoadInt64(&km.failureCount),
		MemoryUsage: MemoryUsage{
			FailureCountsSize: failureCountsSize,
			BlacklistSize:     blacklistedCount,
		},
	}
}

// ResetKeys 重置密钥状态
func (km *KeyManager) ResetKeys() map[string]interface{} {
	beforeCount := 0
	km.blacklistedKeys.Range(func(key, value interface{}) bool {
		beforeCount++
		return true
	})
	
	km.blacklistedKeys = sync.Map{}
	km.keyFailureCounts = sync.Map{}
	
	logrus.Infof("🔄 密钥状态已重置，清除了 %d 个黑名单密钥", beforeCount)
	
	km.keysMutex.RLock()
	totalKeys := len(km.keys)
	km.keysMutex.RUnlock()
	
	return map[string]interface{}{
		"success":      true,
		"message":      fmt.Sprintf("已清除 %d 个黑名单密钥", beforeCount),
		"clearedCount": beforeCount,
		"totalKeys":    totalKeys,
	}
}

// GetBlacklistDetails 获取黑名单详情
func (km *KeyManager) GetBlacklistDetails() *BlacklistInfo {
	var blacklistDetails []BlacklistDetail
	
	km.keysMutex.RLock()
	keys := km.keys
	keyPreviews := km.keyPreviews
	km.keysMutex.RUnlock()
	
	for i, key := range keys {
		if _, blacklisted := km.blacklistedKeys.Load(key); blacklisted {
			blacklistDetails = append(blacklistDetails, BlacklistDetail{
				Index:      i,
				LineNumber: i + 1,
				KeyPreview: keyPreviews[i],
				FullKey:    key,
			})
		}
	}
	
	return &BlacklistInfo{
		TotalBlacklisted: len(blacklistDetails),
		TotalKeys:        len(keys),
		HealthyKeys:      len(keys) - len(blacklistDetails),
		BlacklistedKeys:  blacklistDetails,
	}
}

// setupMemoryCleanup 设置内存清理机制
func (km *KeyManager) setupMemoryCleanup() {
	km.cleanupTicker = time.NewTicker(10 * time.Minute)
	
	go func() {
		for {
			select {
			case <-km.cleanupTicker.C:
				km.performMemoryCleanup()
			case <-km.stopCleanup:
				km.cleanupTicker.Stop()
				return
			}
		}
	}()
}

// performMemoryCleanup 执行内存清理
func (km *KeyManager) performMemoryCleanup() {
	km.keysMutex.RLock()
	maxSize := len(km.keys) * 2
	if maxSize < 1000 {
		maxSize = 1000
	}
	km.keysMutex.RUnlock()
	
	currentSize := 0
	km.keyFailureCounts.Range(func(key, value interface{}) bool {
		currentSize++
		return true
	})
	
	if currentSize > maxSize {
		logrus.Infof("🧹 清理失败计数缓存 (%d -> %d)", currentSize, maxSize)
		
		// 简单策略：清理一半的失败计数
		cleared := 0
		target := currentSize - maxSize
		
		km.keyFailureCounts.Range(func(key, value interface{}) bool {
			if cleared < target {
				km.keyFailureCounts.Delete(key)
				cleared++
			}
			return cleared < target
		})
	}
}

// Close 关闭密钥管理器
func (km *KeyManager) Close() {
	close(km.stopCleanup)
}