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

# 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 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_url: Optional[str] = Form(None),
    session_id: str = Form(""),
):
    """Download a source file via aria2c. Returns JSON with phase + filename."""
    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()

    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 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,
    }

    def _fetch_cl():
        try:
            hr = http_requests.head(url, timeout=5, allow_redirects=True)
            cl = hr.headers.get("Content-Length")
            if cl and 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": None,
    })


@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")
async def session_analyze(
    request: Request,
    vmid: int = Form(...),
    filename: str = Form(...),
    vm_name: str = Form(""),
    session_id: str = Form(""),
):
    """Phase 2 — start async backend analysis, return analysis_id for polling."""
    err = _validate_vmid(vmid)
    if err:
        return JSONResponse({"phase": "error", "error": err}, status_code=400)

    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 JSONResponse({"phase": "error", "error": f"Backend analysis failed: {exc.detail}"}, status_code=502)

    return JSONResponse({
        "phase": "analyzing",
        "analysis_id": analysis.get("analysis_id"),
        "vmid": vmid,
        "filename": filename,
        "vm_name": vm_name.strip(),
    })


@app.get("/session/analyze/status/{analysis_id}")
async def session_analyze_status(analysis_id: str):
    """Poll backend for async analysis status."""
    try:
        return api.get_analysis(analysis_id)
    except ApiError as exc:
        return JSONResponse({"error": exc.detail}, status_code=502)


@app.get("/session/sessions")
async def session_list():
    """Proxy to backend to list resumable sessions."""
    try:
        return api.get_sessions()
    except ApiError as exc:
        return JSONResponse({"sessions": []})


@app.get("/session/configure/{analysis_id}", response_class=HTMLResponse)
async def session_configure(
    request: Request,
    analysis_id: str,
    vmid: str = "",
    source_filename: str = "",
    vm_name: str = "",
    session_id: str = "",
):
    """Full page with VM configuration + disk selection after analysis."""
    try:
        status = api.get_analysis(analysis_id)
        result = status.get("result")
        if not result:
            return render("configure.html", request=request, analysis_id=analysis_id,
                          source_filename=source_filename, session_id=session_id,
                          disks=[], error="Analysis result not ready.")
    except ApiError as exc:
        return render("configure.html", request=request, analysis_id=analysis_id,
                      source_filename=source_filename, session_id=session_id,
                      disks=[], error=f"Failed: {exc.detail}")

    pools = []
    try:
        sp = api.get_storage_pools()
        pools = sp.get("pools", [])
    except Exception:
        pass

    analysis = result
    name = vm_name.strip()
    if not name:
        name = (analysis.get("os_type") or "vm")

    disks = analysis.get("all_disks", [])
    if not disks:
        disks = [{"filename": analysis.get("filename", ""), "format": analysis.get("disk_format", ""), "size_gb": analysis.get("disk_size_gb", 0)}]

    suggested_vmid = vmid if vmid else ""

    return render("configure.html", request=request,
                  analysis_id=analysis_id, source_filename=source_filename,
                  session_id=session_id, vm_name=name, storage_pools=pools,
                  disks=disks, suggested_vmid=suggested_vmid)


@app.post("/session/configure/submit")
async def configure_submit(request: Request):
    """Receive config + disk selection, submit job to backend."""
    body = await request.json()
    try:
        result = api.create_job(body)
        return result
    except ApiError as exc:
        return JSONResponse({"error": exc.detail}, status_code=502)


@app.get("/session/status/{job_id}", response_class=HTMLResponse)
async def session_status_page(request: Request, job_id: str):
    """Full page showing job progress with polling."""
    return render("job_polling.html", request=request, job_id=job_id,
                  vmid=0, vm_name="", session_id="")


@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(""),
    has_multiple_disks: str = Form("0"),
    boot_disk_index: str = Form("0"),
):
    """Build the job payload and submit to the backend."""
    err = _validate_vmid(vmid)
    if err:
        return render("_analysis.html", request=request, error=err)

    # Determine boot disk from multi-disk selection
    boot_file = source_filename
    boot_fmt = disk_format

    if has_multiple_disks == "1":
        idx = int(boot_disk_index) if boot_disk_index.isdigit() else 0
        form_data = await request.form()
        boot_file = form_data.get(f"disk_{idx}_file", source_filename)
        boot_fmt = form_data.get(f"disk_{idx}_fmt", disk_format)

        additional_disks = []
        for i in range(10):
            dfile = form_data.get(f"disk_{i}_file")
            if dfile and i != idx:
                additional_disks.append({
                    "disk_type": "image_file",
                    "source_filename": dfile,
                    "format": form_data.get(f"disk_{i}_fmt", "vmdk"),
                })
    else:
        additional_disks = []

    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": boot_file,
            "format": boot_fmt,
        },
        "additional_disks": additional_disks,
    }
    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(session_id))
    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)

    # ── Probe file size via SSH ────────────────────────────────────
    file_size = 0
    try:
        probe_cmd = [
            "sshpass", "-e",
            "ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
            "-p", str(scp_port), f"{scp_user}@{scp_host}",
            f"stat -c%s '{scp_path}' 2>/dev/null || ls -l '{scp_path}' 2>/dev/null | awk '{{print $5}}'"
        ]
        probe = subprocess.run(probe_cmd, env=env, capture_output=True, text=True, timeout=15)
        if probe.returncode == 0 and probe.stdout.strip().isdigit():
            file_size = int(probe.stdout.strip())
            logger.info("Remote file size: %.1f GiB", file_size / (1024**3))
    except Exception:
        logger.info("Could not probe remote file size")

    # ── Start SCP ──────────────────────────────────────────────────
    try:
        proc = subprocess.Popen(cmd, env=env, stdout=subprocess.DEVNULL,
                                stderr=subprocess.PIPE, text=True)
    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(),
        "total_size": file_size,
    }

    return JSONResponse({
        "phase": "pulling",
        "filename": filename,
        "session_id": session_id,
        "total_size_gb": round(file_size / (1024**3), 1) if file_size else None,
    })


@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:
            stderr_output = ""
            try:
                stderr_output = proc.stderr.read()[:500] if proc.stderr else ""
            except Exception:
                pass
            if dest.exists():
                dest.unlink(missing_ok=True)
            logger.error("SCP pull failed (rc=%d): %s", poll, stderr_output)
            return JSONResponse({
                "phase": "error",
                "error": f"SCP pull failed (exit code {poll}). {stderr_output}",
                "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)
    total = info.get("total_size", 0)

    eta_str = None
    pct = 0
    if total > 0 and current_bytes > 0:
        pct = min(round((current_bytes / total) * 100), 99)
        remaining = total - current_bytes
        eta = remaining / max(current_bytes / elapsed, 1)
        if eta > 60:
            eta_str = f"~{round(eta / 60)} min remaining"

    return JSONResponse({
        "phase": "pulling",
        "file_size_bytes": current_bytes,
        "file_size_gb": round(current_bytes / (1024**3), 1),
        "speed_mbps": speed_mbps,
        "pct": pct,
        "eta": eta_str,
        "total_size_gb": round(total / (1024**3), 1) if total else None,
    })


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