"""Safe edits of an image's control files: preview, backup, atomic write, undo.

Every save on the image pages (Drivers, Unattend, OS and Packages) follows
the same three steps, and this module is the shared implementation:

  1. plan_json(path, data) / plan_text(path, text) render what WOULD be
     written, in the file's own on-disk format, and diff it against the
     current file. Nothing is written. The page shows plan["diff"].
  2. commit(image_type, scope, plans, action) backs every changed file up
     next to itself as <name>.bak-<YYYYmmdd-HHMMSS>-webui, writes the new
     text atomically (temp file in the same dir + os.replace, original
     owner and mode kept) and remembers the save as the undo point for
     (image_type, scope). Pass expect={path: sha256} (from the preview) to
     refuse the save when a file changed after the preview was shown.
  3. undo(image_type, scope) puts the backups of the last commit back, but
     only while every file still holds exactly what that commit wrote.

The on-disk format is never changed by a save: JSON keeps its indent, BOM,
line endings, trailing newline and ASCII escaping (detect_json_format).
A file that does not round-trip byte for byte (a hand-edited vendor file
with mixed tabs and spaces, say) is not re-rendered whole: for a top-level
JSON list the entries that did not change keep their exact original text
and only new or edited entries are rendered (_splice_list), so the diff
shows the real change. Only when that is impossible does a plan say
"reformat": True, so the page can warn that whitespace will move.

commit() and undo() hold one lock (UPDATES_DIR/webui-undo/.lock) from the
stale check to the last write, so two saves cannot interleave, and a
commit that fails half way puts the files it already wrote back.

Undo records live in UPDATES_DIR/webui-undo/<image>--<scope>.json; only
the last save per scope can be undone (the .bak files stay on the share
for anything older).
"""

import contextlib
import difflib
import fcntl
import hashlib
import json
import os
import re
import shutil
import tempfile
from datetime import datetime

from services import updates_common as uc

BACKUP_TAG = "webui"
_BOM = "\ufeff"


class Conflict(Exception):
    """A file changed on disk since the preview or since the last save."""


# --- format ------------------------------------------------------------------

def encode(text):
    """Bytes as written: UTF-8, a leading \\ufeff becomes a real BOM."""
    if text.startswith(_BOM):
        return b"\xef\xbb\xbf" + text[1:].encode("utf-8")
    return text.encode("utf-8")


def sha256_text(text):
    return hashlib.sha256(encode(text)).hexdigest()


def file_sha256(path):
    """sha256 of the file's bytes ("" hashed when missing)."""
    try:
        with open(path, "rb") as fh:
            return hashlib.sha256(fh.read()).hexdigest()
    except OSError:
        return hashlib.sha256(b"").hexdigest()


def read_text(path):
    """Current file content as text ("" when missing). A BOM is kept as the
    leading \\ufeff character so it survives a round trip."""
    try:
        with open(path, "rb") as fh:
            raw = fh.read()
    except OSError:
        return ""
    if raw.startswith(b"\xef\xbb\xbf"):
        return _BOM + raw[3:].decode("utf-8", errors="replace")
    return raw.decode("utf-8", errors="replace")


def detect_json_format(text):
    """How an existing JSON file is laid out. Defaults (for a new file) are
    indent 4, no BOM, LF, trailing newline, non-ASCII kept as is."""
    fmt = {"indent": 4, "bom": False, "newline": "\n", "trailing_newline": True,
           "ensure_ascii": False}
    if not text:
        return fmt
    fmt["bom"] = text.startswith(_BOM)
    body = text[1:] if fmt["bom"] else text
    fmt["newline"] = "\r\n" if "\r\n" in body else "\n"
    fmt["trailing_newline"] = body.endswith("\n")
    lines = body.splitlines()
    if len(lines) < 2:
        # One line: compact JSON, unless it is an empty "[]" / "{}", which
        # says nothing about layout (keep the indent 4 default).
        if not re.fullmatch(r"\s*(\[\s*\]|\{\s*\})\s*", body):
            fmt["indent"] = None
    else:
        m = re.match(r"[ \t]+", lines[1])
        if m:
            ws = m.group(0)
            fmt["indent"] = "\t" if ws.startswith("\t") else len(ws)
    # \uXXXX escapes and no raw non-ASCII: the writer escaped everything.
    fmt["ensure_ascii"] = bool(re.search(r"\\u[0-9a-fA-F]{4}", body)) and body.isascii()
    return fmt


