SonicForgeStudio/app/api/v1/auth.py

import uuid
import time
from fastapi import APIRouter, HTTPException, Header, Depends
from pydantic import BaseModel, EmailStr
from typing import Optional
from app.models.user import get_db_connection
from app.core.auth import hash_password, verify_password, create_token, decode_token, seed_admin

router = APIRouter()

class LoginRequest(BaseModel):
    username: Optional[str] = "admin"
    password: str

class RegisterRequest(BaseModel):
    username: str
    email: str
    password: str

class ChangePasswordRequest(BaseModel):
    old_password: str
    new_password: str

def get_current_user(authorization: Optional[str] = Header(None)):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Thiếu Token xác thực hoặc Token không hợp lệ")
    token = authorization.split(" ")[1]
    payload = decode_token(token)
    if not payload:
        raise HTTPException(status_code=401, detail="Token đã hết hạn hoặc không hợp lệ")
    return payload

@router.post("/login")
async def login(req: LoginRequest):
    conn = get_db_connection()
    cursor = conn.cursor()
    
    username = (req.username or "").strip()
    if not username:
        username = "admin"
        
    password = (req.password or "").strip()
    
    # Case-insensitive search by username or email
    cursor.execute("SELECT * FROM users WHERE LOWER(username) = LOWER(?) OR LOWER(email) = LOWER(?)", (username, username))
    user = cursor.fetchone()
    
    # Auto-heal seed_admin if admin record missing
    if not user and username.lower() == "admin":
        conn.close()
        seed_admin()
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM users WHERE username = 'admin'")
        user = cursor.fetchone()

    conn.close()
    
    if not user or not user["is_active"]:
        raise HTTPException(status_code=400, detail="Tài khoản hoặc mật khẩu không chính xác")
        
    if not verify_password(password, user["hashed_password"]):
        raise HTTPException(status_code=400, detail="Tài khoản hoặc mật khẩu không chính xác")
        
    token = create_token(user["id"], user["username"], user["role"], user["must_change_password"])
    
    return {
        "access_token": token,
        "user": {
            "id": user["id"],
            "username": user["username"],
            "email": user["email"],
            "role": user["role"],
            "must_change_password": bool(user["must_change_password"])
        }
    }

@router.post("/register")
async def register(req: RegisterRequest):
    username = req.username.strip()
    email = req.email.strip()
    password = req.password.strip()
    
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute("SELECT id FROM users WHERE LOWER(username) = LOWER(?) OR LOWER(email) = LOWER(?)", (username, email))
    if cursor.fetchone():
        conn.close()
        raise HTTPException(status_code=400, detail="Tên người dùng hoặc Email đã tồn tại")
        
    user_id = str(uuid.uuid4())
    hashed_pwd = hash_password(password)
    now = time.time()
    
    cursor.execute("""
    INSERT INTO users (id, username, email, hashed_password, role, must_change_password, created_at, is_active)
    VALUES (?, ?, ?, ?, 'standard', 0, ?, 1)
    """, (user_id, username, email, hashed_pwd, now))
    
    cursor.execute("""
    INSERT INTO user_quotas (user_id, storage_limit_mb, max_tracks)
    VALUES (?, 500, 16)
    """, (user_id,))
    
    conn.commit()
    conn.close()
    
    token = create_token(user_id, username, "standard", False)
    return {
        "access_token": token,
        "user": {
            "id": user_id,
            "username": username,
            "email": email,
            "role": "standard",
            "must_change_password": False
        }
    }

@router.post("/change-password")
async def change_password(req: ChangePasswordRequest, current_user: dict = Depends(get_current_user)):
    user_id = current_user["user_id"]
    old_pwd = req.old_password.strip()
    new_pwd = req.new_password.strip()
    
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute("SELECT hashed_password FROM users WHERE id = ?", (user_id,))
    user = cursor.fetchone()
    if not user or not verify_password(old_pwd, user["hashed_password"]):
        conn.close()
        raise HTTPException(status_code=400, detail="Mật khẩu hiện tại không chính xác")
        
    new_hashed = hash_password(new_pwd)
    cursor.execute("""
    UPDATE users SET hashed_password = ?, must_change_password = 0 WHERE id = ?
    """, (new_hashed, user_id))
    
    conn.commit()
    
    cursor.execute("SELECT * FROM users WHERE id = ?", (user_id,))
    updated_user = cursor.fetchone()
    conn.close()
    
    new_token = create_token(updated_user["id"], updated_user["username"], updated_user["role"], False)
    return {
        "message": "Đổi mật khẩu thành công!",
        "access_token": new_token
    }

@router.get("/profile")
async def get_profile(current_user: dict = Depends(get_current_user)):
    user_id = current_user["user_id"]
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute("""
    SELECT u.id, u.username, u.email, u.role, u.must_change_password, q.storage_limit_mb, q.max_tracks
    FROM users u
    LEFT JOIN user_quotas q ON u.id = q.user_id
    WHERE u.id = ?
    """, (user_id,))
    row = cursor.fetchone()
    
    cursor.execute("SELECT SUM(size_bytes) as total_used FROM projects WHERE user_id = ?", (user_id,))
    used_row = cursor.fetchone()
    used_bytes = used_row["total_used"] if used_row and used_row["total_used"] else 0
    used_mb = round(used_bytes / (1024 * 1024), 2)
    
    conn.close()
    
    if not row:
        raise HTTPException(status_code=404, detail="Không tìm thấy thông tin tài khoản")
        
    return {
        "id": row["id"],
        "username": row["username"],
        "email": row["email"],
        "role": row["role"],
        "must_change_password": bool(row["must_change_password"]),
        "quota": {
            "storage_limit_mb": row["storage_limit_mb"] or 500,
            "used_mb": used_mb,
            "max_tracks": row["max_tracks"] or 16
        }
    }