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