def render_json(data, fmt):
    text = json.dumps(data, indent=fmt["indent"], ensure_ascii=fmt["ensure_ascii"])
    if fmt["trailing_newline"]:
        text += "\n"
    if fmt["newline"] != "\n":
        text = text.replace("\n", fmt["newline"])
    if fmt["bom"]:
        text = _BOM + text
    return text


# --- plan ----------------------------------------------------------------------

def _diff(old, new, name):
    return "".join(difflib.unified_diff(
        old.splitlines(keepends=True), new.splitlines(keepends=True),
        fromfile=f"a/{name}", tofile=f"b/{name}", n=3))


def plan_text(path, new_text, label=None):
    """What saving new_text to path would change. Keys: path, name, old,
    new, changed, diff, sha_before, existed, reformat (always False here)."""
    old = read_text(path)
    name = label or os.path.basename(path)
    return {
        "path": path,
        "name": name,
        "old": old,
        "new": new_text,
        "changed": old != new_text,
        "diff": _diff(old, new_text, name) if old != new_text else "",
        "sha_before": file_sha256(path),
        "existed": os.path.isfile(path),
        "reformat": False,
    }


_WS = re.compile(r"[ \t\r\n]*")


def _key(obj):
    """Order-sensitive identity of a JSON value (key order matters on disk)."""
    return json.dumps(obj, ensure_ascii=False)


def _splice_object(s, span, old_obj, new_obj, fmt):
    """Text of new_obj made from the old object's text at span, when both
    are objects with the same keys in the same order: every member keeps
    its text and only changed values are replaced (rendered on one line
    when scalar, else with fmt at that member's indentation). None when
    the shapes differ."""
    if not (isinstance(old_obj, dict) and isinstance(new_obj, dict)
            and list(old_obj) == list(new_obj)):
        return None
    dec = json.JSONDecoder()
    start, end = span
    i = _WS.match(s, start + 1).end()
    out, pos = [], start
    try:
        while s[i:i + 1] == '"':
            key, k_end = dec.raw_decode(s, i)
            j = _WS.match(s, k_end).end()
            if s[j:j + 1] != ":":
                return None
            v_start = _WS.match(s, j + 1).end()
            _val, v_end = dec.raw_decode(s, v_start)
            if _key(old_obj.get(key)) != _key(new_obj.get(key)):
                line_start = s.rfind("\n", 0, i) + 1
                indent = re.match(r"[ \t]*", s[line_start:i]).group(0)
                text = json.dumps(new_obj[key], indent=fmt["indent"],
                                  ensure_ascii=fmt["ensure_ascii"])
                out.append(s[pos:v_start])
                out.append(text.replace("\n", fmt["newline"] + indent))
                pos = v_end
            j = _WS.match(s, v_end).end()
            if s[j:j + 1] == ",":
                i = _WS.match(s, j + 1).end()
            else:
                break
    except ValueError:
        return None
    out.append(s[pos:end])
    return "".join(out)


