app/routers/drive/storage.py

import shutil
import os
import asyncio
from fastapi import (
    APIRouter, 
    UploadFile, 
    File, 
    HTTPException, 
    Header, 
    status, 
    Depends, 
    WebSocket, 
    WebSocketDisconnect
)

# Import Auth System
from app.routers.auth.system import get_current_user 

# Router Configuration
router = APIRouter(
    prefix="/drive",
    tags=["Cloud Drive"]
)

# Settings
UPLOAD_DIR = "uploaded_files"
MAX_FILE_SIZE = 1 * 1024 * 1024 * 1024  # 1GB Limit

# 🔐 Auth Toggle
AUTH_ENABLED = False  # <<< set False to disable authentication

# Auth Wrapper (returns dependency or None)
def auth_dependency():
    return Depends(get_current_user) if AUTH_ENABLED else None

# Ensure the upload directory exists
os.makedirs(UPLOAD_DIR, exist_ok=True)

# ==========================================
# 🛠️ Helper: Safe Username Extraction
# ==========================================
def get_safe_username(user_obj):
    if not user_obj:
        return "anonymous"
    if hasattr(user_obj, "username"):
        return user_obj.username
    return user_obj.get("username", "anonymous")


# ==========================================
# 📤 Standard File Upload
# ==========================================
@router.post("/upload/")
async def upload_file(
    file: UploadFile = File(...),
    content_length: int = Header(None),
    current_user: object = auth_dependency()
):
    if content_length and content_length > MAX_FILE_SIZE:
        raise HTTPException(
            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
            detail="File is too large. Maximum limit is 1GB."
        )

    username = get_safe_username(current_user)
    file_path = os.path.join(UPLOAD_DIR, file.filename)

    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")
    finally:
        file.file.close()

    return {
        "user": username,
        "filename": file.filename,
        "path": file_path,
        "message": "File uploaded successfully"
    }


# ==========================================
# 🚀 Remote Upload (Aria2)
# ==========================================
@router.post("/remote-upload/")
async def remote_url_upload(
    url: str,
    custom_filename: str = None,
    current_user: object = auth_dependency()
):
    username = get_safe_username(current_user)

    command = ["aria2c", "-x", "6", "-s", "6", "-d", UPLOAD_DIR, url]
    if custom_filename:
        command.extend(["-o", custom_filename])

    try:
        process = await asyncio.create_subprocess_exec(
            *command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()

        if process.returncode != 0:
            raise HTTPException(status_code=500, detail=f"Aria2 Error: {stderr.decode()}")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"System Error: {str(e)}")

    return {
        "user": username,
        "source": url,
        "filename": custom_filename or "Auto-detected",
        "message": "Download completed via Aria2"
    }


# ==========================================
# 📡 WebSocket Remote Upload (No Auth Applied)
# ==========================================
@router.websocket("/ws/remote-upload/")
async def websocket_remote_upload(websocket: WebSocket):
    await websocket.accept()
    
    try:
        data = await websocket.receive_json()
        target_url = data.get("url")
        custom_filename = data.get("filename")

        if not target_url:
            await websocket.send_text("❌ Error: No URL provided.")
            await websocket.close()
            return

        command = ["aria2c", "-x", "6", "-s", "6", "-d", UPLOAD_DIR, target_url]
        if custom_filename:
            command.extend(["-o", custom_filename])

        await websocket.send_text(f"🚀 Starting download: {target_url}")

        process = await asyncio.create_subprocess_exec(
            *command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        while True:
            line = await process.stdout.readline()
            if not line:
                if process.returncode is not None:
                    break
                await asyncio.sleep(0.1)
                continue

            decoded_line = line.decode().strip()
            if decoded_line:
                await websocket.send_text(decoded_line)

        await process.wait()
        if process.returncode == 0:
            await websocket.send_text("✅ Download Complete.")
        else:
            await websocket.send_text("❌ Download Failed.")

        await websocket.close()

    except WebSocketDisconnect:
        print("WS: Client disconnected")
    except Exception as e:
        try:
            await websocket.send_text(f"🔥 Critical Error: {str(e)}")
            await websocket.close()
        except:
            pass