onever_drive/backend/app/services/upload_service.py

import math
from datetime import datetime, timezone
from typing import Tuple, List, Dict, Any, Optional
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from app.models.models import Client, BackupJob, BackupSession, BackupChunk, BackupFile
from app.storage.local import storage_provider
from app.services.event_service import log_event
from app.services.retention_service import apply_retention_policy
from app.ws.manager import ws_manager

async def create_or_resume_session(
    db: AsyncSession,
    client: Client,
    filename: str,
    file_size: int,
    sha256_full: str,
    chunk_size: int = 4 * 1024 * 1024,
    job_id: Optional[int] = None
) -> Tuple[BackupSession, List[int]]:
    """
    Initializes a new upload session or resumes an existing incomplete session
    for the specified file and hash.
    """
    total_chunks = max(1, math.ceil(file_size / chunk_size))

    # Check for existing incomplete session for this client and file hash
    query = (
        select(BackupSession)
        .where(
            BackupSession.client_id == client.id,
            BackupSession.sha256_full == sha256_full,
            BackupSession.file_size == file_size,
            BackupSession.status.in_(["PENDING", "UPLOADING"])
        )
    )
    result = await db.execute(query)
    session = result.scalar_one_or_none()

    if session:
        # Resume existing session
        received_chunks = await storage_provider.get_received_chunks(session.session_code)
        session.received_chunks_count = len(received_chunks)
        await db.commit()
        await db.refresh(session)
        return session, received_chunks

    # Create new upload session
    session = BackupSession(
        client_id=client.id,
        job_id=job_id,
        filename=filename,
        file_size=file_size,
        chunk_size=chunk_size,
        total_chunks=total_chunks,
        received_chunks_count=0,
        sha256_full=sha256_full,
        status="UPLOADING",
        started_at=datetime.now(timezone.utc)
    )
    db.add(session)
    await db.commit()
    await db.refresh(session)

    # Initialize temporary storage
    await storage_provider.init_session_storage(session.session_code)

    await log_event(
        db=db,
        event_type="BACKUP_STARTED",
        message=f"Upload session initiated for '{filename}' ({file_size / (1024*1024):.2f} MB, {total_chunks} chunks).",
        severity="INFO",
        client_id=client.id,
        job_id=job_id,
        details={"session_code": session.session_code, "total_chunks": total_chunks}
    )

    return session, []

async def process_chunk_upload(
    db: AsyncSession,
    session: BackupSession,
    chunk_index: int,
    chunk_data: bytes,
    chunk_sha256: Optional[str] = None
) -> Dict[str, Any]:
    """
    Saves a chunk to temporary storage, registers chunk in database,
    and broadcasts live progress telemetry.
    """
    if session.status not in ["PENDING", "UPLOADING"]:
        raise ValueError(f"Cannot upload chunk: session is currently in state {session.status}")

    if chunk_index < 0 or chunk_index >= session.total_chunks:
        raise ValueError(f"Invalid chunk_index {chunk_index}. Session total chunks: {session.total_chunks}")

    # Save to storage (performs chunk SHA-256 verification if provided)
    await storage_provider.save_chunk(
        session_code=session.session_code,
        chunk_index=chunk_index,
        chunk_data=chunk_data,
        expected_sha256=chunk_sha256
    )

    # Record in database
    result = await db.execute(
        select(BackupChunk).where(
            BackupChunk.session_id == session.id,
            BackupChunk.chunk_index == chunk_index
        )
    )
    chunk_rec = result.scalar_one_or_none()

    if not chunk_rec:
        chunk_rec = BackupChunk(
            session_id=session.id,
            chunk_index=chunk_index,
            chunk_size=len(chunk_data),
            sha256=chunk_sha256 or "",
            is_received=True,
            received_at=datetime.now(timezone.utc)
        )
        db.add(chunk_rec)
    else:
        chunk_rec.is_received = True
        chunk_rec.received_at = datetime.now(timezone.utc)

    # Count received chunks
    received_list = await storage_provider.get_received_chunks(session.session_code)
    session.received_chunks_count = len(received_list)
    await db.commit()

    progress_pct = round((session.received_chunks_count / session.total_chunks) * 100, 2)

    # Broadcast live telemetry over WebSocket
    await ws_manager.broadcast("UPLOAD_PROGRESS", {
        "session_code": session.session_code,
        "filename": session.filename,
        "client_id": session.client_id,
        "chunk_index": chunk_index,
        "received_chunks": session.received_chunks_count,
        "total_chunks": session.total_chunks,
        "progress_percent": progress_pct
    })

    return {
        "chunk_index": chunk_index,
        "is_received": True,
        "total_received": session.received_chunks_count,
        "total_chunks": session.total_chunks,
        "progress_percent": progress_pct
    }