def _splice_list(old, data, fmt):
    """New text for a top-level JSON list that keeps the exact original
    text of every entry that did not change (and the separators between
    neighbours that stay neighbours); new or edited entries are rendered
    with fmt at the first entry's indentation. None when old is not a
    plain JSON list or the result would not parse back to data."""
    if not isinstance(data, list) or not data:
        return None
    s = old
    dec = json.JSONDecoder()
    i = _WS.match(s, 1 if s.startswith(_BOM) else 0).end()
    if s[i:i + 1] != "[":
        return None
    head_end = i + 1
    i = _WS.match(s, head_end).end()
    if s[i:i + 1] == "]":
        return None
    spans, items = [], []
    try:
        while True:
            obj, end = dec.raw_decode(s, i)
            spans.append((i, end))
            items.append(obj)
            j = _WS.match(s, end).end()
            if s[j:j + 1] == ",":
                i = _WS.match(s, j + 1).end()
                continue
            if s[j:j + 1] == "]":
                close = j
                break
            return None
    except ValueError:
        return None
    lead = s[head_end:spans[0][0]]
    nl = fmt["newline"]
    if "\n" in lead:
        item_indent = lead.rsplit("\n", 1)[-1]
        std_sep = "," + nl + item_indent
    else:
        item_indent, std_sep = "", ", "

    sm = difflib.SequenceMatcher(None, [_key(o) for o in items], [_key(o) for o in data],
                                 autojunk=False)
    # src: old entry identical to new entry n; origin: the old entry new
    # entry n stands in for (identical, or edited in a one-for-one replace).
    src = [None] * len(data)
    origin = [None] * len(data)
    for tag, i1, i2, j1, j2 in sm.get_opcodes():
        if tag == "equal" or (tag == "replace" and i2 - i1 == j2 - j1):
            for k in range(j2 - j1):
                origin[j1 + k] = i1 + k
                if tag == "equal":
                    src[j1 + k] = i1 + k

    def fresh(obj):
        text = json.dumps(obj, indent=fmt["indent"], ensure_ascii=fmt["ensure_ascii"])
        return text.replace("\n", nl + item_indent)

    def edited(n, obj):
        # An entry that is one of the old ones with values changed (same
        # keys, same order) keeps its own layout: only the values move.
        k = origin[n]
        if src[n] is not None:
            return s[spans[k][0]:spans[k][1]]
        if k is not None:
            text = _splice_object(s, spans[k], items[k], obj, fmt)
            if text is not None:
                return text
        return fresh(obj)

    out = [s[:head_end], lead]
    for n, obj in enumerate(data):
        if n:
            a, b = origin[n - 1], origin[n]
            out.append(s[spans[a][1]:spans[b][0]] if a is not None and b == a + 1 else std_sep)
        out.append(edited(n, obj))
    out.append(s[spans[-1][1]:close])
    out.append(s[close:])
    text = "".join(out)
    try:
        if _key(json.loads(text.lstrip(_BOM))) != _key(data):
            return None
    except ValueError:
        return None
    return text


def plan_json(path, data, label=None):
    """plan_text for a JSON document, rendered in the file's own format.
    A file that does not round-trip keeps its unchanged entries verbatim
    (_splice_list); "reformat" is set only when that is not possible."""
    old = read_text(path)
    fmt = detect_json_format(old)
    if old:
        try:
            same = render_json(json.loads(old.lstrip(_BOM)), fmt) == old
        except ValueError:
            same = False
        if not same:
            spliced = _splice_list(old, data, fmt)
            if spliced is not None:
                return plan_text(path, spliced, label)
            plan = plan_text(path, render_json(data, fmt), label)
            plan["reformat"] = True
            return plan
    return plan_text(path, render_json(data, fmt), label)


def public(plan):
    """The part of a plan a page needs (no full file texts)."""
    return {k: plan[k] for k in ("name", "changed", "diff", "sha_before", "existed", "reformat")}


# --- write ---------------------------------------------------------------------

def _stamp():
    return datetime.now().strftime("%Y%m%d-%H%M%S")


def _match_owner(path, ref_stat):
    if ref_stat is None:
        return
    try:
        os.chown(path, ref_stat.st_uid, ref_stat.st_gid)
    except (PermissionError, OSError):
        pass  # not root (local tests): the file stays ours


def _ref_stat(path):
    """Stat to copy owner/mode from: the file, else its directory."""
    for p in (path, os.path.dirname(path)):
        try:
            return os.stat(p)
        except OSError:
            continue
    return None


def atomic_write(path, text):
    """Write text (see encode) or bytes via a temp file in the same
    directory, keeping the old file's owner and mode. A symlinked control
    file is written through to its target (os.replace on the link itself
    would turn it into a separate copy)."""
    path = os.path.realpath(path)
    folder = os.path.dirname(path)
    os.makedirs(folder, exist_ok=True)
    existed = os.path.isfile(path)
    ref = _ref_stat(path)
    data = text if isinstance(text, bytes) else encode(text)
    fd, tmp = tempfile.mkstemp(dir=folder, prefix=".webui-")
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
            fh.flush()
            os.fsync(fh.fileno())
        if existed:
            shutil.copymode(path, tmp)
        else:
            os.chmod(tmp, 0o664)
        _match_owner(tmp, ref)
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.remove(tmp)
        raise


