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=["*"],
)

# Sửa phần khởi tạo Redis client và thêm biến để kiểm tra Redis status
redis_available = False
try:
    redis_client = redis.Redis(host='localhost', port=6379, db=0)
    redis_client.ping()  # Kiểm tra kết nối
    redis_available = True
except redis.ConnectionError:
    print("Redis không khả dụng - hệ thống sẽ chạy không có cache")
    redis_client = None

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")

    remaining_votes = voter['number_of_votes']
    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
        voter_collection.update_one(
            {'_id': ObjectId(vote_request.voter_id)}, 
            {'$inc': {'number_of_votes': -vote_request.vote_count}}
        )
        pdf_collection.update_one(
            {'_id': ObjectId(vote_request.file_id)}, 
            {'$inc': {'votes': vote_request.vote_count}}
        )
        
        # Chỉ xóa cache nếu Redis khả dụng
        if redis_available:
            try:
                redis_client.delete('all_files')
                redis_client.delete(f'voter:{vote_request.voter_id}')
            except redis.RedisError:
                print("Không thể xóa cache Redis")
        
        # 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:
        # Chỉ check cache nếu Redis khả dụng
        if redis_available:
            try:
                cached_voter = redis_client.get(f'voter:{id}')
                if cached_voter:
                    return json.loads(cached_voter)
            except redis.RedisError:
                print("Không thể đọc cache Redis")

        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']
        }
        
        # Chỉ lưu cache nếu Redis khả dụng
        if redis_available:
            try:
                redis_client.setex(f'voter:{id}', CACHE_EXPIRE_TIME, json.dumps(voter_data))
            except redis.RedisError:
                print("Không thể lưu cache Redis")

        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 if available
    if redis_available and redis_client:
        redis_client.close()

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