main.py

import os, time
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Response, Request
from fastapi.staticfiles import StaticFiles
from authlib.integrations.starlette_client import OAuth
from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import RedirectResponse, JSONResponse
from fastapi.responses import FileResponse
from fastapi import Header
from fastapi import Query
import jwt
import datetime
import httpx

JWT_SECRET = os.environ.get("JWT_SECRET")

app = FastAPI()

# 环境变量 (必须在 HF Space Settings 中配置)
CLIENT_ID = os.environ.get("ROBLOX_CLIENT_ID")
CLIENT_SECRET = os.environ.get("ROBLOX_CLIENT_SECRET")
SESSION_KEY = os.environ.get("SESSION_SECRET", "tony_safe_2026")
WHITE_LIST = os.environ.get("WHITE_LIST", "").split(",")

app.add_middleware(SessionMiddleware, secret_key=SESSION_KEY)

oauth = OAuth()
oauth.register(
    name='roblox',
    client_id=CLIENT_ID,
    client_secret=CLIENT_SECRET,
    server_metadata_url='https://apis.roblox.com/oauth/.well-known/openid-configuration',
    client_kwargs={'scope': 'openid profile'}
)

# 全局流状态
current_batch = b""
last_update_ts = 0
packet_id_counter = 0  # 💡 物理消除“闪回”的全局序号

# --- 身份验证接口 ---
@app.get('/login')
async def login(request: Request):
    redirect_uri = str(request.url_for('auth')).replace("http://", "https://")
    return await oauth.roblox.authorize_redirect(request, redirect_uri)


@app.get('/auth')
async def auth(request: Request):
    code = request.query_params.get("code")

    if not code:
        return JSONResponse({"error": "missing code"}, status_code=400)

    async with httpx.AsyncClient() as client:
        token_resp = await client.post(
            "https://apis.roblox.com/oauth/v1/token",
            data={
                "grant_type": "authorization_code",
                "code": code,
                "client_id": CLIENT_ID,
                "client_secret": CLIENT_SECRET,
                "redirect_uri": "https://tonyd365-tonys-multifunctional-hall.hf.space/auth"
            },
            headers={
                "Content-Type": "application/x-www-form-urlencoded"
            }
        )

    token_data = token_resp.json()

    access_token = token_data.get("access_token")

    if not access_token:
        return JSONResponse(token_data, status_code=400)

    # 获取用户信息
    async with httpx.AsyncClient() as client:
        user_resp = await client.get(
            "https://apis.roblox.com/oauth/v1/userinfo",
            headers={"Authorization": f"Bearer {access_token}"}
        )

    userinfo = user_resp.json()

    if not userinfo:
        return JSONResponse({"error": "userinfo missing"}, status_code=400)

    user_data = {
        'id': str(userinfo.get('sub')),
        'name': userinfo.get('preferred_username', 'Roblox Player'),
        'picture': userinfo.get('picture', ''),
        'exp': int((datetime.datetime.utcnow() + datetime.timedelta(hours=2)).timestamp())
    }

    jwt_token = jwt.encode(user_data, JWT_SECRET, algorithm="HS256")

    if isinstance(jwt_token, bytes):
        jwt_token = jwt_token.decode('utf-8')

    return RedirectResponse(
        url=f"https://tonys-virtual-stage-display.pages.dev/?token={jwt_token}"
    )

@app.get('/logout')
async def logout(request: Request):
    request.session.pop('user', None)
    return RedirectResponse(url='https://tonys-virtual-stage-display.pages.dev/')

@app.get('/api/user')
async def get_user_status(authorization: str = Header(None)):
    if not authorization:
        return JSONResponse({"logged_in": False})

    try:
        token = authorization.replace("Bearer ", "")
        user = jwt.decode(token, JWT_SECRET, algorithms=["HS256"])

        return JSONResponse({
            "logged_in": True,
            "user": user,
            "is_whitelisted": True if user["id"] in WHITE_LIST else False
        })
    except:
        return JSONResponse({"logged_in": False})

# --- 推流 WebSocket ---
@app.websocket("/ws/producer")
async def websocket_endpoint(websocket: WebSocket, token: str = Query(None)):
    if not token:
        await websocket.close(code=1008)
        return
    try:
        user = jwt.decode(token, JWT_SECRET, algorithms=["HS256"])
    except:
        await websocket.close(code=1008)
        return

    if user['id'] not in WHITE_LIST:
        await websocket.close(code=1008)
        return

    await websocket.accept()
    global current_batch, last_update_ts, packet_id_counter
    try:
        while True:
            data = await websocket.receive_bytes()
            if data:
                current_batch = data
                last_update_ts = time.time()
                packet_id_counter += 1 # 💡 关键：包序号递增
    except WebSocketDisconnect: pass

# --- Roblox 拉取接口 ---
@app.get("/poll")
async def poll():
    return Response(
        content=current_batch, 
        media_type="application/octet-stream",
        headers={
            "X-Timestamp": str(last_update_ts),
            "X-Packet-ID": str(packet_id_counter), # 💡 返回当前最新包序号
            "Cache-Control": "no-cache"
        }
    )

static_path = os.path.join(os.path.dirname(__file__), "static")
if not os.path.exists(static_path): os.makedirs(static_path)

#@app.get("/privacy")
#async def privacy():
#    return FileResponse(os.path.join(static_path,"privacy.html"))

#@app.get("/tos")
#async def privacy():
#    return FileResponse(os.path.join(static_path,"tos.html"))

#app.mount("/", StaticFiles(directory=static_path, html=True), name="static")

@app.get("/")
async def blankget(request: Request):
    return RedirectResponse(url='https://tonys-virtual-stage-display.pages.dev/')

@app.head("/")
async def HeadGet(request: Request):
    return JSONResponse({"status": "OK"})