Update main.py

main.py

@@ -1,4 +1,4 @@
-import os, io, uuid, json, wave
+import os, io, uuid, json, wave, math
 import aiohttp
 import aiofiles
 import asyncio
@@ -84,7 +84,6 @@ else:
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 Base = declarative_base()
 
-# SQLAlchemy Models
 # SQLAlchemy Models
 class UserModel(Base):
     __tablename__ = "users"
@@ -161,6 +160,25 @@ class Transcription(TranscriptionBase):
     class Config:
         from_attributes = True
 
+# Add this with other SQLAlchemy models
+class UserCreditModel(Base):
+    __tablename__ = "user_credits"
+    
+    id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(Integer, ForeignKey("users.id"))
+    minutes_used = Column(Integer, default=0)
+    minutes_quota = Column(Integer, default=10)  # Default to 10 minutes
+    last_updated = Column(String, default=lambda: datetime.utcnow().isoformat())
+
+class UserCredit(BaseModel):
+    user_id: int
+    minutes_used: int
+    minutes_quota: int
+    last_updated: str
+    
+    class Config:
+        from_attributes = True
+
 # Add these utilities
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
@@ -185,7 +203,7 @@ os.makedirs("uploads", exist_ok=True)
 
 # Mount static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
-app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
+app.mount("/uploads", StaticFiles(directory=UPLOAD_DIR), name="uploads")
 
 # Initialize database
 def init_db():
@@ -426,21 +444,6 @@ async def login_for_access_token(
             detail=f"Internal server error: {str(e)}"
         )
 
-# @app.post("/token", response_model=Token)
-# async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
-#     user = authenticate_user(db, form_data.username, form_data.password)
-#     if not user:
-#         raise HTTPException(
-#             status_code=401,
-#             detail="Incorrect username or password",
-#             headers={"WWW-Authenticate": "Bearer"},
-#         )
-#     access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-#     access_token = create_access_token(
-#         data={"sub": user.username}, expires_delta=access_token_expires
-#     )
-#     return {"access_token": access_token, "token_type": "bearer"}
-
 @app.post("/api/signup")
 async def signup(user: UserCreate, db: Session = Depends(get_db)):
     try:
@@ -482,17 +485,6 @@ async def signup(user: UserCreate, db: Session = Depends(get_db)):
             detail=f"An error occurred: {str(e)}"
         )
 
