from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from pydantic import BaseModel
import pyotp
import qrcode
import io
import base64
from app.core.database import get_db
from app.core.deps import get_current_user
from app.models.models import User

router = APIRouter(prefix="/auth/2fa", tags=["双因素认证"])


class TwoFASetupResponse(BaseModel):
    secret: str
    qr_code: str  # base64 encoded QR code image
    otpauth_url: str


class TwoFAVerifyRequest(BaseModel):
    code: str


class TwoFALoginRequest(BaseModel):
    temp_token: str
    code: str


@router.post("/setup", response_model=TwoFASetupResponse)
async def setup_2fa(
    user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """生成 2FA 密钥和二维码"""
    secret = pyotp.random_base32()
    
    totp = pyotp.TOTP(secret)
    otpauth_url = totp.provisioning_uri(
        name=user.email,
        issuer_name="债务管理系统"
    )
    
    # 生成 QR 码
    qr = qrcode.QRCode(version=1, box_size=10, border=5)
    qr.add_data(otpauth_url)
    qr.make(fit=True)
    img = qr.make_image(fill_color="black", back_color="white")
    
    # 转换为 base64
    buffer = io.BytesIO()
    img.save(buffer, format="PNG")
    qr_base64 = base64.b64encode(buffer.getvalue()).decode()
    
    # 临时存储 secret（验证后才会真正保存）
    user.two_factor_secret = secret
    await db.commit()
    
    return TwoFASetupResponse(
        secret=secret,
        qr_code=f"data:image/png;base64,{qr_base64}",
        otpauth_url=otpauth_url
    )


@router.post("/verify")
async def verify_2fa_setup(
    req: TwoFAVerifyRequest,
    user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """验证 2FA 设置（首次验证后启用）"""
    if not user.two_factor_secret:
        raise HTTPException(status_code=400, detail="请先执行 2FA 设置")
    
    totp = pyotp.TOTP(user.two_factor_secret)
    if not totp.verify(req.code, valid_window=1):
        raise HTTPException(status_code=400, detail="验证码错误")
    
    # 启用 2FA
    user.two_factor_enabled = 1
    await db.commit()
    
    return {"message": "双因素认证已启用"}


@router.post("/disable")
async def disable_2fa(
    req: TwoFAVerifyRequest,
    user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """禁用 2FA（需要验证当前验证码）"""
    if not user.two_factor_enabled:
        raise HTTPException(status_code=400, detail="未启用双因素认证")
    
    if not user.two_factor_secret:
        raise HTTPException(status_code=400, detail="2FA 配置异常")
    
    totp = pyotp.TOTP(user.two_factor_secret)
    if not totp.verify(req.code, valid_window=1):
        raise HTTPException(status_code=400, detail="验证码错误")
    
    user.two_factor_secret = None
    user.two_factor_enabled = 0
    await db.commit()
    
    return {"message": "双因素认证已禁用"}


@router.post("/verify-login")
async def verify_2fa_login(
    req: TwoFALoginRequest,
    db: AsyncSession = Depends(get_db),
):
    """验证 2FA 登录（密码验证后，2FA 验证前）"""
    from app.core.security import decode_token, create_access_token, create_refresh_token
    
    # 验证临时 token
    payload = decode_token(req.temp_token)
    if not payload or payload.get("type") != "2fa_pending":
        raise HTTPException(status_code=401, detail="无效的验证请求")
    
    user_id = int(payload["sub"])
    result = await db.execute(select(User).where(User.id == user_id))
    user = result.scalar_one_or_none()
    
    if not user or not user.two_factor_secret:
        raise HTTPException(status_code=401, detail="用户状态异常")
    
    # 验证 2FA
    totp = pyotp.TOTP(user.two_factor_secret)
    if not totp.verify(req.code, valid_window=1):
        raise HTTPException(status_code=401, detail="验证码错误")
    
    # 返回正式 token
    return {
        "access_token": create_access_token(user.id),
        "refresh_token": create_refresh_token(user.id),
        "token_type": "bearer"
    }


@router.get("/status")
async def get_2fa_status(
    user: User = Depends(get_current_user),
):
    """获取 2FA 状态"""
    return {
        "enabled": bool(user.two_factor_enabled),
        "has_secret": bool(user.two_factor_secret)
    }