trade-assistant/backend/app/api/v1/auth.py

from fastapi import APIRouter, Depends, HTTPException, status, Header
from fastapi.security import OAuth2PasswordRequestForm
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from typing import Annotated, Optional
import uuid
from app.database import get_db
from app.models.user import User
from app.core.security import hash_password, verify_password, create_access_token, create_refresh_token, decode_token
from pydantic import BaseModel, EmailStr
from datetime import datetime, timedelta

router = APIRouter()


class RegisterRequest(BaseModel):
    phone: str
    password: str
    username: str = ""


class LoginResponse(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str = "bearer"
    user: dict


class RefreshRequest(BaseModel):
    refresh_token: str


@router.post("/register")
async def register(data: RegisterRequest, db: Annotated[AsyncSession, Depends(get_db)] = None):
    existing = await db.execute(select(User).where(User.phone == data.phone))
    if existing.scalar_one_or_none():
        raise HTTPException(status_code=400, detail="Phone already registered")

    user = User(
        phone=data.phone,
        username=data.username or data.phone,
        password_hash=hash_password(data.password),
        tier="free",
    )
    db.add(user)
    await db.flush()

    return {
        "id": str(user.id),
        "phone": user.phone,
        "username": user.username,
        "tier": user.tier,
        "role": user.role,
    }


@router.post("/login", response_model=LoginResponse)
async def login(
    form: Annotated[OAuth2PasswordRequestForm, Depends()],
    db: Annotated[AsyncSession, Depends(get_db)] = None,
):
    result = await db.execute(select(User).where(User.phone == form.username))
    user = result.scalar_one_or_none()

    if not user or not verify_password(form.password, user.password_hash):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid credentials",
        )

    return LoginResponse(
        access_token=create_access_token({"sub": str(user.id), "tier": user.tier, "role": user.role}),
        refresh_token=create_refresh_token({"sub": str(user.id)}),
        user={
            "id": str(user.id),
            "phone": user.phone,
            "username": user.username,
            "tier": user.tier,
        },
    )


@router.post("/login/guest")
async def guest_login():
    guest_id = f"guest_{uuid.uuid4().hex[:12]}"
    access_token = create_access_token(
        {"sub": guest_id, "tier": "guest", "role": "guest", "is_guest": True},
        expires_delta=timedelta(hours=24)
    )
    refresh_token = create_refresh_token({"sub": guest_id, "is_guest": True})
    
    return LoginResponse(
        access_token=access_token,
        refresh_token=refresh_token,
        token_type="bearer",
        user={
            "id": guest_id,
            "phone": None,
            "username": "游客用户",
            "tier": "guest",
            "is_guest": True,
        },
    )


@router.post("/refresh")
async def refresh(data: RefreshRequest):
    payload = decode_token(data.refresh_token)
    if not payload or payload.get("type") != "refresh":
        raise HTTPException(status_code=401, detail="Invalid refresh token")

    return {
        "access_token": create_access_token({"sub": payload["sub"]}),
        "token_type": "bearer",
    }


@router.get("/me")
async def get_me(
    authorization: Optional[str] = Header(None, alias="Authorization"),
    db: Annotated[AsyncSession, Depends(get_db)] = None,
):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing token")

    payload = decode_token(authorization[7:])
    if not payload:
        raise HTTPException(status_code=401, detail="Invalid token")

    result = await db.execute(select(User).where(User.id == payload["sub"]))
    user = result.scalar_one_or_none()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    return {
        "id": str(user.id),
        "phone": user.phone,
        "username": user.username,
        "tier": user.tier,
        "role": user.role,
        "settings": user.settings,
        "created_at": user.created_at.isoformat() if user.created_at else None,
    }


class SettingsUpdate(BaseModel):
    preferred_translate_provider: str = None
    reply_tone: str = None
    timezone: str = None
    languages: list = None


class WeChatLoginRequest(BaseModel):
    code: str
    encrypted_data: str = ""
    iv: str = ""


@router.post("/wechat-login")
async def wechat_login(data: WeChatLoginRequest, db: Annotated[AsyncSession, Depends(get_db)] = None):
    from app.services.wechat import wechat_service

    session = await wechat_service.code2session(data.code)
    if not session:
        raise HTTPException(status_code=400, detail="WeChat login failed")

    openid = session.get("openid")
    result = await db.execute(select(User).where(User.wechat_openid == openid))
    user = result.scalar_one_or_none()

    if not user:
        user = User(
            wechat_openid=openid,
            username=f"wx_{openid[-8:]}",
            tier="free",
        )
        db.add(user)
        await db.flush()

    return LoginResponse(
        access_token=create_access_token({"sub": str(user.id), "tier": user.tier, "role": user.role}),
        refresh_token=create_refresh_token({"sub": str(user.id)}),
        user={
            "id": str(user.id),
            "phone": user.phone,
            "username": user.username,
            "tier": user.tier,
            "role": user.role,
        },
    )


@router.patch("/settings")
async def update_settings(
    data: SettingsUpdate,
    authorization: Optional[str] = Header(None, alias="Authorization"),
    db: Annotated[AsyncSession, Depends(get_db)] = None,
):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing token")

    payload = decode_token(authorization[7:])
    if not payload:
        raise HTTPException(status_code=401, detail="Invalid token")

    result = await db.execute(select(User).where(User.id == payload["sub"]))
    user = result.scalar_one_or_none()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")

    settings = user.settings or {}
    for key, value in data.dict(exclude_unset=True).items():
        if value is not None:
            settings[key] = value

    user.settings = settings
    await db.flush()

    return {"settings": user.settings}