ShooterHub/storage.py

"""
File storage helpers.

All paths are relative to Config.STORAGE_ROOT (a Docker volume at /app/storage).
Layout:
  csvs/{user_id}/{analysis_id}_{filename}
  pdfs/{user_id}/{analysis_id}_report.pdf
  equipment_photos/{user_id}/{item_id}_{uuid}.jpg
  session_photos/{user_id}/{session_id}_{uuid}.jpg
"""

import io
import uuid
from pathlib import Path

from flask import current_app
from PIL import Image


def _root() -> Path:
    return Path(current_app.config["STORAGE_ROOT"])


def _ensure(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path


# ---------------------------------------------------------------------------
# Analysis files
# ---------------------------------------------------------------------------

def _to_python(obj):
    """Recursively convert numpy scalars/arrays to plain Python types for JSON storage."""
    if isinstance(obj, dict):
        return {k: _to_python(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_to_python(v) for v in obj]
    if hasattr(obj, "item"):   # numpy scalar → int/float
        return obj.item()
    if hasattr(obj, "tolist"):  # numpy array → list
        return obj.tolist()
    return obj


def save_analysis(*, user, csv_bytes: bytes, pdf_bytes: bytes, overall: dict,
                  group_stats: list, filename: str, session_id: int | None = None,
                  is_public: bool = False) -> int:
    """Persist a completed analysis for a logged-in user. Returns the new Analysis.id."""
    from extensions import db
    from models import Analysis

    overall = _to_python(overall)
    group_stats = _to_python(group_stats)
    shot_count = int(overall.get("count", 0))
    group_count = len(group_stats)

    analysis = Analysis(
        user_id=user.id,
        session_id=session_id,
        is_public=is_public,
        title=_default_title(filename),
        csv_path="",          # filled in below
        pdf_path="",
        overall_stats=overall,
        group_stats=group_stats,
        shot_count=shot_count,
        group_count=group_count,
    )
    db.session.add(analysis)
    db.session.flush()  # assigns analysis.id without committing

    csv_dir = _ensure(_root() / "csvs" / str(user.id))
    pdf_dir = _ensure(_root() / "pdfs" / str(user.id))

    safe_name = Path(filename).name.replace(" ", "_")
    csv_rel = f"csvs/{user.id}/{analysis.id}_{safe_name}"
    pdf_rel = f"pdfs/{user.id}/{analysis.id}_report.pdf"

    (csv_dir / f"{analysis.id}_{safe_name}").write_bytes(csv_bytes)
    (pdf_dir / f"{analysis.id}_report.pdf").write_bytes(pdf_bytes)

    analysis.csv_path = csv_rel
    analysis.pdf_path = pdf_rel
    db.session.commit()

    return analysis.id


def _default_title(filename: str) -> str:
    stem = Path(filename).stem.replace("_", " ").replace("-", " ")
    return stem[:255] if stem else "Analysis"


# ---------------------------------------------------------------------------
# Equipment photos
# ---------------------------------------------------------------------------

MAX_PHOTO_DIM = 1200
PHOTO_QUALITY = 85


def _save_photo(file_storage, dest_dir: Path, prefix: str) -> str:
    """
    Validate, resize, re-encode as JPEG, and save an uploaded photo.
    Returns the path relative to STORAGE_ROOT.
    Raises ValueError on invalid image data.
    """
    try:
        img = Image.open(file_storage)
        img.verify()
        file_storage.seek(0)
        img = Image.open(file_storage)
        img = img.convert("RGB")
    except Exception as exc:
        raise ValueError(f"Invalid image file: {exc}") from exc

    img.thumbnail((MAX_PHOTO_DIM, MAX_PHOTO_DIM), Image.LANCZOS)

    _ensure(dest_dir)
    unique = uuid.uuid4().hex
    filename = f"{prefix}_{unique}.jpg"

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=PHOTO_QUALITY, optimize=True)
    (dest_dir / filename).write_bytes(buf.getvalue())

    return str(dest_dir.relative_to(_root()) / filename)


def save_equipment_photo(user_id: int, item_id: int, file_storage) -> str:
    dest = _root() / "equipment_photos" / str(user_id)
    return _save_photo(file_storage, dest, str(item_id))


def save_session_photo(user_id: int, session_id: int, file_storage) -> str:
    dest = _root() / "session_photos" / str(user_id)
    return _save_photo(file_storage, dest, str(session_id))


def save_avatar(user_id: int, file_storage) -> str:
    dest = _root() / "avatars" / str(user_id)
    return _save_photo(file_storage, dest, "avatar")


def save_analysis_group_photo(user_id: int, analysis_id: int, group_index: int,
                              file_storage) -> str:
    dest = _root() / "analysis_group_photos" / str(user_id)
    return _save_photo(file_storage, dest, f"{analysis_id}_g{group_index}")


def rotate_photo(rel_path: str, degrees: int) -> None:
    """Rotate a stored JPEG in-place. degrees is clockwise (90, -90, 180)."""
    path = _root() / rel_path
    if not path.exists():
        return
    img = Image.open(path).convert("RGB")
    # Pillow rotates counter-clockwise; negate to get clockwise behaviour
    rotated = img.rotate(-degrees, expand=True)
    buf = io.BytesIO()
    rotated.save(buf, format="JPEG", quality=PHOTO_QUALITY, optimize=True)
    path.write_bytes(buf.getvalue())