set/main.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import jwt
import json
from datetime import datetime, timedelta
from typing import List, Dict, Optional
import asyncio

from models import *
from game_logic import SetGame, Card
from database import (
    users_db, games_db, active_connections,
    get_user_by_nickname, create_user, get_user_by_token,
    create_game, get_game, get_all_games, update_game,
    add_player_to_game, remove_player_from_game,
    broadcast_game_state
)
from auth import (
    SECRET_KEY, ALGORITHM, create_access_token,
    verify_password, get_password_hash, verify_token
)
from websocket_manager import ConnectionManager

manager = ConnectionManager()

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Server starting...")
    yield
    print("Server shutting down...")
    await manager.disconnect_all()

app = FastAPI(
    title="Set Game Server",
    description="Сервер для игры Set с WebSocket поддержкой",
    version="1.0.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class BaseResponse:
    @staticmethod
    def success(data: Dict = None, **kwargs):
        response = {
            "success": True,
            "exception": None
        }
        if data:
            response.update(data)
        response.update(kwargs)
        return response
    
    @staticmethod
    def error(message: str):
        return {
            "success": False,
            "exception": {
                "message": message
            }
        }

# Эндпоинты

@app.post("/user/register", response_model=RegisterResponse)
async def register(user: RegisterRequest):

    if get_user_by_nickname(user.nickname):
        raise HTTPException(
            status_code=400,
            detail=BaseResponse.error("User already exists")
        )
    
    access_token = create_access_token(data={"sub": user.nickname})
    create_user(user.nickname, user.password, access_token)
    
    return BaseResponse.success(
        nickname=user.nickname,
        accessToken=access_token
    )

@app.post("/set/room/create")
async def create_game_room(request: GameRequest):
    
    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    game_id = create_game(request.accessToken)
    
    return BaseResponse.success(gameId=game_id)

@app.post("/set/room/list")
async def list_games(request: GameRequest):
    
    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    games = get_all_games()
    games_list = [{"id": game_id} for game_id in games.keys()]
    
    return BaseResponse.success(games=games_list)

@app.post("/set/room/enter")
async def enter_game(request: EnterGameRequest):
    
    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    game = get_game(request.gameId)
    if not game:
        raise HTTPException(
            status_code=404,
            detail=BaseResponse.error("Game not found")
        )
    
    success = add_player_to_game(request.gameId, request.accessToken, user["nickname"])
    if not success:
        raise HTTPException(
            status_code=400,
            detail=BaseResponse.error("Cannot join game")
        )
    
    await broadcast_game_state(request.gameId)
    
    return BaseResponse.success(gameId=request.gameId)

@app.post("/set/field")
async def get_game_field(request: GameRequest):

    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    game = None
    for game_id, game_data in games_db.items():
        if request.accessToken in game_data["players"]:
            game = game_data
            break
    
    if not game:
        raise HTTPException(
            status_code=400,
            detail=BaseResponse.error("User not in game")
        )
    
    field_cards = game["game"].get_field()
    score = game["players"].get(request.accessToken, {}).get("score", 0)
    status = "ongoing" if not game["game"].is_game_over() else "ended"
    
    return {
        "cards": field_cards,
        "status": status,
        "score": score
    }

@app.post("/set/pick")
async def pick_cards(request: PickRequest):

    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    game = None
    game_id = None
    for g_id, game_data in games_db.items():
        if request.accessToken in game_data["players"]:
            game = game_data
            game_id = g_id
            break
    
    if not game:
        raise HTTPException(
            status_code=400,
            detail=BaseResponse.error("User not in game")
        )
    
    is_set = game["game"].check_set(request.cards)
    
    if is_set:
        game["game"].remove_cards(request.cards)
        
        current_score = game["players"][request.accessToken].get("score", 0)
        game["players"][request.accessToken]["score"] = current_score + 1
        
        update_game(game_id, game)
        
        await broadcast_game_state(game_id)
    
    score = game["players"][request.accessToken].get("score", 0)
    
    return {
        "isSet": is_set,
        "score": score
    }

@app.post("/set/add")
async def add_cards(request: GameRequest):

    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    game = None
    game_id = None
    for g_id, game_data in games_db.items():
        if request.accessToken in game_data["players"]:
            game = game_data
            game_id = g_id
            break
    
    if not game:
        raise HTTPException(
            status_code=400,
            detail=BaseResponse.error("User not in game")
        )
    
    game["game"].add_cards(3) 
    
    update_game(game_id, game)
    
    await broadcast_game_state(game_id)
    
    return BaseResponse.success()

@app.post("/set/scores")
async def get_scores(request: GameRequest):

    user = get_user_by_token(request.accessToken)
    if not user:
        raise HTTPException(
            status_code=401,
            detail=BaseResponse.error("Invalid token")
        )
    
    game = None
    for game_data in games_db.values():
        if request.accessToken in game_data["players"]:
            game = game_data
            break
    
    if not game:
        raise HTTPException(
            status_code=400,
            detail=BaseResponse.error("User not in game")
        )
    
    users = []
    for player_data in game["players"].values():
        users.append({
            "name": player_data["nickname"],
            "score": player_data.get("score", 0)
        })
    
    return BaseResponse.success(users=users)

@app.websocket("/set")
async def websocket_endpoint(websocket: WebSocket):
    
    await manager.connect(websocket)
    
    try:
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)
            
            access_token = message.get("accessToken")
            if not access_token:
                await websocket.send_text(json.dumps(
                    BaseResponse.error("Token required")
                ))
                continue
            
            user = get_user_by_token(access_token)
            if not user:
                await websocket.send_text(json.dumps(
                    BaseResponse.error("Invalid token")
                ))
                continue
            
            game_id = None
            for g_id, game_data in games_db.items():
                if access_token in game_data["players"]:
                    game_id = g_id
                    break
            
            if not game_id:
                await websocket.send_text(json.dumps(
                    BaseResponse.error("User not in game")
                ))
                continue
            
            manager.register_user(websocket, access_token, game_id)
            
            game = games_db[game_id]
            field_cards = game["game"].get_field()
            status = "ongoing" if not game["game"].is_game_over() else "ended"
            score = game["players"][access_token].get("score", 0)
            
            await websocket.send_text(json.dumps({
                "type": "game_state",
                "cards": field_cards,
                "status": status,
                "score": score
            }))
            
    except WebSocketDisconnect:
        await manager.disconnect(websocket)
    except Exception as e:
        print(f"WebSocket error: {e}")
        await manager.disconnect(websocket)

@app.get("/health")
async def health_check():
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "users_count": len(users_db),
        "games_count": len(games_db)
    }

@app.get("/")
async def root():
    return {
        "message": "Set Game Server",
        "version": "1.0.0",
        "docs": "/docs",
        "endpoints": [
            "/user/register",
            "/set/room/create",
            "/set/room/list",
            "/set/room/enter",
            "/set/field",
            "/set/pick",
            "/set/add",
            "/set/scores",
            "/set (WebSocket)"
        ]
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)