"""Clonezilla images on the clonezilla share, grouped by machine number.

The boot menu's "Clonezilla: back up this PC" (playbook/clonezilla-scripts/
pxe-cz.sh) saves <machine>-<YYYY-MM-DD-HHMM>[-n] image directories with a
pxe-meta.txt (key=value: machine, serial, model, disk, date, Clonezilla
version, bits, save/check result). Images made by hand (manual mode) have
only Clonezilla's own files; their machine number comes from the directory
name and serial/model from Clonezilla's Info-dmi.txt when it is there.

Everything here is defensive: the share is written by PCs over SMB, so a
half-written, unreadable or odd directory is listed (or skipped) but never
raises out of list_images(). Paths from the browser go through
resolve_image(), which only accepts a direct child of the share that is a
real (not symlinked) Clonezilla image directory.
"""
from __future__ import annotations

import json
import os
import re
import shutil
import tarfile
import tempfile
import time
from datetime import datetime, timezone

import config

MACHINE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9-]{0,31}$")
# <machine>-YYYY-MM-DD-HHMM, optionally -n when two land in one minute.
NAME_RE = re.compile(r"^(?P<machine>[A-Za-z0-9][A-Za-z0-9-]*?)-(?P<stamp>\d{4}-\d{2}-\d{2}-\d{4})(?:-(?P<n>\d{1,3}))?$")
META_FILE = "pxe-meta.txt"
META_MAX = 64 * 1024
UNASSIGNED = ""            # group key for images with no machine number
IN_PROGRESS_SECONDS = 6 * 3600

DEFAULT_SETTINGS = {"keep_last": 0, "warn_percent": 85}


class ImageError(Exception):
    """A refused action (bad name, missing image, clash); the message is shown."""


def share():
    return config.CLONEZILLA_SHARE


# --- detection and metadata -------------------------------------------------

def _listdir(path):
    try:
        return os.listdir(path)
    except OSError:
        return []


def is_image_dir(path):
    """Clonezilla image dir: a 'disk' or 'parts' file plus an image file
    (partclone *-ptcl-img*, dd *.dd-img*, ntfsclone/partimage *-img.*) or
    Clonezilla's clonezilla-img log."""
    if os.path.islink(path) or not os.path.isdir(path):
        return False
    names = _listdir(path)
    if not ("disk" in names or "parts" in names):
        return False
    for n in names:
        if "-ptcl-img" in n or ".dd-img" in n or "-img." in n or n == "clonezilla-img":
            return True
    return False


def _read_text(path, limit=META_MAX):
    try:
        with open(path, "rb") as f:
            return f.read(limit).decode("utf-8", "replace")
    except OSError:
        return ""


def read_meta(path):
    """pxe-meta.txt as a dict ({} when missing or unreadable)."""
    out = {}
    for line in _read_text(os.path.join(path, META_FILE)).splitlines():
        if "=" not in line or line.startswith("#"):
            continue
        k, v = line.split("=", 1)
        k = k.strip()
        if k and k not in out:
            out[k] = v.strip()
    return out


def _dmi_fallback(path):
    """Serial, model and maker from Clonezilla's Info-dmi.txt (dmidecode
    output), System Information block."""
    text = _read_text(os.path.join(path, "Info-dmi.txt"), 512 * 1024)
    out = {}
    block = False
    for line in text.splitlines():
        s = line.strip()
        if s == "System Information":
            block = True
            continue
        if block:
            if not line.startswith(("\t", " ")) and s:
                break
            for key, name in (("serial", "Serial Number"), ("model", "Product Name"), ("manufacturer", "Manufacturer")):
                if s.startswith(name + ":") and key not in out:
                    out[key] = s.split(":", 1)[1].strip()
    return out


def _version_fallback(path):
    """'clonezilla-live-3.3.3-37-amd64' from clonezilla-img ("Saved by ...")."""
    m = re.search(r"Saved by (clonezilla-live-[A-Za-z0-9._-]+?)\.?\s", _read_text(os.path.join(path, "clonezilla-img"), 16 * 1024) + " ")
    return m.group(1) if m else ""


def _bits_from_version(v):
    v = (v or "").lower()
    if "amd64" in v or "x86_64" in v:
        return "64"
    if "i686" in v or "i386" in v or "i586" in v:
        return "32"
    return ""


def _words(path, name):
    return _read_text(os.path.join(path, name), 4096).split()


def dir_size(path):
    """Bytes on disk, like du (st_blocks), symlinks not followed."""
    total = 0
    for root, dirs, files in os.walk(path, followlinks=False):
        for n in files:
            try:
                st = os.lstat(os.path.join(root, n))
            except OSError:
                continue
            total += getattr(st, "st_blocks", 0) * 512 or st.st_size
    return total