async def get_session_status_info(
    db: AsyncSession,
    session: BackupSession
) -> Dict[str, Any]:
    """Returns detailed session status and lists of received / missing chunks."""
    received = await storage_provider.get_received_chunks(session.session_code)
    received_set = set(received)
    missing = [i for i in range(session.total_chunks) if i not in received_set]
    progress_pct = round((len(received) / session.total_chunks) * 100, 2)

    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,
        "missing_chunks": missing,
        "status": session.status,
        "progress_percent": progress_pct
    }

async def complete_session(
    db: AsyncSession,
    session: BackupSession
) -> BackupFile:
    """
    Assembles chunks into final storage, verifies SHA-256 integrity,
    updates client stats, applies retention policy, and logs completion.
    """
    # Fetch client and job codes for directory naming
    client_res = await db.execute(select(Client).where(Client.id == session.client_id))
    client = client_res.scalar_one_or_none()
    if not client:
        raise ValueError(f"Client {session.client_id} not found")

    job_code = "DEFAULT"
    if session.job_id:
        job_res = await db.execute(select(BackupJob).where(BackupJob.id == session.job_id))
        job = job_res.scalar_one_or_none()
        if job:
            job_code = job.job_code

    session.status = "ASSEMBLING"
    await db.commit()

    try:
        # Assemble and verify streaming SHA-256
        rel_path, final_sha256, total_bytes = await storage_provider.assemble_file(
            session_code=session.session_code,
            client_code=client.client_code,
            job_code=job_code,
            filename=session.filename,
            total_chunks=session.total_chunks,
            expected_sha256=session.sha256_full
        )

        session.status = "SUCCESS"
        session.completed_at = datetime.now(timezone.utc)

        # Create BackupFile record
        backup_file = BackupFile(
            client_id=client.id,
            job_id=session.job_id,
            session_id=session.id,
            filename=session.filename,
            relative_path=rel_path,
            file_size=total_bytes,
            sha256=final_sha256,
            retention_tag="DAILY",
            is_active=True,
            created_at=datetime.now(timezone.utc)
        )
        db.add(backup_file)

        # Update client storage and last backup timestamp
        client.storage_used_bytes += total_bytes
        client.last_backup_at = datetime.now(timezone.utc)

        await db.commit()
        await db.refresh(backup_file)

        # Log completion event
        await log_event(
            db=db,
            event_type="BACKUP_COMPLETED",
            message=f"Backup successfully verified & stored: '{session.filename}' ({total_bytes / (1024*1024):.2f} MB). SHA-256: {final_sha256[:16]}...",
            severity="INFO",
            client_id=client.id,
            job_id=session.job_id,
            details={"sha256": final_sha256, "file_size": total_bytes, "path": rel_path}
        )

        # Apply retention policy if associated with a job
        if session.job_id:
            await apply_retention_policy(db, session.job_id)

        # Broadcast completion
        await ws_manager.broadcast("UPLOAD_COMPLETED", {
            "session_code": session.session_code,
            "filename": session.filename,
            "client_id": client.id,
            "file_size": total_bytes,
            "sha256": final_sha256,
            "status": "SUCCESS"
        })

        return backup_file

    except Exception as ex:
        session.status = "FAILED"
        session.error_message = str(ex)
        await db.commit()

        await log_event(
            db=db,
            event_type="BACKUP_FAILED",
            message=f"Backup assembly/verification failed for '{session.filename}': {str(ex)}",
            severity="ERROR",
            client_id=client.id,
            job_id=session.job_id,
            details={"error": str(ex)}
        )
        raise ex