def backup(path, tag=BACKUP_TAG, stamp=None):
    """Copy path to <path>.bak-<stamp>-<tag> (owner kept). None if missing."""
    if not os.path.isfile(path):
        return None
    dst = f"{path}.bak-{stamp or _stamp()}-{tag}"
    n = 1
    while os.path.exists(dst):  # two saves in one second
        n += 1
        dst = f"{path}.bak-{stamp or _stamp()}-{tag}{n}"
    shutil.copy2(path, dst)
    _match_owner(dst, _ref_stat(path))
    return dst


# --- undo records ------------------------------------------------------------------

def _undo_path(image_type, scope):
    safe = re.sub(r"[^A-Za-z0-9._-]", "_", f"{image_type}--{scope}")
    return os.path.join(uc.state_dir("webui-undo"), safe + ".json")


@contextlib.contextmanager
def _locked():
    """One lock for every commit/undo (saves are rare; threads and worker
    processes both honour it, flock is per open file)."""
    lock = os.path.join(uc.state_dir("webui-undo"), ".lock")
    with open(lock, "a") as fh:
        fcntl.flock(fh, fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(fh, fcntl.LOCK_UN)


def _restore(path, bak):
    """Put a file back as its backup holds it, or remove it (bak None)."""
    if bak:
        with open(bak, "rb") as fh:
            atomic_write(path, fh.read())
    elif os.path.lexists(path):
        os.remove(path)


def last_save(image_type, scope):
    """The undo point for (image_type, scope): {"ts", "action", "files":
    [{"path", "name", "backup", "existed", "sha_after"}]} or None."""
    return uc.read_json(_undo_path(image_type, scope), None)


def commit(image_type, scope, plans, action="", expect=None):
    """Back up and write every changed plan; record the undo point.
    Returns the list of changed plans (empty when nothing changed, and then
    the previous undo point is kept). Raises Conflict when expect is given
    and a file no longer has the previewed hash. When a write fails part
    way, the files already written are put back before the error is raised."""
    expect = expect or {}
    with _locked():
        changed = [p for p in plans if p["changed"]]
        stale = [p["name"] for p in plans
                 if p["name"] in expect and expect[p["name"]] != file_sha256(p["path"])]
        if stale:
            raise Conflict("changed on disk since the preview: " + ", ".join(stale)
                           + ". Reload the page and review again.")
        if not changed:
            return []
        stamp = _stamp()
        files = []
        try:
            for p in changed:
                bak = backup(p["path"], stamp=stamp)
                files.append({"path": p["path"], "name": p["name"], "backup": bak,
                              "existed": bak is not None, "sha_after": sha256_text(p["new"])})
                atomic_write(p["path"], p["new"])
            uc.write_json(_undo_path(image_type, scope),
                          {"ts": uc.now(), "action": action, "files": files})
        except BaseException:
            for f in reversed(files):
                try:
                    _restore(f["path"], f["backup"])
                except OSError:
                    pass  # the .bak file is still next to it
            raise
    return changed


def undo(image_type, scope):
    """Restore the backups of the last commit for (image_type, scope).
    The current files are themselves backed up (tag webui-undo) first, so
    an undo can be undone by hand. Returns the record that was undone."""
    with _locked():
        rec = last_save(image_type, scope)
        if not rec:
            raise Conflict("nothing to undo")
        moved = [f["name"] for f in rec["files"]
                 if file_sha256(f["path"]) != f["sha_after"]]
        if moved:
            raise Conflict("changed since that save, not undone: " + ", ".join(moved)
                           + ". The backups are still next to the files.")
        missing = [f["name"] for f in rec["files"]
                   if f["existed"] and not os.path.isfile(f["backup"] or "")]
        if missing:
            raise Conflict("backup missing for: " + ", ".join(missing))
        stamp = _stamp()
        for f in rec["files"]:
            backup(f["path"], tag=BACKUP_TAG + "-undo", stamp=stamp)
            _restore(f["path"], f["backup"] if f["existed"] else None)
        os.remove(_undo_path(image_type, scope))
    return rec