vm-bench/frontend/app.py

"""
VM Bench — Proxmox Image Conversion Frontend

FastAPI web application serving as the user GUI.
Runs in the vm-bench LXC on port 5000.
Communicates with the backend at http://10.2.0.2:9000.

Routes:
  GET  /                       New session form
  POST /session/start           Upload + analyse source image
  POST /session/confirm         Submit conversion job
  GET  /session/status/{id}     Poll job status (JSON)
  POST /session/cleanup/{id}    Clean up / reuse staging files
"""

from __future__ import annotations

import logging
from logging.handlers import RotatingFileHandler
import os
import shutil
import subprocess
import threading
import time
import uuid
from pathlib import Path
from typing import Optional

import requests as http_requests

from fastapi import FastAPI, Form, Request, UploadFile, File
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from jinja2 import Environment, FileSystemLoader

from api_client import ApiClient, ApiError, BACKEND_URL

# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------
app = FastAPI(title="VM Bench — Frontend", version="1.0.0")

# ── Request logging middleware (logs before handler, catches silent failures) ──
@app.middleware("http")
async def log_requests(request: Request, call_next):
    cl = request.headers.get("content-length", "?")
    logger.info("→ %s %s (body %s bytes)", request.method, request.url.path, cl)
    try:
        response = await call_next(request)
        # Prevent browser caching (especially important for JS-heavy pages)
        response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
        response.headers["Pragma"] = "no-cache"
        response.headers["Expires"] = "0"
        return response
    except Exception:
        logger.exception("← %s %s FAILED", request.method, request.url.path)
        raise

BASE = Path(__file__).parent
app.mount("/static", StaticFiles(directory=str(BASE / "static")), name="static")

_jinja = Environment(loader=FileSystemLoader(str(BASE / "templates")), autoescape=True)
STAGING_ROOT = Path("/mnt/converter/tmp")

def _staging(session_id: str = "") -> Path:
    """Session-aware staging directory. Returns /mnt/converter/tmp/{guid}/ or /mnt/converter/tmp/"""
    return STAGING_ROOT / session_id if session_id else STAGING_ROOT

def _get_session_id(request: Request) -> str:
    """Extract session ID from cookie or query param."""
    sid = request.cookies.get("vm_bench_sid", "")
    if not sid:
        sid = request.query_params.get("session_id", "")
    return sid

api = ApiClient()

logger = logging.getLogger("vm-bench")
logger.setLevel(logging.INFO)

# Console handler → systemd journal via stdout
_ch = logging.StreamHandler()
_ch.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
logger.addHandler(_ch)

# File handler → shared log directory on /mnt/converter
LOG_DIR = Path("/mnt/converter/logs")
LOG_DIR.mkdir(parents=True, exist_ok=True)
_fh = RotatingFileHandler(
    LOG_DIR / "vm-bench.log", maxBytes=10 * 1024 * 1024, backupCount=5,
)
_fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s"))
logger.addHandler(_fh)

logger.info("Frontend starting — log file: %s", LOG_DIR / "vm-bench.log")

# Config
VM_ID_MIN = 21000
VM_ID_MAX = 21100
DEFAULT_STORAGE = "local-lvm"
MIN_FREE_DISK_GB = 2                    # keep 2 GB headroom
DOWNLOAD_TIMEOUT = 14400                 # 4 hours absolute max (safety net)
SPEED_CHECK_AFTER = 30                   # wait N seconds before judging speed
MAX_ETA_SECONDS = 3600                   # kill download if ETA > 1 hour
UPLOAD_PROGRESS_INTERVAL = 1024**3       # log every 1 GiB during upload

# Track active background downloads for progress polling
_active_downloads: dict[str, dict] = {}

def _validate_vmid(vmid: int) -> Optional[str]:
    if not (VM_ID_MIN <= vmid <= VM_ID_MAX):
        return f"VM ID must be between {VM_ID_MIN} and {VM_ID_MAX}."
    return None

