app.py

from fastapi import FastAPI, HTTPException, Form, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi_socketio import SocketManager
from typing import Optional
import pymongo
from bson.objectid import ObjectId
import redis
import json
from redis.exceptions import RedisError
from pydantic import BaseModel, validator
from pymongo.errors import ConnectionFailure

# Initialize FastAPI app
app = FastAPI()
socket_manager = SocketManager(app=app, cors_allowed_origins="*", mount_location="/socket.io")

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

# Initialize Redis client
redis_client = redis.Redis(host='localhost', port=6379, db=0)
CACHE_EXPIRE_TIME = 300  # 5 minutes

# MongoDB connection
mongo_url = "mongodb+srv://ip6ofme:JL1S4hjMWRoko8AJ@cluster0.x0vo0.mongodb.net/"
client = pymongo.MongoClient(mongo_url)
db = client["test"]
pdf_collection = db["PdfDetails"]
voter_collection = db["Voters"]

# Kiểm tra kết nối MongoDB khi khởi động
try:
    # The ismaster command is cheap and does not require auth.
    client.admin.command('ismaster')
except ConnectionFailure:
    print("Server not available")
    raise

# Pydantic models for request validation
class VoterRegistration(BaseModel):
    name: str
    group: str
    role: str

class VoteRequest(BaseModel):
    voter_id: str
    file_id: str
    vote_count: Optional[int] = 1

    class Config:
        @validator('vote_count')
        def validate_vote_count(cls, v):
            if v <= 0:
                raise ValueError('Vote count must be greater than 0')
            return v

@app.post("/upload-files")
async def upload_file(
    title: str = Form(...),
    group: str = Form(...),
    url: str = Form(...)
):
    new_pdf = {
        'title': title,
        'group': group,
        'url': url,
        'votes': 0
    }
    result = pdf_collection.insert_one(new_pdf)
    return {'status': 'ok', 'id': str(result.inserted_id)}

@app.get("/get-votes")
async def get_votes(file_id: str):
    try:
        file = pdf_collection.find_one({'_id': ObjectId(file_id)})
        if not file:
            raise HTTPException(status_code=404, detail="File not found")
        return {'status': 'ok', 'votes': file.get('votes', 0)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/get-files")
async def get_files():
    try:
        files = pdf_collection.find({})
        file_list = [
            {
                'id': str(file['_id']),
                'title': file['title'],
                'group': file['group'],
                'url': file['url'],
                'votes': file.get('votes', 0)
            }
            for file in files
        ]
        return {'status': 'ok', 'data': file_list}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/register-voter")
async def register_voter(voter: VoterRegistration):
    # Xác định số lượng vote tối đa dựa vào role
    max_votes = 10 if voter.role == 'judge' else 2
    
    new_voter = {
        'name': voter.name,
        'group': voter.group,
        'role': voter.role,
        'number_of_votes': max_votes  # Khởi tạo với số vote tối đa
    }
    result = voter_collection.insert_one(new_voter)
    return {'status': 'ok', 'id': str(result.inserted_id)}

@app.post("/vote-by-voter")
async def vote_by_voter(vote_request: VoteRequest):
    if vote_request.vote_count <= 0:
        raise HTTPException(status_code=400, detail="Vote count must be greater than 0")

    voter = voter_collection.find_one({'_id': ObjectId(vote_request.voter_id)})
    if not voter:
        raise HTTPException(status_code=404, detail="Voter not found")

    # Kiểm tra số vote còn lại
    remaining_votes = voter['number_of_votes']  # Số vote còn lại
    if vote_request.vote_count > remaining_votes:
        raise HTTPException(status_code=400, detail=f"Not enough votes remaining. You have {remaining_votes} votes left")

    try:
        # Cập nhật MongoDB - giảm số vote còn lại
        voter_collection.update_one(
            {'_id': ObjectId(vote_request.voter_id)}, 
            {'$inc': {'number_of_votes': -vote_request.vote_count}}  # Giảm số vote còn lại
        )
        pdf_collection.update_one(
            {'_id': ObjectId(vote_request.file_id)}, 
            {'$inc': {'votes': vote_request.vote_count}}
        )
        
        # Clear cache
        redis_client.delete('all_files')
        redis_client.delete(f'voter:{vote_request.voter_id}')
        
        # Emit socket event
        updated_file = pdf_collection.find_one({'_id': ObjectId(vote_request.file_id)})
        await socket_manager.emit('vote_update', {
            'file_id': vote_request.file_id,
            'votes': updated_file.get('votes', 0)
        })
        
        return {'status': 'ok', 'message': f'Vote recorded successfully with {vote_request.vote_count} votes'}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/get-voter")
async def get_voter(id: str):
    try:
        # Check cache
        try:
            cached_voter = redis_client.get(f'voter:{id}')
            if cached_voter:
                return json.loads(cached_voter)
        except RedisError as redis_err:
            print(f"Redis error: {redis_err}")
            pass

        voter = voter_collection.find_one({'_id': ObjectId(id)})
        if not voter:
            raise HTTPException(status_code=404, detail="Voter not found")
        
        voter_data = {
            'status': 'ok',
            'name': voter['name'],
            'group': voter['group'],
            'role': voter['role'],
            'number_of_votes': voter['number_of_votes']  # Số này giờ là số vote còn lại
        }
        
        try:
            redis_client.setex(f'voter:{id}', CACHE_EXPIRE_TIME, json.dumps(voter_data))
        except RedisError as redis_err:
            print(f"Redis caching error: {redis_err}")
            pass

        return voter_data
    except Exception as e:
        print(f"Error in get_voter: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@socket_manager.on('connect')
async def handle_connect(sid, environ):
    print(f'Client connected: {sid}')

@socket_manager.on('disconnect')
async def handle_disconnect(sid):
    print(f'Client disconnected: {sid}')

@app.get("/")
async def index():
    return {'status': 'ok', 'message': 'Server is running'}

@app.on_event("shutdown")
async def shutdown_event():
    # Close MongoDB connection
    client.close()
    # Close Redis connection
    redis_client.close()

if _name_ == "_main_":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)