"""The Engineering Layer payload of the engineer image.

The gea-engineer unattend stages and runs
Deploy/Applications/extra_engineering/Install-EngineeringLayer.ps1 (the
playbook deploys it from the repo). It installs, from beside itself:

    engineering_layer/Engineering_Install_v<version>.exe   (~530 MB)
    prereqs/redistrib/..., prereqs/sxs/*.cab                 (VC++ / .NET 3.5)

A new site has none of that; this module lets it be uploaded from the
engineer image page (Engineering Layer tab):

  * Engineering_Install_v<ver>.exe  replaces the installer, keeps prereqs/
  * a .zip with engineering_layer/ (and usually prereqs/) at its top level,
    or under one extra_engineering/ folder  replaces both

The upload streams to disk, is checked (an MZ executable named
Engineering_Install_v<ver>.exe; a zip with no absolute or .. paths), and
only then swapped in: the current engineering_layer/ and prereqs/ move to
UPDATES_DIR/englayer/archive/<timestamp>/ (the two newest are kept) and the
new ones are renamed into place. Install-EngineeringLayer.ps1 is never
touched (it belongs to the repo); it installs the newest
Engineering_Install_v*.exe it finds.
"""

import fcntl
import os
import re
import shutil
import tempfile
import zipfile
from datetime import datetime

import config
from services import updates_common as uc

IMAGE = "gea-engineer"
EXE = re.compile(r"^Engineering_Install_v(\d+(?:\.\d+)*)\.exe$", re.IGNORECASE)
PARTS = ("engineering_layer", "prereqs")
KEEP_ARCHIVES = 2


def root():
    return os.path.join(config.SAMBA_SHARE, IMAGE, "Deploy", "Applications", "extra_engineering")


def _state(*parts):
    return uc.state_dir("englayer", *parts)


def _version_key(v):
    return tuple(int(x) for x in v.split("."))


def _installers(folder):
    out = []
    try:
        names = os.listdir(folder)
    except OSError:
        return out
    for n in names:
        m = EXE.match(n)
        p = os.path.join(folder, n)
        if m and os.path.isfile(p):
            st = os.stat(p)
            out.append({"name": n, "version": m.group(1), "size": st.st_size, "mtime": st.st_mtime})
    return sorted(out, key=lambda x: _version_key(x["version"]), reverse=True)


def _tree_size(path):
    total, count = 0, 0
    for dirpath, _, files in os.walk(path):
        for f in files:
            try:
                total += os.path.getsize(os.path.join(dirpath, f))
                count += 1
            except OSError:
                pass
    return total, count


def status():
    r = root()
    inst = _installers(os.path.join(r, "engineering_layer"))
    pre = os.path.join(r, "prereqs")
    pre_size, pre_files = _tree_size(pre) if os.path.isdir(pre) else (0, 0)
    archives = []
    adir = _state("archive")
    for name in sorted(os.listdir(adir), reverse=True):
        p = os.path.join(adir, name)
        if os.path.isdir(p):
            a_inst = _installers(os.path.join(p, "engineering_layer"))
            archives.append({"name": name, "version": a_inst[0]["version"] if a_inst else None,
                             "prereqs": os.path.isdir(os.path.join(p, "prereqs"))})
    return {
        "root": r,
        "image_exists": os.path.isdir(os.path.join(config.SAMBA_SHARE, IMAGE)),
        "installers": inst,
        "active": inst[0] if inst else None,
        "prereqs": {"present": pre_files > 0, "files": pre_files, "size": pre_size,
                    "sxs": os.path.isdir(os.path.join(pre, "sxs")),
                    "redistrib": os.path.isdir(os.path.join(pre, "redistrib"))},
        "script": os.path.isfile(os.path.join(r, "Install-EngineeringLayer.ps1")),
        "archives": archives,
    }


class _Lock:
    def __enter__(self):
        self.fd = os.open(os.path.join(_state(), ".lock"), os.O_CREAT | os.O_RDWR, 0o644)
        try:
            fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except OSError:
            os.close(self.fd)
            raise RuntimeError("another Engineering Layer upload is being applied") from None
        return self

    def __exit__(self, *exc):
        fcntl.flock(self.fd, fcntl.LOCK_UN)
        os.close(self.fd)
        return False


def _check_name(name):
    name = os.path.basename(str(name or "")).strip()
    if EXE.match(name) or name.lower().endswith(".zip"):
        return name
    raise ValueError("upload Engineering_Install_v<version>.exe, or a .zip of the "
                     "extra_engineering payload (engineering_layer/ and prereqs/)")


def _receive(stream, name, length):
    """Stream to UPDATES_DIR/englayer/uploads/ (never held in memory)."""
    folder = _state("uploads")
    if length:
        free = shutil.disk_usage(folder).free
        if 2 * length > free - (1 << 30):  # the upload plus its unpacked copy
            raise RuntimeError(f"not enough disk space for {length // (1 << 20)} MB "
                               f"({free // (1 << 20)} MB free)")
    fd, part = tempfile.mkstemp(dir=folder, prefix=".upload.", suffix=".part")
    written = 0
    try:
        with os.fdopen(fd, "wb") as f:
            while True:
                block = stream.read(1 << 20)
                if not block:
                    break
                f.write(block)
                written += len(block)
        if length and written != length:
            raise RuntimeError(f"upload incomplete: {written} of {length} bytes")
        if not written:
            raise RuntimeError("empty upload")
        return part, written
    except BaseException:
        if os.path.exists(part):
            os.remove(part)
        raise


