import redis
import json
import time
from config_manager import global_config


class RedisManager:
    def __init__(self):
        self.config = global_config.get_config()
        self.redis = redis.StrictRedis(
            host=self.config.REDIS_HOST,
            port=self.config.REDIS_PORT,
            db=self.config.REDIS_DB,
            password=self.config.REDIS_PASSWORD,
            decode_responses=True,
            max_connections=100,
            socket_timeout=5
        )

    def cache_group_data(self, group_id, data_type, data):
        """缓存群组相关数据"""
        key = f"{self.config.REDIS_KEY_PREFIX}:group:{group_id}:{data_type}"
        if isinstance(data, list):
            self.redis.delete(key)
            if data:
                self.redis.sadd(key, *data)
        else:
            self.redis.set(key, json.dumps(data))

    def get_group_data(self, group_id, data_type):
        """获取群组相关数据"""
        key = f"{self.config.REDIS_KEY_PREFIX}:group:{group_id}:{data_type}"
        data_type = data_type.lower()

        if data_type in ['users', 'sms_accounts']:
            return self.redis.smembers(key) or []
        else:
            data = self.redis.get(key)
            return json.loads(data) if data else None

    def add_user_message(self, user_id, role, content):
        """添加用户消息到Redis"""
        key = f"{self.config.REDIS_KEY_PREFIX}:user_messages:{user_id}:messages"
        message = {
            "role": role,
            "content": content,
            "timestamp": time.time()
        }
        self.redis.rpush(key, json.dumps(message))
        self.redis.expire(key, self.config.MESSAGE_EXPIRE)

    def add_conversation_id(self, user_id, conversation_id):
        """添加会话id"""
        key = f"{self.config.REDIS_KEY_PREFIX}:user_messages:{user_id}:conversation_id"
        self.redis.set(key, conversation_id)
        self.redis.expire(key, self.config.MESSAGE_EXPIRE)

    def get_user_messages(self, user_id):
        """获取用户的所有消息"""
        key = f"{self.config.REDIS_KEY_PREFIX}:user_messages:{user_id}:messages"
        messages = self.redis.lrange(key, 0, -1)
        return [json.loads(msg) for msg in messages] if messages else []

    def get_user_conversation_id(self, user_id):
        """获取用户的conversation_id"""
        key = f"{self.config.REDIS_KEY_PREFIX}:user_messages:{user_id}:conversation_id"
        return self.redis.get(key)

    def clear_user_messages(self, user_id):
        """清除用户的消息缓存"""
        key_prefix = f"{self.config.REDIS_KEY_PREFIX}:user_messages:{user_id}"
        key1 = f"{key_prefix}:messages"
        key2 = f"{key_prefix}:conversation_id"
        self.redis.delete(key1)
        self.redis.delete(key2)

    def add_stream(self, stream_id):
        """添加流列表"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}"
        self.redis.set(key, 0)
        self.redis.expire(key, 180)

    def add_stream_chunks(self, stream_id, chunks):
        """添加流列表"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:chunks"
        self.redis.lpush(key, *chunks)
        self.redis.expire(key, 180)

    def get_stream_chunk(self, stream_id):
        """获取流块"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:chunks"
        return self.redis.rpop(key)

    def add_stream_status(self, stream_id, status):
        """添加ai请求的状态"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:status"
        self.redis.set(key, status)
        self.redis.expire(key, 180)

    def get_stream_status(self, stream_id):
        """添加ai请求的状态"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:status"
        return self.redis.get(key)

    def add_stream_lock(self, stream_id, lock):
        """添加lock"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:lock"
        self.redis.set(key, lock)
        self.redis.expire(key, 180)

    def add_stream_conversation_id(self, stream_id, conversation_id):
        """添加ai请求的状态"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:conversation_id"
        self.redis.set(key, conversation_id)
        self.redis.expire(key, 180)

    def get_stream_conversation_id(self, stream_id):
        """添加ai请求的状态"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:conversation_id"
        return self.redis.get(key)

    def exists_stream(self, stream_id):
        """检查键是否存在"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:lock"
        return self.redis.exists(key)

    def exists_stream_chunks(self, stream_id):
        """检查键是否存在"""
        key = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:chunks"
        return self.redis.exists(key)

    def delete_stream_chunks(self, stream_id):
        """删除stream_id"""
        pattern = f"{self.config.REDIS_KEY_PREFIX}:stream:{stream_id}:*"
        count = 0
        for key in self.redis.scan_iter(match=pattern):
            self.redis.delete(key)
            count += 1

    def add_wecom_msgid(self, msgid):
        """添加企业微信回调msgid"""
        key = f"{self.config.REDIS_KEY_PREFIX}:msgid:{msgid}"
        self.redis.set(key, 0)
        self.redis.expire(key, self.config.MSGID_EXPIRE)

    def update_wecom_msgid(self, msgid, ai_answer):
        """更新企业微信回调msgid"""
        key = f"{self.config.REDIS_KEY_PREFIX}:msgid:{msgid}"
        self.redis.set(key, ai_answer)
        self.redis.expire(key, self.config.MSGID_EXPIRE)

    def delete_wecom_msgid(self, msgid):
        """删除企业微信回调msgid"""
        key = f"{self.config.REDIS_KEY_PREFIX}:msgid:{msgid}"
        self.redis.delete(key)

    def get_wecom_msgid(self, msgid):
        """获取msgid"""
        key = f"{self.config.REDIS_KEY_PREFIX}:msgid:{msgid}"
        return self.redis.get(key)

    def exists_wecom_msgid(self, msgid):
        """msgid是否存在"""
        key = f"{self.config.REDIS_KEY_PREFIX}:msgid:{msgid}"
        return self.redis.exists(key)