main.py

import os
import io
import jwt
import uuid
import bcrypt
import zipfile
import psycopg2
import tempfile
from PIL import Image
from minio import Minio
from functools import wraps
from urllib.parse import quote
from dotenv import load_dotenv
from minio.error import S3Error
from typing import List, Optional
from psycopg2.extras import DictCursor
from datetime import datetime, timedelta, timezone
from fastapi.responses import StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from model import UserRegister, UserLogin, DownloadRequest, TokenData
from fastapi import FastAPI, HTTPException, UploadFile, File, status, Depends, Security, Response, Query


load_dotenv()

app = FastAPI()

security = HTTPBearer()

SECRET_KEY = os.getenv('SECRET_KEY')
ALGORITHM = os.getenv('ALGORITHM')
ACCESS_TOKEN_EXPIRE_MINUTES = 3
REFRESH_TOKEN_EXPIRE_DAYS = 1

minio_client = Minio(
    "minio:9000",
    access_key=os.getenv('MINIO_ROOT_USER'),
    secret_key=os.getenv('MINIO_ROOT_PASSWORD'),
    secure=False
)

db_conn = psycopg2.connect(
    dbname=os.getenv('POSTGRES_DB'),
    user=os.getenv('POSTGRES_USER'),
    password=os.getenv('POSTGRES_PASSWORD'),
    host="postgres",
    port="5432"
)

def create_jwt(data: dict, token_type: str = "access") -> str:
    to_encode = data.copy()
    expire = datetime.now(timezone.utc) + (timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES) if token_type == "access" else timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

def verify_token(token: str, credentials_exception):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        email = payload.get("sub")
        if not email:
            raise credentials_exception
        return TokenData(email=email)
    except jwt.PyJWTError:
        raise credentials_exception

def get_file_category(filename: str) -> str:
    match os.path.splitext(filename)[1].lower():
        case '.jpg' | '.png' | '.gif' | '.bmp' | '.tiff' | '.webp' | '.svg' | '.raw' | '.ico' | '.heic' | '.jpeg':
            return "Pictures/"
        case '.mp4' | '.avi' | '.mkv' | '.mov' | '.wmv' | '.flv' | '.webm' | '.mpeg' | '.3gp' | '.vob':
            return "Movies/"
        case '.mp3' | '.wav' | '.flac' | '.aac' | '.ogg' | '.wma' | '.m4a' | '.opus' | '.alac' | '.aiff':
            return "Music/"
        case _:
            return "Documents/"

def s3_handler(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        bucket_name = kwargs.get("bucket")
        current_user = kwargs.get("current_user")
        if bucket_name and not minio_client.bucket_exists(bucket_name):
            raise HTTPException(status_code=404, detail=f"Bucket '{bucket_name}' does not exist.")
        if bucket_name != current_user:
            raise HTTPException(status_code=403, detail="Access denied")
        try:
            return await func(*args, **kwargs)
        except S3Error as exc:
            raise HTTPException(status_code=500, detail=str(exc))
    return wrapper

@app.post("/sign_up/")
def register_user(user: UserRegister):
    cursor = db_conn.cursor()
    cursor.execute("SELECT * FROM users WHERE email = %s", (user.email,))
    if cursor.fetchone():
        raise HTTPException(status_code=400, detail="This user already exists")
    user_uuid = str(uuid.uuid4())
    cursor.execute("INSERT INTO users (email, password, uuid) VALUES (%s, %s, %s)", (user.email, user.bcrypt12rounds_password, user_uuid))
    db_conn.commit()
    cursor.close()
    if not minio_client.bucket_exists(user_uuid):
        minio_client.make_bucket(user_uuid)
    for folder in ["Pictures/", "Movies/", "Music/", "Documents/"]:
            minio_client.put_object(user_uuid, folder, io.BytesIO(b""), 0)
    return {"message": f"User registered successfully!", "bucket": user_uuid}

@app.post("/login")
def login(user: UserLogin):
    cursor = db_conn.cursor(cursor_factory=DictCursor)
    cursor.execute("SELECT * FROM users WHERE email = %s", (user.email,))
    db_user = cursor.fetchone()
    if not db_user or not bcrypt.checkpw(user.password.encode('utf-8'), db_user["password"].encode('utf-8')):
        raise HTTPException(status_code=400, detail="Incorrect email or password")
    access_token = create_jwt(data={"sub": user.email}, token_type="access")
    refresh_token = create_jwt(data={"sub": user.email}, token_type="refresh")
    cursor.execute("UPDATE users SET refresh_token = %s, expiration_date = %s WHERE email = %s", (refresh_token, datetime.now(timezone.utc) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS), user.email))
    db_conn.commit()
    cursor.close()
    return {"access_token": access_token, "token_type": "bearer", "refresh_token": refresh_token, "bucket": db_user["uuid"]}

