from functools import wraps
import time
from flask import current_app
import logging
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy import text
from scripts.models import db
from collections import defaultdict

# 简单的内存缓存实现
class QueryCache:
    def __init__(self, max_size=100, ttl=300):  # 默认缓存300秒
        self.cache = {}
        self.max_size = max_size
        self.ttl = ttl
        self.access_times = defaultdict(int)

    def get(self, key):
        if key in self.cache:
            item = self.cache[key]
            if time.time() - item['time'] < self.ttl:
                self.access_times[key] += 1
                return item['value']
            else:
                # 缓存过期，删除
                del self.cache[key]
                if key in self.access_times:
                    del self.access_times[key]
        return None

    def set(self, key, value):
        # 如果缓存满了，删除最少访问的项
        if len(self.cache) >= self.max_size:
            # 找出访问次数最少的键
            if self.access_times:
                min_key = min(self.access_times, key=self.access_times.get)
                if min_key in self.cache:
                    del self.cache[min_key]
                del self.access_times[min_key]

        self.cache[key] = {
            'value': value,
            'time': time.time()
        }
        self.access_times[key] = 1
        return value

# 创建全局缓存实例
query_cache = QueryCache()

# 数据库操作装饰器 - 用于缓存结果
def cache_query(ttl=300):
    """
    缓存查询结果的装饰器
    :param ttl: 缓存有效期（秒）
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 生成缓存键
            cache_key = f"{func.__name__}:{str(args)}:{str(kwargs)}"
            
            # 尝试从缓存获取
            cached_result = query_cache.get(cache_key)
            if cached_result is not None:
                return cached_result
            
            # 执行查询
            result = func(*args, **kwargs)
            
            # 设置缓存
            query_cache.set(cache_key, result)
            return result
        return wrapper
    return decorator

# 数据库操作装饰器 - 用于异常处理
def db_operation(retries=3, retry_delay=0.5):
    """
    数据库操作的装饰器，提供异常处理和重试功能
    :param retries: 重试次数
    :param retry_delay: 重试延迟(秒)
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_error = None
            for attempt in range(retries):
                try:
                    return func(*args, **kwargs)
                except SQLAlchemyError as e:
                    last_error = e
                    current_app.logger.warning(f"数据库操作失败: {str(e)}，尝试重试 ({attempt+1}/{retries})")
                    
                    # 回滚会话
                    db.session.rollback()
                    
                    # 如果不是最后一次尝试，则等待后重试
                    if attempt < retries - 1:
                        time.sleep(retry_delay)
            
            # 所有重试都失败
            current_app.logger.error(f"数据库操作失败(已重试{retries}次): {str(last_error)}")
            raise last_error
        return wrapper
    return decorator

# 辅助函数：清除特定表相关的所有缓存
def clear_table_cache(table_name):
    """
    清除与特定表相关的所有缓存
    :param table_name: 表名
    """
    keys_to_delete = []
    for key in list(query_cache.cache.keys()):
        if table_name.lower() in key.lower():
            keys_to_delete.append(key)
    
    for key in keys_to_delete:
        if key in query_cache.cache:
            del query_cache.cache[key]
        if key in query_cache.access_times:
            del query_cache.access_times[key]

# 数据库连接重试机制
def retry_database_connection(db_instance, app, max_retries=10, retry_delay=5):
    """
    尝试连接数据库，在连接失败时进行重试
    :param db_instance: SQLAlchemy数据库实例
    :param app: Flask应用实例
    :param max_retries: 最大重试次数
    :param retry_delay: 重试间隔(秒)
    :return: 连接成功返回True，失败返回False
    """
    import sqlalchemy.exc as sa_exc

    for attempt in range(max_retries):
        try:
            app.logger.info(f"尝试连接数据库... (尝试 {attempt + 1}/{max_retries})")

            # 使用with语句确保应用上下文正确
            with app.app_context():
                # 尝试建立连接 - SQLAlchemy 2.0 兼容方式
                with db_instance.engine.connect() as connection:
                    connection.execute(text("SELECT 1"))
                    connection.commit()  # 确保事务提交

            app.logger.info("数据库连接成功！")
            return True

        except (sa_exc.OperationalError, sa_exc.DatabaseError, Exception) as e:
            app.logger.warning(f"数据库连接失败: {str(e)}")

            if attempt < max_retries - 1:
                app.logger.info(f"等待 {retry_delay} 秒后重试...")
                time.sleep(retry_delay)
            else:
                app.logger.error(f"数据库连接失败，已重试 {max_retries} 次，放弃连接")
                return False

    return False

# 通用查询函数，包含异常处理和缓存
@db_operation()
@cache_query()
def get_all_distinct_values(model, column_name):
    """
    获取指定模型指定列的所有不同值
    :param model: 模型类
    :param column_name: 列名
    :return: 所有不同值的列表
    """
    column = getattr(model, column_name)
    results = db.session.query(column).distinct().all()
    return [result[0] for result in results if result[0]]