def _new_from_exe(part, name, work):
    with open(part, "rb") as f:
        if f.read(2) != b"MZ":
            raise ValueError(f"{name} is not a Windows executable")
    layer = os.path.join(work, "engineering_layer")
    os.makedirs(layer)
    shutil.move(part, os.path.join(layer, name))
    return ("engineering_layer",)


def _new_from_zip(part, work):
    try:
        zf = zipfile.ZipFile(part)
    except zipfile.BadZipFile:
        raise ValueError("the .zip file is damaged or not a zip") from None
    with zf:
        names = [i.filename.replace("\\", "/") for i in zf.infolist()]
        for n in names:
            if n.startswith("/") or ":" in n.split("/")[0] or ".." in n.split("/"):
                raise ValueError(f"unsafe path in the zip: {n}")
        # Accept the payload at the top, or inside one extra_engineering/ folder.
        prefix = ""
        if not any(n.split("/")[0] in PARTS for n in names):
            tops = {n.split("/")[0] for n in names if n}
            if len(tops) == 1:
                prefix = tops.pop() + "/"
        exes = [n for n in names if n.startswith(prefix + "engineering_layer/")
                and EXE.match(n.rsplit("/", 1)[-1])]
        if not exes:
            raise ValueError("the zip has no engineering_layer/Engineering_Install_v<version>.exe")
        found = set()
        for info in zf.infolist():
            n = info.filename.replace("\\", "/")
            if not n.startswith(prefix):
                continue
            rel = n[len(prefix):]
            top = rel.split("/")[0]
            if top not in PARTS:
                continue  # Install-EngineeringLayer.ps1 etc.: the repo owns those
            found.add(top)
            dest = os.path.join(work, *rel.split("/"))
            if info.is_dir():
                os.makedirs(dest, exist_ok=True)
                continue
            os.makedirs(os.path.dirname(dest), exist_ok=True)
            with zf.open(info) as src, open(dest, "wb") as out:
                shutil.copyfileobj(src, out, 1 << 20)
    with open(os.path.join(work, exes[0][len(prefix):]), "rb") as f:
        if f.read(2) != b"MZ":
            raise ValueError(f"{exes[0]} is not a Windows executable")
    return tuple(p for p in PARTS if p in found)


def _fix_modes(path):
    for dirpath, dirs, files in os.walk(path):
        os.chmod(dirpath, 0o755)
        for f in files:
            os.chmod(os.path.join(dirpath, f), 0o644)


def upload(stream, name, length=None):
    """Receive one upload and swap it in. Returns a summary dict."""
    name = _check_name(name)
    if not os.path.isdir(os.path.join(config.SAMBA_SHARE, IMAGE)):
        raise RuntimeError(f"no {IMAGE} image on this server yet; import it first")
    part, size = _receive(stream, name, length)
    work = tempfile.mkdtemp(dir=_state("uploads"), prefix=".new.")
    try:
        if EXE.match(name):
            parts = _new_from_exe(part, name, work)
        else:
            parts = _new_from_zip(part, work)
        _fix_modes(work)
        r = root()
        os.makedirs(r, exist_ok=True)
        stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        archive = os.path.join(_state("archive"), stamp)
        with _Lock():
            moved = []
            for p in parts:
                cur = os.path.join(r, p)
                if os.path.exists(cur):
                    os.makedirs(archive, exist_ok=True)
                    shutil.move(cur, os.path.join(archive, p))
                    moved.append(p)
            for p in parts:
                shutil.move(os.path.join(work, p), os.path.join(r, p))
        _prune()
        inst = _installers(os.path.join(r, "engineering_layer"))
        return {"name": name, "size": size, "replaced": list(parts),
                "archived": stamp if moved else None,
                "version": inst[0]["version"] if inst else None}
    finally:
        shutil.rmtree(work, ignore_errors=True)
        if os.path.exists(part):
            os.remove(part)


def _prune():
    adir = _state("archive")
    old = sorted((d for d in os.listdir(adir) if os.path.isdir(os.path.join(adir, d))), reverse=True)
    for d in old[KEEP_ARCHIVES:]:
        shutil.rmtree(os.path.join(adir, d), ignore_errors=True)


def restore(stamp):
    """Put an archived payload back (the current one is archived in turn)."""
    if not re.fullmatch(r"\d{8}-\d{6}", str(stamp or "")):
        raise ValueError("no such archive")
    src = os.path.join(_state("archive"), stamp)
    if not os.path.isdir(src):
        raise ValueError("no such archive")
    r = root()
    now = datetime.now().strftime("%Y%m%d-%H%M%S")
    back = os.path.join(_state("archive"), now)
    with _Lock():
        parts = [p for p in PARTS if os.path.isdir(os.path.join(src, p))]
        for p in parts:
            cur = os.path.join(r, p)
            if os.path.exists(cur):
                os.makedirs(back, exist_ok=True)
                shutil.move(cur, os.path.join(back, p))
        for p in parts:
            shutil.move(os.path.join(src, p), os.path.join(r, p))
        shutil.rmtree(src, ignore_errors=True)
    _prune()
    return parts