from ninja import Router, Form
from django.contrib.auth import get_user_model
from rest_framework_simplejwt.tokens import RefreshToken
from django.db.models import Q

from invites.models import RegistrationCode

auth_router = Router(tags=["认证"])
User = get_user_model()


@auth_router.post("/register")
def register(
    request,
    username: str = Form(...),
    password: str = Form(...),
    email: str = Form(...),
    code: str = Form(None)
):
    if User.objects.filter(username=username).exists():
        return {"success": False, "message": "用户名已存在"}

    user = User(username=username, email=email, role="user")
    user.set_password(password)
    user.save()

    if code:
        try:
            reg = RegistrationCode.objects.get(code=code)
            if not reg.is_available():
                return {"success": False, "message": "注册码已达使用上限"}
            user.authorized_websites.set(reg.manager.managed_websites.all())
            reg.used_count += 1
            reg.save()
        except RegistrationCode.DoesNotExist:
            return {"success": False, "message": "注册码无效"}

    refresh = RefreshToken.for_user(user)

    return {
        "success": True,
        "message": "注册成功",
        "user": {
            "id": user.id,
            "username": user.username,
            "role": user.role,
        },
        "token": {
            "access": str(refresh.access_token),
            "refresh": str(refresh),
        }
    }

@auth_router.post("/login")
def login(
    request,
    username: str = Form(...),
    password: str = Form(...),
):
    user = User.objects.filter(Q(username=username) | Q(email=username)).first()

    if not user or not user.check_password(password):
        return {"success": False, "message": "用户名或密码错误"}

    if not user.is_active:
        return {"success": False, "message": "账号未激活"}

    refresh = RefreshToken.for_user(user)
    return {
        "success": True,
        "message": "登录成功",
        "user": {
            "id": user.id,
            "username": user.username,
            "role": user.role,
        },
        "token": {
            "access": str(refresh.access_token),
            "refresh": str(refresh),
        }
    }