MileVault/backend/app/api/upload.py

import os
import zipfile
from pathlib import Path
from fastapi import APIRouter, Depends, UploadFile, File, HTTPException, BackgroundTasks
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel

from app.core.database import get_db
from app.core.security import get_current_user
from app.core.config import settings
from app.models.user import User
from app.workers.tasks import process_activity_file, process_garmin_health_zip, analyze_strava_export

router = APIRouter()

MAX_FILE_SIZE = 500 * 1024 * 1024            # 500 MB upload cap
MAX_EXTRACT_SIZE = 4 * 1024 * 1024 * 1024    # 4 GB total uncompressed cap (zip-bomb guard)
_CHUNK = 1024 * 1024

_TASK_OWNER_TTL = 86400   # 24h — long enough to outlive any upload's polling


def _remember_task_owner(task_id: str, user_id: int) -> None:
    """Record which user a pollable task belongs to, so the status endpoint can
    refuse to surface another user's task result (the Celery task id is the only
    thing the client presents). Best-effort: Redis hiccups must not fail uploads."""
    if not task_id:
        return
    try:
        import redis as redis_lib
        redis_lib.Redis.from_url(settings.redis_url).set(
            f"upload_task_owner:{task_id}", user_id, ex=_TASK_OWNER_TTL
        )
    except Exception:
        pass


def _task_owner(task_id: str) -> int | None:
    try:
        import redis as redis_lib
        v = redis_lib.Redis.from_url(settings.redis_url).get(f"upload_task_owner:{task_id}")
        if v is None:
            return None
        return int(v.decode() if isinstance(v, (bytes, bytearray)) else v)
    except Exception:
        return None


def _safe_name(filename: str) -> str:
    """Reduce an uploaded filename to a safe basename — no path traversal."""
    name = os.path.basename((filename or "").replace("\\", "/"))
    if not name or name in (".", ".."):
        raise HTTPException(status_code=400, detail="Invalid filename")
    return name


def save_upload(upload: UploadFile, dest_dir: Path) -> Path:
    """Stream an upload to disk under dest_dir, enforcing the size cap."""
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest = dest_dir / _safe_name(upload.filename)
    size = 0
    with open(dest, "wb") as f:
        while True:
            chunk = upload.file.read(_CHUNK)
            if not chunk:
                break
            size += len(chunk)
            if size > MAX_FILE_SIZE:
                f.close()
                dest.unlink(missing_ok=True)
                raise HTTPException(status_code=413, detail="File exceeds the 500 MB limit")
            f.write(chunk)
    return dest


def _safe_extract(zf: zipfile.ZipFile, dest_dir: Path) -> list[Path]:
    """Extract a zip safely: skip path-traversal members, cap total uncompressed
    bytes (zip-bomb guard). Returns the list of extracted regular-file paths."""
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest_root = dest_dir.resolve()
    total = 0
    extracted: list[Path] = []
    for info in zf.infolist():
        if info.is_dir():
            continue
        target = (dest_root / info.filename).resolve()
        # Reject absolute paths and ../ traversal: the target must stay under dest_root.
        if target != dest_root and dest_root not in target.parents:
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        with zf.open(info) as src, open(target, "wb") as out:
            while True:
                chunk = src.read(_CHUNK)
                if not chunk:
                    break
                total += len(chunk)
                if total > MAX_EXTRACT_SIZE:
                    out.close()
                    target.unlink(missing_ok=True)
                    raise HTTPException(status_code=413, detail="Archive expands beyond the size limit")
                out.write(chunk)
        extracted.append(target)
    return extracted


def _gunzip(path: Path) -> Path | None:
    """Decompress a .gz member to a sibling file without the .gz suffix,
    enforcing the same uncompressed-size cap. Returns the new path, or None on
    failure. The .gz is removed once expanded."""
    import gzip
    out_path = path.with_suffix("")  # strips the trailing .gz
    try:
        total = 0
        with gzip.open(path, "rb") as src, open(out_path, "wb") as out:
            while True:
                chunk = src.read(_CHUNK)
                if not chunk:
                    break
                total += len(chunk)
                if total > MAX_EXTRACT_SIZE:
                    out.close()
                    out_path.unlink(missing_ok=True)
                    return None
                out.write(chunk)
    except (OSError, EOFError):
        out_path.unlink(missing_ok=True)
        return None
    finally:
        path.unlink(missing_ok=True)
    return out_path


