from io import BytesIO
from captcha.image import ImageCaptcha
from uuid import uuid4
from utils.redis_client import get_redis, set_token_in_redis
import redis.asyncio as redis
import random
import string


def generate_captcha(length: int = 4) -> tuple[str, bytes]:
    """生成验证码图片和验证码文本"""
    # 生成随机验证码
    characters = string.digits + string.ascii_uppercase
    captcha_text = ''.join(random.choices(characters, k=length))
    
    # 生成验证码图片
    image = ImageCaptcha(width=150, height=50)
    image_bytes = image.generate(captcha_text)
    
    # 将图片转换为字节流
    buffer = BytesIO()
    buffer.write(image_bytes.read())
    buffer.seek(0)
    
    return captcha_text, buffer.getvalue()


async def save_captcha_to_redis(redis_conn: redis.Redis, captcha_id: str, captcha_text: str, expire_seconds: int = 300):
    """将验证码保存到Redis"""
    await redis_conn.set(f"captcha:{captcha_id}", captcha_text, ex=expire_seconds)


async def get_captcha_from_redis(redis_conn: redis.Redis, captcha_id: str) -> str | None:
    """从Redis获取验证码"""
    return await redis_conn.get(f"captcha:{captcha_id}")


async def delete_captcha_from_redis(redis_conn: redis.Redis, captcha_id: str):
    """从Redis删除验证码"""
    await redis_conn.delete(f"captcha:{captcha_id}")


async def validate_captcha(redis_conn: redis.Redis, captcha_id: str, user_input: str) -> bool:
    """验证验证码是否正确"""
    stored_captcha = await get_captcha_from_redis(redis_conn, captcha_id)
    if not stored_captcha:
        return False
    
    # 忽略大小写比较
    is_valid = stored_captcha.upper() == user_input.upper()
    
    # 无论验证成功与否，都删除验证码（防止重复使用）
    await delete_captcha_from_redis(redis_conn, captcha_id)
    
    return is_valid


async def create_captcha(redis_conn: redis.Redis) -> tuple[str, bytes]:
    """创建验证码并保存到Redis"""
    # 生成验证码
    captcha_text, image_bytes = generate_captcha()
    
    # 生成唯一ID
    captcha_id = str(uuid4())
    
    # 保存到Redis，有效期5分钟
    await save_captcha_to_redis(redis_conn, captcha_id, captcha_text, 300)
    
    return captcha_id, image_bytes