-# Modify the existing routes to require authentication
-@app.get("/api/me", response_model=User)
-async def read_users_me(current_user: User = Depends(get_current_user)):
-    return current_user
-
-@app.get("/", response_class=HTMLResponse)
-async def read_index():
-    async with aiofiles.open("index.html", mode="r") as f:
-        content = await f.read()
-    return HTMLResponse(content=content)
-
 # Helper function to split audio into chunks
 def split_audio_chunks(
     file_path: str, 
@@ -648,6 +640,244 @@ def combine_results(results: list, response_format: str):
     
     return final
 
+def get_user_upload_dir(user_id):
+    """Get the upload directory for a specific user"""
+    user_dir = os.path.join(UPLOAD_DIR, str(user_id))
+    os.makedirs(user_dir, exist_ok=True)
+    return user_dir
+
+@app.get("/api/me", response_model=User)
+async def read_users_me(current_user: User = Depends(get_current_user)):
+    return current_user
+
+@app.get("/", response_class=HTMLResponse)
+async def read_index():
+    async with aiofiles.open("index.html", mode="r") as f:
+        content = await f.read()
+    return HTMLResponse(content=content)
+
+@app.get("/api/credits")
+async def get_user_credits(
+    current_user: User = Depends(get_current_user),
+    db: Session = Depends(get_db)
+):
+    try:
+        # Get user credit information
+        user_credit = db.query(UserCreditModel).filter(
+            UserCreditModel.user_id == current_user.id
+        ).first()
+        
+        if not user_credit:
+            # Create new credit record if it doesn't exist
+            user_credit = UserCreditModel(
+                user_id=current_user.id,
+                minutes_used=0,
+                minutes_quota=10,
+                last_updated=datetime.utcnow().isoformat()
+            )
+            db.add(user_credit)
+            db.commit()
+            db.refresh(user_credit)
+        
+        return {
+            "user_id": user_credit.user_id,
+            "minutes_used": user_credit.minutes_used,
+            "minutes_quota": user_credit.minutes_quota,
+            "minutes_remaining": max(0, user_credit.minutes_quota - user_credit.minutes_used),
+            "last_updated": user_credit.last_updated
+        }
+    except Exception as e:
+        print(f"Error getting user credits: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/api/admin/credits")
+async def get_all_user_credits(
+    current_user: User = Depends(get_current_active_admin),
+    db: Session = Depends(get_db)
+):
+    try:
+        # Get all user credit information
+        user_credits = db.query(UserCreditModel).all()
+        
+        # Get usernames for each user_id
+        user_data = {}
+        for user in db.query(UserModel).all():
+            user_data[user.id] = user.username
+        
+        return [
+            {
+                "user_id": credit.user_id,
+                "username": user_data.get(credit.user_id, "Unknown"),
+                "minutes_used": credit.minutes_used,
+                "minutes_quota": credit.minutes_quota,
+                "minutes_remaining": max(0, credit.minutes_quota - credit.minutes_used),
+                "last_updated": credit.last_updated
+            }
+            for credit in user_credits
+        ]
+    except Exception as e:
+        print(f"Error getting all user credits: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.put("/api/admin/credits/{user_id}")
+async def update_user_credits(
+    user_id: int,
+    credit_update: dict,
+    current_user: User = Depends(get_current_active_admin),
+    db: Session = Depends(get_db)
+):
+    try:
+        # Validate input
+        if not any(field in credit_update for field in ["minutes_used", "minutes_quota"]):
+            raise HTTPException(status_code=400, detail="At least one of minutes_used or minutes_quota is required")
+        
+        # Get user credit record
+        user_credit = db.query(UserCreditModel).filter(
+            UserCreditModel.user_id == user_id
+        ).first()
+        
+        if not user_credit:
+            # Create new credit record if it doesn't exist
+            user = db.query(UserModel).filter(UserModel.id == user_id).first()
+            if not user:
+                raise HTTPException(status_code=404, detail="User not found")
+                
+            minutes_used = credit_update.get("minutes_used", 0)
+            minutes_quota = credit_update.get("minutes_quota", 60)  # Default quota
+            
+            if not isinstance(minutes_used, int) or minutes_used < 0:
+                raise HTTPException(status_code=400, detail="minutes_used must be a non-negative integer")
+            
+            if not isinstance(minutes_quota, int) or minutes_quota < 0:
+                raise HTTPException(status_code=400, detail="minutes_quota must be a non-negative integer")
+                
+            user_credit = UserCreditModel(
+                user_id=user_id,
+                minutes_used=minutes_used,
+                minutes_quota=minutes_quota,
+                last_updated=datetime.utcnow().isoformat()
+            )
+            db.add(user_credit)
+        else:
+            # Update existing credit record
+            if "minutes_used" in credit_update:
+                minutes_used = credit_update["minutes_used"]
+                if not isinstance(minutes_used, int) or minutes_used < 0:
+                    raise HTTPException(status_code=400, detail="minutes_used must be a non-negative integer")
+                user_credit.minutes_used = minutes_used
+                
+            if "minutes_quota" in credit_update:
+                minutes_quota = credit_update["minutes_quota"]
+                if not isinstance(minutes_quota, int) or minutes_quota < 0:
+                    raise HTTPException(status_code=400, detail="minutes_quota must be a non-negative integer")
+                user_credit.minutes_quota = minutes_quota
+                
+            user_credit.last_updated = datetime.utcnow().isoformat()
+        
+        db.commit()
+        db.refresh(user_credit)
+        
+        # Get username for response
+        user = db.query(UserModel).filter(UserModel.id == user_id).first()
+        username = user.username if user else "Unknown"
+        
+        return {
+            "user_id": user_credit.user_id,
+            "username": username,
+            "minutes_used": user_credit.minutes_used,
+            "minutes_quota": user_credit.minutes_quota,
+            "minutes_remaining": max(0, user_credit.minutes_quota - user_credit.minutes_used),
+            "last_updated": user_credit.last_updated
+        }
+    except Exception as e:
+        print(f"Error updating user credits: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/admin/credits/{user_id}/reset-quota")
+async def reset_user_quota(
+    user_id: int,
+    quota_data: dict,
+    current_user: User = Depends(get_current_active_admin),
+    db: Session = Depends(get_db)
+):
+    try:
+        new_quota = quota_data.get("new_quota")
+        if new_quota is None:
+            raise HTTPException(status_code=400, detail="new_quota field is required")
+            
+        if not isinstance(new_quota, int) or new_quota < 0:
+            raise HTTPException(status_code=400, detail="new_quota must be a non-negative integer")
+        
+        # Get user credit record
+        user_credit = db.query(UserCreditModel).filter(
+            UserCreditModel.user_id == user_id
+        ).first()
+        
+        if not user_credit:
+            # Create new credit record if it doesn't exist
+            user = db.query(UserModel).filter(UserModel.id == user_id).first()
+            if not user:
+                raise HTTPException(status_code=404, detail="User not found")
+                
+            user_credit = UserCreditModel(
+                user_id=user_id,
+                minutes_used=0,
+                minutes_quota=new_quota,
+                last_updated=datetime.utcnow().isoformat()
+            )
+            db.add(user_credit)
+        else:
+            # Update existing credit record
+            user_credit.minutes_quota = new_quota
+            user_credit.last_updated = datetime.utcnow().isoformat()
+        
+        db.commit()
+        db.refresh(user_credit)
+        
+        # Get username for response
+        user = db.query(UserModel).filter(UserModel.id == user_id).first()
+        username = user.username if user else "Unknown"
+        
+        return {
+            "user_id": user_credit.user_id,
+            "username": username,
+            "minutes_used": user_credit.minutes_used,
+            "minutes_quota": user_credit.minutes_quota,
+            "minutes_remaining": max(0, user_credit.minutes_quota - user_credit.minutes_used),
+            "last_updated": user_credit.last_updated
+        }
+    except Exception as e:
+        print(f"Error resetting user quota: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/upload-audio")
+async def upload_audio(
+    request: Request, 
+    current_user: User = Depends(get_current_user),
+    db: Session = Depends(get_db)
+):
+    try:
+        form_data = await request.form()
+        file = form_data["file"]
+        
+        file_id = str(uuid.uuid4())
+        filename = f"{file_id}_{file.filename}"
+        
+        # Get user-specific upload directory
+        user_upload_dir = get_user_upload_dir(current_user.id)
+        
+        # Create the full file path in the user's directory
+        file_path = os.path.join(user_upload_dir, filename)
+        print(f"file_path: ${file_path}")
+        
+        contents = await file.read()
+        async with aiofiles.open(file_path, "wb") as f:
+            await f.write(contents)
+            
+        return JSONResponse({"filename": filename})
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.post("/api/upload")
 async def create_upload_file(
     current_user: User = Depends(get_current_user),
@@ -696,6 +926,40 @@ async def create_upload_file(
         if start_sample >= end_sample:
             raise HTTPException(400, "Invalid time range")
 
+        # Calculate audio duration in minutes
+        audio_duration_seconds = (end_sample - start_sample) / sample_rate
+        audio_duration_minutes = math.ceil(audio_duration_seconds / 60)  # Round up to nearest minute
+
+        # Check user quota
+        user_credit = db.query(UserCreditModel).filter(
+            UserCreditModel.user_id == current_user.id
+        ).first()
+        
+        if not user_credit:
+            # Create new credit record if it doesn't exist
+            user_credit = UserCreditModel(
+                user_id=current_user.id,
+                minutes_used=0,
+                minutes_quota=60,  # Default quota
+                last_updated=datetime.utcnow().isoformat()
+            )
+            db.add(user_credit)
+            db.commit()
+            db.refresh(user_credit)
+        
+        # Check if user has enough quota
+        minutes_remaining = max(0, user_credit.minutes_quota - user_credit.minutes_used)
+        if audio_duration_minutes > minutes_remaining:
+            raise HTTPException(
+                status_code=403, 
+                detail=f"Quota exceeded. You have {minutes_remaining} minutes remaining, but this audio requires {audio_duration_minutes} minutes."
+            )
+
+        # Update user credits
+        user_credit.minutes_used += audio_duration_minutes
+        user_credit.last_updated = datetime.utcnow().isoformat()
+        db.commit()        
+
         # Slice the audio data
         audio_data = audio_data[start_sample:end_sample]
         
@@ -703,6 +967,12 @@ async def create_upload_file(
         sliced_path = f"{temp_path}_sliced.flac"
         sf.write(sliced_path, audio_data, sample_rate, format='FLAC')
 
+        # check sliced_path file size if it exceeded the limit
+        file_size = os.path.getsize(sliced_path)
+        if file_size > MAX_FILE_SIZE:
+            print(f"Slice: {sliced_path} exceeds the limit of {MAX_FILE_SIZE / (1024 * 1024)}MB")
+            # raise HTTPException(400, f"File size exceeds the limit of {MAX_FILE_SIZE / (1024 * 1024)}MB")
+
         # Split into FLAC chunks
         chunks = split_audio_chunks(
             sliced_path, 
@@ -728,6 +998,14 @@ async def create_upload_file(
         # Combine results
         combined = combine_results(results, response_format)
 
+        # Add credit usage information to response
+        combined["metadata"]["credit_usage"] = {
+            "minutes_used": audio_duration_minutes,
+            "total_minutes_used": user_credit.minutes_used,
+            "minutes_quota": user_credit.minutes_quota,
+            "minutes_remaining": max(0, user_credit.minutes_quota - user_credit.minutes_used)
+        }
+
         # Cleanup
         for chunk in chunks:
             os.remove(chunk["path"])
@@ -744,29 +1022,6 @@ async def create_upload_file(
             os.remove(sliced_path)
         raise HTTPException(500, str(e))
 
-@app.post("/api/upload-audio")
-async def upload_audio(
-    request: Request, 
-    current_user: User = Depends(get_current_user),
-    db: Session = Depends(get_db)
-):
-    try:
-        form_data = await request.form()
-        file = form_data["file"]
-        
-        file_id = str(uuid.uuid4())
-        filename = f"{file_id}_{file.filename}"
-        file_path = os.path.join("uploads", filename)
-        print(f"file_path: ${file_path}")
-        
-        contents = await file.read()
-        async with aiofiles.open(file_path, "wb") as f:
-            await f.write(contents)
-            
-        return JSONResponse({"filename": filename})
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
 @app.post("/api/upload-audio-chunk")
 async def upload_audio_chunk(
     file: UploadFile = File(...),
@@ -812,8 +1067,11 @@ async def finalize_audio_upload(
     metadata_path = os.path.join(upload_path, "metadata.txt")
     filename = f"{uuid.uuid4()}.mp3"  # Default filename
     
+    # Get user-specific upload directory
+    user_upload_dir = get_user_upload_dir(current_user.id)
+    
     # Combine all chunks into final file
-    output_path = os.path.join(UPLOAD_DIR, filename)
+    output_path = os.path.join(user_upload_dir, filename)
     
     with open(output_path, "wb") as outfile:
         # Get all chunks sorted by offset
@@ -851,24 +1109,52 @@ async def get_audio_files(
     current_user: User = Depends(get_current_user)
 ):
     try:
-        # Scan uploads directory for audio files
-        upload_dir = os.path.join(os.getcwd(), UPLOAD_DIR)
         file_list = []
         
-        for filename in os.listdir(upload_dir):
-            file_path = os.path.join(upload_dir, filename)
-            if os.path.isfile(file_path) and not filename.startswith('.'):
-                # Get file stats
-                file_stats = os.stat(file_path)
-                file_size = os.path.getsize(file_path)
-                
-                file_list.append({
-                    "id": filename,  # Use filename as ID
-                    "filename": filename,
-                    "original_filename": filename,
-                    "size": file_size,
-                    "uploaded_at": datetime.fromtimestamp(file_stats.st_mtime).isoformat()
-                })
+        if current_user.is_admin:
+            # Admin can see all files
+            for user_id in os.listdir(UPLOAD_DIR):
+                user_dir = os.path.join(UPLOAD_DIR, user_id)
+                if os.path.isdir(user_dir):
+                    # Get username for this user_id
+                    db = next(get_db())
+                    user = db.query(UserModel).filter(UserModel.id == int(user_id)).first()
+                    username = user.username if user else f"User {user_id}"
+                    
+                    for filename in os.listdir(user_dir):
+                        file_path = os.path.join(user_dir, filename)
+                        if os.path.isfile(file_path) and not filename.startswith('.'):
+                            # Get file stats
+                            file_stats = os.stat(file_path)
+                            file_size = os.path.getsize(file_path)
+                            
+                            file_list.append({
+                                "id": filename,
+                                "filename": filename,
+                                "original_filename": filename,
+                                "size": file_size,
+                                "uploaded_at": datetime.fromtimestamp(file_stats.st_mtime).isoformat(),
+                                "user_id": int(user_id),
+                                "username": username
+                            })
+        else:
+            # Regular users can only see their own files
+            user_dir = get_user_upload_dir(current_user.id)
+            
+            for filename in os.listdir(user_dir):
+                file_path = os.path.join(user_dir, filename)
+                if os.path.isfile(file_path) and not filename.startswith('.'):
+                    # Get file stats
+                    file_stats = os.stat(file_path)
+                    file_size = os.path.getsize(file_path)
+                    
+                    file_list.append({
+                        "id": filename,
+                        "filename": filename,
+                        "original_filename": filename,
+                        "size": file_size,
+                        "uploaded_at": datetime.fromtimestamp(file_stats.st_mtime).isoformat()
+                    })
         
         # Sort by upload date (newest first)
         file_list.sort(key=lambda x: x["uploaded_at"], reverse=True)
@@ -881,11 +1167,25 @@ async def get_audio_files(
 @app.delete("/api/audio-files/{filename}")
 async def delete_audio_file(
     filename: str,
-    current_user: User = Depends(get_current_user)
+    user_id: int = None,
+    current_user: User = Depends(get_current_user),
+    db: Session = Depends(get_db)
 ):
     try:
-        # Check if file exists in uploads directory
-        file_path = os.path.join(os.getcwd(), UPLOAD_DIR, filename)
+        # Determine which user's file to delete
+        target_user_id = user_id if current_user.is_admin and user_id else current_user.id
+        
+        # Check if user exists if admin is deleting another user's file
+        if current_user.is_admin and user_id and user_id != current_user.id:
+            user = db.query(UserModel).filter(UserModel.id == user_id).first()
+            if not user:
+                raise HTTPException(status_code=404, detail="User not found")
+        
+        # Get the user's upload directory
+        user_upload_dir = get_user_upload_dir(target_user_id)
+        
+        # Check if file exists in the user's directory
+        file_path = os.path.join(user_upload_dir, filename)
         if not os.path.exists(file_path):
             raise HTTPException(status_code=404, detail="Audio file not found")