def _parse_stamp(stamp):
    try:
        return datetime.strptime(stamp, "%Y-%m-%d-%H%M")
    except (TypeError, ValueError):
        return None


def norm_machine(m):
    return (m or "").strip().upper()


def image_info(name, with_size=True):
    """One image directory as a dict for the page and the actions."""
    path = os.path.join(share(), name)
    meta = read_meta(path)
    m = NAME_RE.match(name)
    machine = norm_machine(meta.get("machine"))
    source = "meta" if machine else ""
    if not machine and m:
        machine, source = norm_machine(m.group("machine")), "name"
    if machine and not MACHINE_RE.match(machine):
        machine, source = UNASSIGNED, ""
    try:
        mtime = os.stat(path).st_mtime
    except OSError:
        mtime = 0
    when = None
    if meta.get("date_utc"):
        try:
            when = datetime.strptime(meta["date_utc"], "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc).timestamp()
        except ValueError:
            when = None
    if when is None and m and _parse_stamp(m.group("stamp")):
        when = _parse_stamp(m.group("stamp")).timestamp()   # server-local naming time
    if when is None:
        when = mtime
    fb = {} if (meta.get("serial") and meta.get("model")) else _dmi_fallback(path)
    version = meta.get("clonezilla") or _version_fallback(path)
    save = meta.get("save", "")
    check = meta.get("check", "")
    if save == "failed":
        status = "failed"
    elif check == "ok":
        status = "ok"
    elif check == "failed":
        status = "check-failed"
    elif meta:
        status = "unchecked"
    elif time.time() - mtime < IN_PROGRESS_SECONDS:
        status = "in-progress"
    else:
        status = "manual"
    disk_size = meta.get("disk_size_bytes", "")
    return {
        "name": name,
        "machine": machine,
        "machine_source": source,
        "when": when,
        "mtime": mtime,
        "stamp": m.group("stamp") if m else "",
        "suffix": (m.group("n") or "") if m else "",
        "size": dir_size(path) if with_size else 0,
        "disks": _words(path, "disk"),
        "parts": _words(path, "parts"),
        "serial": meta.get("serial") or fb.get("serial", ""),
        "model": meta.get("model") or fb.get("model", ""),
        "manufacturer": meta.get("manufacturer") or fb.get("manufacturer", ""),
        "clonezilla": version,
        "bits": meta.get("bits") or _bits_from_version(version),
        "firmware": meta.get("firmware", ""),
        "disk_model": meta.get("disk_model", ""),
        "disk_size": int(disk_size) if disk_size.isdigit() else 0,
        "date_text": meta.get("date", ""),
        "check_date": meta.get("check_date", ""),
        "save": save,
        "check": check,
        "status": status,
        "has_meta": bool(meta),
        "good": status in ("ok", "unchecked", "manual"),
    }


def list_images(with_size=True):
    out = []
    for name in sorted(_listdir(share())):
        if name.startswith("."):
            continue
        try:
            if not is_image_dir(os.path.join(share(), name)):
                continue
            out.append(image_info(name, with_size))
        except Exception:  # noqa: BLE001 - an odd dir must never break the page
            continue
    out.sort(key=lambda i: (i["when"], i["suffix"].zfill(3)), reverse=True)
    return out


def list_zips():
    out = []
    for name in sorted(_listdir(share())):
        p = os.path.join(share(), name)
        try:
            if name.lower().endswith(".zip") and os.path.isfile(p) and not os.path.islink(p):
                st = os.stat(p)
                out.append({"filename": name, "machine": os.path.splitext(name)[0],
                            "size": st.st_size, "modified": st.st_mtime})
        except OSError:
            continue
    return out


def grouped(images=None, zips=None):
    """Machines, newest activity first (unassigned last). Each: machine,
    images (newest first), zips, serials, mismatch, latest."""
    images = list_images() if images is None else images
    zips = list_zips() if zips is None else zips
    groups = {}

    def g(key):
        return groups.setdefault(key, {"machine": key, "images": [], "zips": [], "latest": 0})

    for i in images:
        grp = g(i["machine"])
        grp["images"].append(i)
        grp["latest"] = max(grp["latest"], i["when"] or 0)
    for z in zips:
        key = norm_machine(z["machine"]) if MACHINE_RE.match(z["machine"] or "") else UNASSIGNED
        grp = g(key)
        grp["zips"].append(z)
        grp["latest"] = max(grp["latest"], z["modified"] or 0)
    for grp in groups.values():
        serials = []
        for i in grp["images"]:
            s = (i["serial"] or "").strip().upper()
            if s and s not in serials:
                serials.append(s)
        grp["serials"] = serials
        # Rows whose serial is not the newest image's: the likely typo.
        ref = next(((i["serial"] or "").strip().upper() for i in grp["images"] if (i["serial"] or "").strip()), "")
        for i in grp["images"]:
            s = (i["serial"] or "").strip().upper()
            i["serial_differs"] = bool(s and ref and s != ref)
        grp["mismatch"] = len(serials) > 1
        grp["size"] = sum(i["size"] for i in grp["images"]) + sum(z["size"] for z in grp["zips"])
    return sorted(groups.values(), key=lambda x: (x["machine"] == UNASSIGNED, -x["latest"], x["machine"]))


# --- path safety --------------------------------------------------------------

def resolve_image(name):
    """Absolute path of image directory `name` (a direct child of the share),
    or ImageError."""
    if (not name or name in (".", "..") or name.startswith(".") or "/" in name or "\\" in name
            or "\x00" in name or name != os.path.basename(name) or len(name) > 255):
        raise ImageError("Invalid image name.")
    root = os.path.realpath(share())
    path = os.path.join(root, name)
    if os.path.islink(path):
        raise ImageError(f"{name} is a link, not an image directory.")
    if os.path.dirname(os.path.realpath(path)) != root:
        raise ImageError("Invalid image name.")
    if not is_image_dir(path):
        raise ImageError(f"Image not found: {name}")
    return path


# --- actions --------------------------------------------------------------------

def delete_image(name):
    path = resolve_image(name)
    shutil.rmtree(path)
    return name


def _write_meta(path, updates):
    meta_path = os.path.join(path, META_FILE)
    lines = _read_text(meta_path).splitlines()
    keys = set(updates)
    kept = [l for l in lines if l.split("=", 1)[0].strip() not in keys]
    kept += [f"{k}={v}" for k, v in updates.items()]
    fd, tmp = tempfile.mkstemp(dir=path, prefix=".pxe-meta-")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write("\n".join(kept) + "\n")
        os.chmod(tmp, 0o644)
        os.replace(tmp, meta_path)
    except BaseException:
        if os.path.exists(tmp):
            os.remove(tmp)
        raise


def move_image(name, new_machine):
    """Give image `name` to machine `new_machine`: rename the directory to
    <new>-<same stamp> and set machine= in pxe-meta.txt. Returns the new name."""
    path = resolve_image(name)
    new_machine = norm_machine(new_machine)
    if not MACHINE_RE.match(new_machine):
        raise ImageError("Machine number: letters, digits and dash only (up to 32, not starting with a dash).")
    info = image_info(name, with_size=False)
    stamp = info["stamp"] or datetime.fromtimestamp(info["when"] or time.time()).strftime("%Y-%m-%d-%H%M")
    existing = {n.lower() for n in _listdir(share())}
    base = f"{new_machine}-{stamp}"
    new = base if not info["suffix"] else f"{base}-{info['suffix']}"
    n = 2
    # SMB is case-insensitive: a case twin is the same name to Windows.
    while new.lower() in existing and new.lower() != name.lower():
        new = f"{base}-{n}"
        n += 1
    new_path = os.path.join(os.path.dirname(path), new)
    if new != name:
        if os.path.lexists(new_path) and new.lower() != name.lower():
            raise ImageError(f"{new} already exists.")
        os.rename(path, new_path)
    updates = {"machine": new_machine}
    if info["machine"] and info["machine"] != new_machine:
        updates["moved_from"] = info["machine"]
    _write_meta(new_path, updates)
    return new


# --- tar download (streamed, no temp copy) -------------------------------------

CHUNK = 1024 * 1024


def tar_stream(name):
    """Generator of a POSIX (pax) tar of image `name`, read file by file in
    1 MiB chunks: nothing is staged. A file that shrinks while it is read is
    padded with zeros so the archive stays well-formed."""
    path = resolve_image(name)

    def header(ti):
        return ti.tobuf(tarfile.PAX_FORMAT, "utf-8", "surrogateescape")

    def gen():
        written = 0
        for root, dirs, files in os.walk(path, followlinks=False):
            dirs.sort()
            rel_root = os.path.relpath(root, path)
            arc_root = name if rel_root == "." else f"{name}/{rel_root}"
            ti = tarfile.TarInfo(arc_root)
            ti.type = tarfile.DIRTYPE
            ti.mode = 0o755
            try:
                ti.mtime = int(os.lstat(root).st_mtime)
            except OSError:
                ti.mtime = int(time.time())
            buf = header(ti)
            written += len(buf)
            yield buf
            for n in sorted(files):
                full = os.path.join(root, n)
                try:
                    st = os.lstat(full)
                except OSError:
                    continue
                if not os.path.isfile(full) or os.path.islink(full):
                    continue   # sockets, fifos, links: not part of an image
                ti = tarfile.TarInfo(f"{arc_root}/{n}")
                ti.size = st.st_size
                ti.mtime = int(st.st_mtime)
                ti.mode = 0o644
                try:
                    f = open(full, "rb")
                except OSError:
                    continue
                buf = header(ti)
                written += len(buf)
                yield buf
                left = ti.size
                with f:
                    while left > 0:
                        chunk = f.read(min(CHUNK, left))
                        if not chunk:
                            break
                        left -= len(chunk)
                        written += len(chunk)
                        yield chunk
                while left > 0:
                    pad = min(CHUNK, left)
                    left -= pad
                    written += pad
                    yield b"\0" * pad
                rem = ti.size % tarfile.BLOCKSIZE
                if rem:
                    pad = tarfile.BLOCKSIZE - rem
                    written += pad
                    yield b"\0" * pad
        end = b"\0" * (tarfile.BLOCKSIZE * 2)
        written += len(end)
        rem = written % tarfile.RECORDSIZE
        if rem:
            end += b"\0" * (tarfile.RECORDSIZE - rem)
        yield end

    return gen()


# --- disk space -----------------------------------------------------------------

def disk_space(warn_percent=None):
    warn = settings()["warn_percent"] if warn_percent is None else warn_percent
    try:
        du = shutil.disk_usage(share())
    except OSError:
        return None
    pct = (du.used / du.total * 100) if du.total else 0
    return {"total": du.total, "used": du.used, "free": du.free, "percent": pct,
            "warn_percent": warn, "warn": pct >= warn}


# --- settings and retention ---------------------------------------------------

def _settings_path():
    return config.CLONEZILLA_STATE


def settings():
    out = dict(DEFAULT_SETTINGS)
    try:
        with open(_settings_path(), encoding="utf-8") as f:
            saved = json.load(f)
        if isinstance(saved, dict):
            for k in DEFAULT_SETTINGS:
                if isinstance(saved.get(k), int) and not isinstance(saved.get(k), bool):
                    out[k] = saved[k]
    except (OSError, ValueError):
        pass
    return out


def save_settings(keep_last, warn_percent):
    try:
        keep_last = int(keep_last)
        warn_percent = int(warn_percent)
    except (TypeError, ValueError):
        raise ImageError("Settings must be whole numbers.") from None
    if not 0 <= keep_last <= 100:
        raise ImageError("Keep last N images: 0 (off) to 100.")
    if not 50 <= warn_percent <= 99:
        raise ImageError("Disk space warning: 50 to 99 percent.")
    path = _settings_path()
    folder = os.path.dirname(path) or "."
    os.makedirs(folder, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=folder, prefix=".tmp-")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump({"keep_last": keep_last, "warn_percent": warn_percent}, f, indent=1)
        os.chmod(tmp, 0o644)
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.remove(tmp)
        raise
    return settings()


def retention_plan(keep_last=None, images=None):
    """What 'keep the last N images per machine' would delete, never applied
    here. Per machine number (unassigned images are never touched): keep the
    N newest GOOD images (saved and not failed the check; manual images
    count as good), delete every image OLDER than the oldest one kept.
    Newer failed or in-progress images stay, so a backup being written now
    is never removed and never pushes a good one out."""
    keep_last = settings()["keep_last"] if keep_last is None else keep_last
    if not keep_last or keep_last < 1:
        return []
    images = list_images() if images is None else images
    plan = []
    for grp in grouped(images, []):
        if grp["machine"] == UNASSIGNED:
            continue
        imgs = grp["images"]          # newest first
        good = [i for i in imgs if i["good"]]
        if len(good) <= keep_last:
            continue
        oldest_kept = good[keep_last - 1]
        cut = imgs.index(oldest_kept)
        for i in imgs[cut + 1:]:
            plan.append({"machine": grp["machine"], "name": i["name"], "when": i["when"],
                         "size": i["size"], "status": i["status"]})
    return plan


def apply_retention(names, keep_last=None):
    """Delete the posted names that are STILL in the plan (a stale preview
    cannot delete an image that has become one of the N newest)."""
    planned = {p["name"] for p in retention_plan(keep_last)}
    deleted, skipped = [], []
    for n in names:
        if n in planned:
            try:
                delete_image(n)
                deleted.append(n)
            except (ImageError, OSError):
                skipped.append(n)
        else:
            skipped.append(n)
    return deleted, skipped