def _check_disk_space(path: Path, needed_gb: int) -> Optional[str]:
    """Return an error message if <path> has less than <needed_gb> + headroom free."""
    usage = shutil.disk_usage(path.parent if path.is_file() or not path.exists() else path)
    free_gb = usage.free / (1024**3)
    required = needed_gb + MIN_FREE_DISK_GB
    if free_gb < required:
        return (
            f"Insufficient disk space on {path.parent}: "
            f"{free_gb:.1f} GB free, {required:.1f} GB needed "
            f"({needed_gb} GB file + {MIN_FREE_DISK_GB} GB headroom). "
            f"Free up space or use a smaller file."
        )
    return None

def render(name: str, status: int = 200, **ctx) -> HTMLResponse:
    tpl = _jinja.get_template(name)
    return HTMLResponse(tpl.render(**ctx), status_code=status)


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/", response_class=HTMLResponse)
async def index(request: Request, vmid: Optional[int] = None, source: Optional[str] = None):
    """Landing page — new session form. Optionally pre-fills vmid + source for reuse."""
    return render("index.html", request=request, backend_url=BACKEND_URL,
                  prefill_vmid=vmid or "", prefill_source=source or "",
                  prefill_vmname="")


@app.post("/session/upload")
def session_upload(
    request: Request,
    vmid: int = Form(...),
    vm_name: str = Form(""),
    source_type: str = Form("upload"),
    source_file: Optional[UploadFile] = File(None),
    source_url: Optional[str] = Form(None),
    session_id: str = Form(""),
):
    """Phase 1 — acquire the source file. Returns JSON so the frontend can
    show progress, then call /session/analyze separately."""
    sdir = _staging(session_id) / "in"
    sdir.mkdir(parents=True, exist_ok=True)
    err = _validate_vmid(vmid)
    if err:
        return JSONResponse({"phase": "error", "error": err}, status_code=400)

    vm_name = vm_name.strip()

    # ── Upload ──────────────────────────────────────────────────────
    if source_type == "upload":
        if not source_file or not source_file.filename:
            return JSONResponse({"phase": "error", "error": "No file uploaded."}, status_code=400)

        filename = source_file.filename
        dest = sdir / filename

        content_length = request.headers.get("content-length")
        if content_length:
            estimated_gb = int(content_length) / (1024**3)
            err = _check_disk_space(dest, estimated_gb)
            if err:
                return JSONResponse({"phase": "error", "error": err}, status_code=400)

        try:
            logger.info("Receiving upload: %s (%s bytes)", filename, content_length or "unknown")
            written = 0
            with dest.open("wb") as f:
                while True:
                    chunk = source_file.file.read(8 * 1024 * 1024)
                    if not chunk:
                        break
                    f.write(chunk)
                    written += len(chunk)
                    if written % UPLOAD_PROGRESS_INTERVAL < len(chunk):
                        logger.info("Upload progress: %s — %.1f GiB", filename, written / (1024**3))
            file_size_gb = round(written / (1024**3), 1)
            logger.info("Upload complete: %s (%.1f GiB)", filename, file_size_gb)
        except OSError as exc:
            if dest.exists():
                dest.unlink(missing_ok=True)
            return JSONResponse({"phase": "error", "error": f"Upload failed (disk full?): {exc}"}, status_code=500)
        except Exception as exc:
            if dest.exists():
                dest.unlink(missing_ok=True)
            return JSONResponse({"phase": "error", "error": f"Upload failed: {exc}"}, status_code=500)

        return JSONResponse({
            "phase": "staged",
            "filename": filename,
            "vmid": vmid,
            "vm_name": vm_name,
            "file_size_gb": file_size_gb,
         })

    else:
        # ── Download ────────────────────────────────────────────────
        url = (source_url or "").strip()
        if not url:
            return JSONResponse({"phase": "error", "error": "No URL provided."}, status_code=400)

        filename = Path(url).name or f"download_{uuid.uuid4().hex[:8]}"
        dest = sdir / filename

        _active_downloads.pop(filename, None)

        usage = shutil.disk_usage(_staging(session_id))
        free_gb = usage.free / (1024**3)
        if free_gb < 50:
            logger.warning("Low disk: %.1f GB free — download may fail", free_gb)

        logger.info("Starting background download: %s → %s", url, dest)
        try:
        proc = subprocess.Popen(
            ["aria2c", "-x8", "-s8", "-d", str(dest.parent), "-o", dest.name, url],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        except Exception as exc:
            return JSONResponse({"phase": "error", "error": f"Failed to start download: {exc}"}, status_code=500)

        _active_downloads[filename] = {
            "proc": proc, "dest": dest, "vmid": vmid, "vm_name": vm_name,
            "start_time": time.time(), "content_length": 0, "_last_logged_bytes": 0,
        }

        content_length = 0
        def _fetch_cl():
            nonlocal content_length
            try:
                hr = http_requests.head(url, timeout=5, allow_redirects=True)
                cl = hr.headers.get("Content-Length")
                if cl:
                    content_length = int(cl)
                    if filename in _active_downloads:
                        _active_downloads[filename]["content_length"] = int(cl)
                    logger.info("Download size: %.1f GiB", int(cl) / (1024**3))
            except Exception:
                pass
        threading.Thread(target=_fetch_cl, daemon=True).start()

        return JSONResponse({
            "phase": "downloading",
            "filename": filename, "vmid": vmid, "vm_name": vm_name,
            "content_length_gb": round(content_length / (1024**3), 1) if content_length else None,
        })


@app.post("/session/upload-raw")
async def upload_raw(request: Request):
    """Raw streaming upload — bypasses multipart parsing for large files."""
    filename = request.headers.get("X-Filename", "upload.bin")
    vmid_str = request.headers.get("X-VMID", "")
    vm_name = request.headers.get("X-VM-Name", "")
    session_id = request.headers.get("X-Session-ID", "")

    try:
        vmid = int(vmid_str)
    except ValueError:
        return JSONResponse({"phase": "error", "error": "Invalid VM ID."}, status_code=400)

    err = _validate_vmid(vmid)
    if err:
        return JSONResponse({"phase": "error", "error": err}, status_code=400)

    sdir = _staging(session_id) / "in"
    sdir.mkdir(parents=True, exist_ok=True)
    dest = sdir / filename
    content_length = request.headers.get("content-length")
    if content_length:
        estimated_gb = int(content_length) / (1024**3)
        err = _check_disk_space(dest, estimated_gb)
        if err:
            return JSONResponse({"phase": "error", "error": err}, status_code=400)

    logger.info("Raw upload: %s (%s bytes)", filename, content_length or "unknown")
    written = 0
    try:
        with dest.open("wb") as f:
            async for chunk in request.stream():
                f.write(chunk)
                written += len(chunk)
                if written % UPLOAD_PROGRESS_INTERVAL < len(chunk):
                    logger.info("Upload progress: %s — %.1f GiB", filename, written / (1024**3))
        file_size_gb = round(written / (1024**3), 1)
        logger.info("Raw upload complete: %s (%.1f GiB)", filename, file_size_gb)
    except Exception as exc:
        if dest.exists():
            dest.unlink(missing_ok=True)
        logger.exception("Raw upload failed: %s", filename)
        return JSONResponse({"phase": "error", "error": f"Upload failed: {exc}"}, status_code=500)

    return JSONResponse({
        "phase": "staged",
        "filename": filename,
        "vmid": vmid,
        "vm_name": vm_name,
        "file_size_gb": file_size_gb,
    })


@app.get("/session/progress/{filename}")
async def session_progress(filename: str):
    """Poll download progress — returns current file size and phase."""
    info = _active_downloads.get(filename)
    if not info:
        # Check if file exists on disk (download already completed in a
        # previous session, or it was an upload)
        dest = _staging("") / "in" / filename
        if dest.exists():
            return JSONResponse({
                "phase": "complete",
                "file_size_bytes": dest.stat().st_size,
                "file_size_gb": round(dest.stat().st_size / (1024**3), 1),
            })
        return JSONResponse({"phase": "unknown", "error": "No active download for this file."}, status_code=404)

    proc = info["proc"]
    dest = info["dest"]

    # Current bytes on disk
    current_bytes = dest.stat().st_size if dest.exists() else 0

    # Check if process still running
    poll = proc.poll()
    if poll is not None:
        # Process exited
        _active_downloads.pop(filename, None)
        if poll != 0:
            logger.error("Download failed: %s (wget exited with code %d)", filename, poll)
            if dest.exists():
                dest.unlink(missing_ok=True)
            return JSONResponse({
                "phase": "error",
                "error": f"Download failed (wget exited with code {poll}).",
                "file_size_bytes": current_bytes,
            })
        # Success
        final_bytes = dest.stat().st_size
        logger.info("Download complete: %s (%.1f GiB)", filename, final_bytes / (1024**3))
        return JSONResponse({
            "phase": "complete",
            "file_size_bytes": final_bytes,
            "file_size_gb": round(final_bytes / (1024**3), 1),
        })

    # Still downloading — check speed and estimate ETA
    elapsed = max(time.time() - info.get("start_time", 0), 1)
    speed_bps = current_bytes / elapsed
    speed_mbps = round(speed_bps / 1_000_000, 1)

    content_length = info.get("content_length", 0)
    eta_str = None
    too_slow = False
    if content_length > 0 and elapsed > SPEED_CHECK_AFTER:
        remaining = content_length - current_bytes
        eta = remaining / max(speed_bps, 1)  # seconds
        if eta > MAX_ETA_SECONDS:
            too_slow = True
            eta_hours = round(eta / 3600, 1)
            # Kill the download
            try:
                proc.kill()
            except Exception:
                pass
            _active_downloads.pop(filename, None)
            if dest.exists():
                dest.unlink(missing_ok=True)
            logger.warning(
                "Download %s killed: ETA %.1f h at %.1f MB/s (threshold %d h)",
                filename, eta_hours, speed_mbps, MAX_ETA_SECONDS // 3600,
            )
            return JSONResponse({
                "phase": "too_slow",
                "speed_mbps": speed_mbps,
                "eta_hours": eta_hours,
                "content_length_gb": round(content_length / (1024**3), 1),
                "downloaded_gb": round(current_bytes / (1024**3), 1),
                "message": (
                    f"Download would take ~{eta_hours} hours at {speed_mbps} MB/s "
                    f"(file is {round(content_length / (1024**3), 1)} GiB). "
                    f"Consider downloading to your computer manually, then use File Upload."
                ),
            })
        eta_str = f"~{round(eta / 60)} min remaining"

    if too_slow:
        # Already handled above; this line is unreachable but kept for clarity
        pass

    # Log progress periodically (every ~1 GiB)
    last_logged = info.get("_last_logged_bytes", 0)
    if current_bytes - last_logged >= 1024**3:
        info["_last_logged_bytes"] = current_bytes
        logger.info("Downloading: %s — %.1f GiB (%.1f MB/s)%s",
                    filename, current_bytes / (1024**3), speed_mbps,
                    f" ETA {eta_str}" if eta_str else "")

    return JSONResponse({
        "phase": "downloading",
        "file_size_bytes": current_bytes,
        "file_size_gb": round(current_bytes / (1024**3), 1),
        "speed_mbps": speed_mbps,
        "eta": eta_str,
        "content_length_gb": round(content_length / (1024**3), 1) if content_length else None,
    })


@app.post("/session/analyze", response_class=HTMLResponse)
async def session_analyze(
    request: Request,
    vmid: int = Form(...),
    filename: str = Form(...),
    vm_name: str = Form(""),
    session_id: str = Form(""),
):
    """Phase 2 — call the backend /analyze endpoint and render the result."""
    err = _validate_vmid(vmid)
    if err:
        return render("_analysis.html", request=request, error=err)

    # Prepend session subdirectory to the filename for the backend
    # Files are in tmp/{guid}/in/, so the backend needs {guid}/in/filename
    backend_filename = f"{session_id}/in/{filename}" if session_id else filename

    try:
        analysis = api.analyze(vmid=vmid, filename=backend_filename)
    except ApiError as exc:
        return render("_analysis.html", request=request,
                      error=f"Backend analysis failed: {exc.detail}")

    vm_name = vm_name.strip()
    if not vm_name:
        vm_name = (analysis.get("os_type") or "vm") + f"-{vmid}"

    return render("_analysis.html", request=request,
                  vmid=vmid, source_filename=filename,
                  vm_name=vm_name, analysis=analysis)


@app.post("/session/confirm", response_class=HTMLResponse)
async def confirm_session(
    request: Request,
    vmid: int = Form(...),
    source_filename: str = Form(...),
    disk_format: str = Form(...),
    vm_name: str = Form(...),
    cpu_cores: int = Form(2),
    ram_mb: int = Form(4096),
    target_storage: str = Form("local-lvm"),
    target_disk_size_gb: Optional[int] = Form(None),
    auto_detect_boot: str = Form("true"),
    boot_type: str = Form("uefi"),
    session_id: str = Form(""),
):
    """Build the job payload and submit to the backend."""
    # Validate VM ID range
    err = _validate_vmid(vmid)
    if err:
        return render("_analysis.html", request=request, error=err)

    payload = {
        "vmid": vmid,
        "vm_name": vm_name,
        "cpu_cores": cpu_cores,
        "ram_mb": ram_mb,
        "target_storage": target_storage,
        "auto_detect_boot": auto_detect_boot == "true",
        "boot_type": boot_type,
        "session_id": session_id,
        "boot_disk": {
            "disk_type": "image_file",
            "source_filename": source_filename,
            "format": disk_format,
        },
    }
    if target_disk_size_gb is not None:
        payload["target_disk_size_gb"] = target_disk_size_gb

    try:
        job = api.create_job(payload)
    except ApiError as exc:
        return render("_analysis.html", request=request,
                      error=f"Job submission failed: {exc.detail}")

    return render("polling.html", request=request,
                  job_id=job["job_id"], vmid=vmid,
                  vm_name=vm_name, source_filename=source_filename)


@app.get("/session/status/{job_id}")
async def session_status(job_id: str):
    """Poll backend for job status (returns JSON for AJAX polling)."""
    try:
        return api.get_job(job_id)
    except ApiError as exc:
        return JSONResponse({"error": exc.detail, "status": "error"}, status_code=502)


@app.post("/session/cleanup/{job_id}")
async def session_cleanup(job_id: str, request: Request):
    """Forward cleanup request to backend."""
    try:
        body = await request.json()
        delete = body.get("delete_staging_files", False)
        session_id = body.get("session_id", "")
        return api.cleanup_job(job_id, delete, session_id)
    except ApiError as exc:
        return JSONResponse({"error": exc.detail}, status_code=502)


@app.post("/session/clone")
async def session_clone(request: Request):
    """Proxy clone request to backend."""
    try:
        body = await request.json()
        result = api.clone_vm(
            source_vmid=body["source_vmid"],
            target_vmid=body["target_vmid"],
            target_name=body["target_name"],
        )
        return result
    except ApiError as exc:
        return JSONResponse({"error": exc.detail}, status_code=502)


# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------
# Track active SCP pulls for progress polling
_active_scp: dict[str, dict] = {}


# ---------------------------------------------------------------------------
# SCP Pull — large file transfer from remote servers
# ---------------------------------------------------------------------------

@app.get("/scp", response_class=HTMLResponse)
async def scp_page(request: Request):
    """SCP pull page — fetch files directly from remote servers."""
    return render("scp.html", request=request, backend_url=BACKEND_URL,
                  session_id=uuid.uuid4().hex)


@app.post("/scp/start")
def scp_start(
    request: Request,
    scp_host: str = Form(...),
    scp_port: int = Form(22),
    scp_user: str = Form(...),
    scp_pass: str = Form(...),
    scp_path: str = Form(...),
    session_id: str = Form(""),
):
    """Start an SCP pull in the background. Returns JSON with phase + filename."""
    remote = f"{scp_user}@{scp_host}:{scp_path}"
    filename = Path(scp_path).name
    if not filename:
        return JSONResponse({"phase": "error", "error": "Invalid remote path."}, status_code=400)

    sdir = _staging(session_id) / "in"
    sdir.mkdir(parents=True, exist_ok=True)
    dest = sdir / filename

    # Check disk space
    usage = shutil.disk_usage(STAGING)
    free_gb = usage.free / (1024**3)
    if free_gb < 10:
        logger.warning("Low disk: %.1f GB free — SCP pull may fail", free_gb)

    env = os.environ.copy()
    env["SSHPASS"] = scp_pass

    cmd = [
        "sshpass", "-e",
        "scp",
        "-o", "StrictHostKeyChecking=no",
        "-o", "ConnectTimeout=10",
        "-P", str(scp_port),
        remote, str(dest),
    ]

    logger.info("Starting SCP pull: %s → %s", remote, dest)
    try:
        proc = subprocess.Popen(cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception as exc:
        return JSONResponse({"phase": "error", "error": f"Failed to start SCP: {exc}"}, status_code=500)

    key = f"{session_id}/{filename}"
    _active_scp[key] = {
        "proc": proc,
        "dest": dest,
        "start_time": time.time(),
    }

    return JSONResponse({
        "phase": "pulling",
        "filename": filename,
        "session_id": session_id,
    })


@app.get("/scp/progress/{session_id}/{filename}")
def scp_progress(session_id: str, filename: str):
    """Poll SCP progress by checking output file size."""
    key = f"{session_id}/{filename}"
    info = _active_scp.get(key)
    if not info:
        dest = _staging(session_id) / "in" / filename
        if dest.exists():
            return JSONResponse({
                "phase": "complete",
                "file_size_bytes": dest.stat().st_size,
                "file_size_gb": round(dest.stat().st_size / (1024**3), 1),
            })
        return JSONResponse({"phase": "unknown", "error": "No active SCP pull."}, status_code=404)

    proc = info["proc"]
    dest = info["dest"]
    current_bytes = dest.stat().st_size if dest.exists() else 0

    poll = proc.poll()
    if poll is not None:
        _active_scp.pop(key, None)
        if poll != 0:
            if dest.exists():
                dest.unlink(missing_ok=True)
            return JSONResponse({
                "phase": "error",
                "error": f"SCP pull failed (exit code {poll}). Check credentials and remote path.",
                "file_size_bytes": current_bytes,
            })
        final_bytes = dest.stat().st_size
        logger.info("SCP pull complete: %s (%.1f GiB)", filename, final_bytes / (1024**3))
        return JSONResponse({
            "phase": "complete",
            "file_size_bytes": final_bytes,
            "file_size_gb": round(final_bytes / (1024**3), 1),
        })

    elapsed = max(time.time() - info.get("start_time", 0), 1)
    speed_mbps = round(current_bytes / elapsed / 1_000_000, 1)

    return JSONResponse({
        "phase": "pulling",
        "file_size_bytes": current_bytes,
        "file_size_gb": round(current_bytes / (1024**3), 1),
        "speed_mbps": speed_mbps,
    })


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)