onever_drive/backend/app/api/upload.py

from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, status, Header, Request, Query, UploadFile, File
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from app.core.database import get_db
from app.models.models import Client, BackupSession
from app.schemas.schemas import (
    UploadSessionInitRequest, UploadSessionInitResponse,
    UploadSessionStatusResponse, ChunkUploadResponse,
    UploadSessionCompleteResponse
)
from app.api.deps import get_current_client
from app.services.upload_service import (
    create_or_resume_session, process_chunk_upload,
    get_session_status_info, complete_session
)

router = APIRouter(prefix="/upload", tags=["Chunk Upload Engine"])

@router.post("/session", response_model=UploadSessionInitResponse)
async def init_session(
    payload: UploadSessionInitRequest,
    db: AsyncSession = Depends(get_db),
    current_client: Client = Depends(get_current_client)
):
    """
    Initializes a new upload session or resumes an existing incomplete session
    for the specified file. Returns list of previously received chunks so the agent
    only transmits the remaining blocks.
    """
    session, received_chunks = await create_or_resume_session(
        db=db,
        client=current_client,
        filename=payload.filename,
        file_size=payload.file_size,
        sha256_full=payload.sha256,
        chunk_size=payload.chunk_size,
        job_id=payload.job_id
    )

    return {
        "session_code": session.session_code,
        "filename": session.filename,
        "file_size": session.file_size,
        "chunk_size": session.chunk_size,
        "total_chunks": session.total_chunks,
        "received_chunks": received_chunks,
        "status": session.status
    }

@router.post("/{session_code}/chunk", response_model=ChunkUploadResponse)
async def upload_chunk(
    session_code: str,
    request: Request,
    chunk_index: int = Query(..., description="0-indexed chunk number"),
    chunk_sha256: Optional[str] = Query(None, description="SHA-256 hash of this specific chunk"),
    x_chunk_index: Optional[int] = Header(None, alias="X-Chunk-Index"),
    x_chunk_sha256: Optional[str] = Header(None, alias="X-Chunk-SHA256"),
    db: AsyncSession = Depends(get_db),
    current_client: Client = Depends(get_current_client)
):
    """
    Receives and stores a single chunk of data for an active upload session.
    Accepts raw binary body directly via streaming.
    """
    idx = x_chunk_index if x_chunk_index is not None else chunk_index
    sha = x_chunk_sha256 or chunk_sha256

    res = await db.execute(
        select(BackupSession).where(
            BackupSession.session_code == session_code,
            BackupSession.client_id == current_client.id
        )
    )
    session = res.scalar_one_or_none()
    if not session:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Upload session not found")

    chunk_data = await request.body()
    if not chunk_data:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Empty chunk payload")

    try:
        result = await process_chunk_upload(
            db=db,
            session=session,
            chunk_index=idx,
            chunk_data=chunk_data,
            chunk_sha256=sha
        )
        return result
    except ValueError as ex:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(ex))

@router.get("/{session_code}/status", response_model=UploadSessionStatusResponse)
async def get_session_status(
    session_code: str,
    db: AsyncSession = Depends(get_db),
    current_client: Client = Depends(get_current_client)
):
    """
    Returns the current status of an upload session, including the list of received
    chunks and missing chunks for easy re-connection and resumption.
    """
    res = await db.execute(
        select(BackupSession).where(
            BackupSession.session_code == session_code,
            BackupSession.client_id == current_client.id
        )
    )
    session = res.scalar_one_or_none()
    if not session:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Upload session not found")

    return await get_session_status_info(db, session)

@router.post("/{session_code}/complete", response_model=UploadSessionCompleteResponse)
async def complete_upload(
    session_code: str,
    db: AsyncSession = Depends(get_db),
    current_client: Client = Depends(get_current_client)
):
    """
    Triggers sequential file assembly and final SHA-256 integrity verification.
    If integrity passes, the backup file is registered and retention policies are applied.
    """
    res = await db.execute(
        select(BackupSession).where(
            BackupSession.session_code == session_code,
            BackupSession.client_id == current_client.id
        )
    )
    session = res.scalar_one_or_none()
    if not session:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Upload session not found")

    try:
        backup_file = await complete_session(db, session)
        return {
            "session_code": session.session_code,
            "filename": backup_file.filename,
            "relative_path": backup_file.relative_path,
            "file_size": backup_file.file_size,
            "sha256": backup_file.sha256,
            "status": "SUCCESS",
            "completed_at": backup_file.created_at
        }
    except Exception as ex:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=f"Integrity check or assembly failed: {str(ex)}"
        )