@router.post("/activity")
async def upload_activity(
    file: UploadFile = File(...),
    background_tasks: BackgroundTasks = None,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Upload a single .fit, .gpx or .tcx activity file."""
    suffix = Path(file.filename).suffix.lower()
    if suffix not in {".fit", ".gpx", ".tcx"}:
        raise HTTPException(status_code=400, detail="Only .fit, .gpx and .tcx files are supported")

    dest_dir = Path(settings.file_store_path) / str(current_user.id) / "activities"
    dest = save_upload(file, dest_dir)

    # Queue processing
    task = process_activity_file.delay(str(dest), current_user.id, suffix[1:])
    _remember_task_owner(task.id, current_user.id)

    return {"task_id": task.id, "status": "queued", "filename": file.filename}


@router.post("/garmin-export")
async def upload_garmin_export(
    file: UploadFile = File(...),
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """
    Upload a full Garmin Connect data export ZIP.
    Processes all FIT files for activities + wellness data.
    """
    if not file.filename.endswith(".zip"):
        raise HTTPException(status_code=400, detail="Please upload a .zip Garmin export")

    dest_dir = Path(settings.file_store_path) / str(current_user.id) / "exports"
    dest = save_upload(file, dest_dir)

    # Extract (safely) and queue all FIT files
    extract_dir = dest_dir / f"garmin_{dest.stem}"

    task_ids = []
    try:
        with zipfile.ZipFile(dest) as zf:
            extracted = _safe_extract(zf, extract_dir)
    except zipfile.BadZipFile:
        dest.unlink(missing_ok=True)
        raise HTTPException(status_code=400, detail="Uploaded file is not a valid ZIP archive")

    has_health = False
    for path in extracted:
        suffix = path.suffix.lower()
        if suffix == ".fit":
            task = process_activity_file.delay(str(path), current_user.id, "fit")
            task_ids.append(task.id)
        elif suffix == ".json":
            has_health = True   # Garmin wellness data is exported as JSON files
        elif suffix == ".zip":
            # Garmin exports nest activity FIT files inside sub-zips
            # (e.g. DI-Connect-Uploaded-Files/UploadedFiles_*_Part*.zip)
            nested_extract = path.parent / path.stem
            try:
                with zipfile.ZipFile(path) as nzf:
                    nested = _safe_extract(nzf, nested_extract)
            except zipfile.BadZipFile:
                nested = []
            for np in nested:
                if np.suffix.lower() == ".fit":
                    task = process_activity_file.delay(str(np), current_user.id, "fit")
                    task_ids.append(task.id)

    if not task_ids and not has_health:
        raise HTTPException(
            status_code=400,
            detail="No fitness data found in this archive — make sure you uploaded your full Garmin Connect export ZIP",
        )

    # Queue health/wellness data extraction
    health_task = process_garmin_health_zip.delay(str(dest), current_user.id)
    _remember_task_owner(health_task.id, current_user.id)

    return {
        "status": "queued",
        "activity_tasks": len(task_ids),
        "task_id": health_task.id,
    }


def _collect_activity_files(extracted: list[Path]) -> list[Path]:
    """From extracted archive members, gunzip any .gz wrappers (Strava ships most
    activities as <id>.fit.gz / .gpx.gz / .tcx.gz) and return the .fit/.gpx/.tcx
    paths ready to parse. Without this, the gzipped majority is silently skipped."""
    files: list[Path] = []
    for path in extracted:
        if path.suffix.lower() == ".gz":
            inner = _gunzip(path)
            if inner is None:
                continue
            path = inner
        if path.suffix.lower() in (".fit", ".gpx", ".tcx"):
            files.append(path)
    return files


class StravaConfirmIn(BaseModel):
    token: str


@router.post("/strava-export")
async def upload_strava_export(
    file: UploadFile = File(...),
    dry_run: bool = False,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """Upload a Strava bulk export ZIP (activities/ folder with GPX/FIT/TCX files).

    With `dry_run=true` nothing is imported: the archive is extracted and a single
    analysis task reports how many activities are new vs. already present (Garmin
    data is preferred), returning a `token` to confirm the real import with.
    Otherwise activities import straight away (existing data is kept on a match)."""
    if not file.filename.endswith(".zip"):
        raise HTTPException(status_code=400, detail="Please upload a .zip Strava export")

    dest_dir = Path(settings.file_store_path) / str(current_user.id) / "exports"
    dest = save_upload(file, dest_dir)

    extract_dir = dest_dir / f"strava_{dest.stem}"

    try:
        with zipfile.ZipFile(dest) as zf:
            extracted = _safe_extract(zf, extract_dir)
    except zipfile.BadZipFile:
        dest.unlink(missing_ok=True)
        raise HTTPException(status_code=400, detail="Uploaded file is not a valid ZIP archive")

    files = _collect_activity_files(extracted)
    if not files:
        raise HTTPException(
            status_code=400,
            detail="No activity files (.fit, .gpx or .tcx) found in this Strava archive",
        )

    if dry_run:
        # Preview only — classify new/duplicate without writing. The extracted
        # files are kept so the confirm step can import them without re-uploading.
        task = analyze_strava_export.delay([str(p) for p in files], current_user.id)
        _remember_task_owner(task.id, current_user.id)
        return {"status": "analyzing", "task_id": task.id, "token": dest.stem,
                "activity_files": len(files)}

    task_ids = [
        process_activity_file.delay(str(p), current_user.id, p.suffix.lower()[1:],
                                    prefer_existing=True).id
        for p in files
    ]
    polled = task_ids[-1] if task_ids else None
    _remember_task_owner(polled, current_user.id)
    return {
        "status": "queued",
        "activity_tasks": len(task_ids),
        "task_id": polled,
    }


@router.post("/strava-export/confirm")
async def confirm_strava_export(
    body: StravaConfirmIn,
    current_user: User = Depends(get_current_user),
):
    """Import the activities from a previously previewed (dry-run) Strava export.
    `token` is the value returned by the dry-run upload; the already-extracted
    files are imported with Garmin data preferred over Strava duplicates."""
    token = _safe_name(body.token)
    extract_dir = Path(settings.file_store_path) / str(current_user.id) / "exports" / f"strava_{token}"
    if not extract_dir.is_dir():
        raise HTTPException(status_code=404, detail="Preview not found — please upload and preview again")

    files = [p for p in extract_dir.rglob("*")
             if p.is_file() and p.suffix.lower() in (".fit", ".gpx", ".tcx")]
    if not files:
        raise HTTPException(status_code=400, detail="No activity files found for this preview")

    task_ids = [
        process_activity_file.delay(str(p), current_user.id, p.suffix.lower()[1:],
                                    prefer_existing=True).id
        for p in files
    ]
    return {"status": "queued", "activity_tasks": len(task_ids)}


@router.get("/task/{task_id}")
async def check_task_status(
    task_id: str,
    current_user: User = Depends(get_current_user),
):
    """Check the status of an upload processing task."""
    # A task result can carry the owner's activity data, so don't surface another
    # user's task. We fail closed only on a positive owner mismatch; a missing
    # record (Redis down / TTL expired) stays permissive so polling never breaks.
    owner = _task_owner(task_id)
    if owner is not None and owner != current_user.id:
        raise HTTPException(status_code=404, detail="Task not found")

    from app.workers.celery_app import celery_app
    result = celery_app.AsyncResult(task_id)
    return {
        "task_id": task_id,
        "status": result.status,
        "result": result.result if result.ready() else None,
    }