@app.post("/refresh")
def refresh_token(refresh_token: str):
    credentials_exception = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"})
    token_data = verify_token(refresh_token, credentials_exception)
    cursor = db_conn.cursor(cursor_factory=DictCursor)
    cursor.execute("SELECT * FROM users WHERE email = %s", (token_data.email,))
    db_user = cursor.fetchone()
    if not db_user or db_user["refresh_token"] != refresh_token:
        raise credentials_exception
    if datetime.now(timezone.utc) > db_user["expiration_date"]:
        raise HTTPException(status_code=400, detail="Refresh token expired")
    return {"access_token": create_jwt(data={"sub": token_data.email}, token_type="access"), "token_type": "bearer", "bucket": db_user["uuid"]}

async def get_current_user(credentials: HTTPAuthorizationCredentials = Security(security)):
    credentials_exception = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"})
    token = credentials.credentials
    token_data = verify_token(token, credentials_exception)
    cursor = db_conn.cursor(cursor_factory=DictCursor)
    cursor.execute("SELECT uuid FROM users WHERE email = %s", (token_data.email,))
    user = cursor.fetchone()
    if not user:
        raise credentials_exception
    return user["uuid"]

@app.get("/files/", summary="List files in bucket")
@s3_handler
async def list_files(bucket: str, folder: str = "", current_user: str = Depends(get_current_user)):
    if folder and not folder.endswith('/'):
        folder += '/'
    objects = minio_client.list_objects(bucket, prefix=folder, recursive=True)
    files = [
        obj.object_name 
        for obj in objects 
            if not obj.object_name.endswith('/') and obj.object_name != folder
    ]
    return {"files": files}

@app.post("/upload")
@s3_handler
async def upload_file(bucket: str, file: UploadFile = File(...), current_user: str = Depends(get_current_user)):
    object_name = f"{get_file_category(file.filename)}{file.filename}"
    minio_client.put_object(bucket, object_name, file.file, file.size)
    return {"message": f"File '{object_name}' uploaded to '{bucket}'."}

@app.post("/download")
@s3_handler
async def download_file(request: DownloadRequest, current_user: str = Depends(get_current_user)):
    try:
        response = minio_client.get_object(request.bucket, request.file)
        file_content = response.read()
        response.close()
        response.release_conn()
        return StreamingResponse(iter([file_content]), media_type="application/octet-stream", headers={"Content-Disposition": f"attachment; filename*=UTF-8''{quote(request.file)}"})
    except S3Error as exc:
        raise HTTPException(status_code=500, detail=str(exc))

@app.post("/preview")
@s3_handler
async def preview_images(
    bucket: str, 
    files: List[str],
    percentage: int = Query(..., gt=0, le=100, description="Percentage of original size (1-100)"), 
    current_user: str = Depends(get_current_user)
):
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file:
        for file_path in files:
            try:
                with minio_client.get_object(bucket, file_path) as response, \
                     Image.open(io.BytesIO(response.read())) as img:
                    w, h = img.size
                    new_w, new_h = int(w * percentage/100), int(h * percentage/100)
                    img_buffer = io.BytesIO()
                    img.resize((new_w, new_h), Image.LANCZOS).save(img_buffer, img.format or 'JPEG')
                    zip_file.writestr(f"preview_{new_w}x{new_h}_{file_path.split('/')[-1]}", img_buffer.getvalue())
                    
            except S3Error as e:
                raise HTTPException(500, f"Error processing {file_path}: {str(e)}")
            except Exception as e:
                raise HTTPException(500, f"Image processing error: {str(e)}")
    zip_buffer.seek(0)
    return StreamingResponse(zip_buffer, media_type="application/zip",
        headers={"Content-Disposition": "attachment; filename=preview_images.zip"})

@app.delete("/delete")
@s3_handler
async def delete_file(bucket: str, file: str, current_user: str = Depends(get_current_user)):
    minio_client.remove_object(bucket, file)
    return {"message": f"File '{file}' deleted successfully from '{bucket}'."}

@app.delete("/remove")
@s3_handler
async def remove_bucket(bucket: str, current_user: str = Depends(get_current_user)):
    for obj in minio_client.list_objects(bucket, recursive=True):
        minio_client.remove_object(bucket, obj.object_name)
    minio_client.remove_bucket(bucket)
    return {"message": f"Bucket '{bucket}' removed successfully."}

@app.on_event("shutdown")
def shutdown_db_connection():
    if db_conn:
        db_conn.close()
        print("Database connection closed.")