#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0

# Remove the Colab-only "To run this, press Runtime > Run all ..." sentence, which is
# wrong inside Docker, keeping the rest of the header cell. Two modes:
#
#   unsloth_nb_strip_colab.py <a.ipynb> ...              strip in place (idempotent)
#   unsloth_nb_strip_colab.py --state <STATE> --dest <DEST>
#       strip + rehash each notebook whose hash still matches STATE
#
# Safe with refresh: content_sig classifies the intro cell as boilerplate, so the
# body digest is unchanged. Exit code is always 0.
import argparse
import hashlib
import json
import os
import stat
import sys

_INTRO_PREFIX = "to run this, press"

# Baked tqdm widget outputs render as a stuck "Loading widget...", and outputs are not
# in the refresh signature, so dropping them is safe.
_WIDGET_VIEW_MIME = "application/vnd.jupyter.widget-view+json"


def _is_intro_line(line):
    """True for the Colab run announcement, bare or in a one-line HTML comment; only a comment that
    opens and closes on one line matches, so no dangling <!-- can swallow the cell."""
    stripped = line.strip()
    low = stripped.lower()
    if low.startswith(_INTRO_PREFIX):
        return True
    if low.startswith("<!--") and stripped.endswith("-->"):
        return stripped[4:-3].strip().lower().startswith(_INTRO_PREFIX)
    return False


def _stage_metadata(staged, dest):
    """os.replace swaps the directory entry, so the staged inode's root owner would
    become the published file's, undoing the sync script's care to leave a
    bind-mounted notebook host-owned. Best effort."""
    try:
        st = os.stat(dest)
    except OSError:
        return
    try:
        os.chmod(staged, stat.S_IMODE(st.st_mode))
    except OSError:
        pass
    try:
        os.chown(staged, st.st_uid, st.st_gid)
    except (OSError, AttributeError):
        pass


def _strip_lines(lines):
    for i, line in enumerate(lines):
        if _is_intro_line(line):
            out = lines[:i] + lines[i + 1 :]
            if i < len(out) and out[i].strip() == "":
                out = out[:i] + out[i + 1 :]
            return out
    return None


def _strip_cell(cell):
    src = cell.get("source")
    if isinstance(src, str):
        lines = src.splitlines(keepends = True)
        as_str = True
    elif isinstance(src, list):
        lines = list(src)
        as_str = False
    else:
        return False
    new_lines = _strip_lines(lines)
    if new_lines is None:
        return False
    cell["source"] = "".join(new_lines) if as_str else new_lines
    return True


def _strip_intro(nb):
    """Strip the Colab intro from the LEADING markdown block: cells[0] alone misses the
    notebooks that put the badge there and the sentence in cells[1]. The scan stops at
    the first non-markdown cell, so it never reaches prose between code cells."""
    cells = nb.get("cells")
    if not isinstance(cells, list):
        return False
    changed = False
    for cell in cells:
        if not isinstance(cell, dict) or cell.get("cell_type") != "markdown":
            break  # the first code cell ends the header block
        if _strip_cell(cell):
            changed = True
    return changed


def _clean_widgets(nb):
    changed = False
    cells = nb.get("cells")
    if isinstance(cells, list):
        for cell in cells:
            if not isinstance(cell, dict):
                continue
            outs = cell.get("outputs")
            if not isinstance(outs, list):
                continue
            kept = [
                o
                for o in outs
                if not (isinstance(o, dict) and _WIDGET_VIEW_MIME in (o.get("data") or {}))
            ]
            if len(kept) != len(outs):
                cell["outputs"] = kept
                changed = True
    md = nb.get("metadata")
    if isinstance(md, dict) and "widgets" in md:
        del md["widgets"]
        changed = True
    return changed


def _unlink(path):
    try:
        os.remove(path)
    except OSError:
        pass


def _stage_clean(path):
    """Returns (tmp, hash_before, hash_after) for a cleaned copy staged beside the notebook but not
    published, or None if nothing changed or the write failed; split out so migrate can record
    the new hash first."""
    try:
        before = _sha256(path)
        with open(path, "r", encoding = "utf-8") as f:
            nb = json.load(f)
    except Exception:
        return None

    changed = _strip_intro(nb)
    changed = _clean_widgets(nb) or changed
    if not changed:
        return None

    tmp = path + ".tmp"
    try:
        with open(tmp, "w", encoding = "utf-8") as f:
            json.dump(nb, f, indent = 1, ensure_ascii = False)
            f.write("\n")
        return tmp, before, _sha256(tmp)
    except Exception:
        _unlink(tmp)
        return None


def _publish(tmp, path, before):
    """Move the staged copy onto the notebook. False if it was not published."""
    try:
        # JupyterLab is serving the tree, so recheck before replacing.
        if _sha256(path) != before:
            _unlink(tmp)
            return False
        _stage_metadata(tmp, path)
        os.replace(tmp, path)
    except Exception:
        _unlink(tmp)
        return False
    return True


def _write_state(state_path, lines):
    """Replaces the state file durably, returning False on failure; fsync runs before the rename,
    since a crash between them would leave the rename visible with the content unwritten."""
    tmp = state_path + ".tmp"
    try:
        with open(tmp, "w", encoding = "utf-8") as f:
            f.write("\n".join(lines) + "\n")
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, state_path)
    except OSError:
        _unlink(tmp)
        return False
    return True


def _resume(path, rec):
    """Finishes a publish interrupted between the state and notebook renames, True if it did; that
    stop would otherwise read as a user edit, and only that case reproduces the record byte for
    byte."""
    st = _stage_clean(path)
    if st is None:
        return False
    tmp, before, after = st
    if after != rec:
        _unlink(tmp)
        return False
    return _publish(tmp, path, before)


def strip_notebook(path, staged = None):
    """True if the notebook was modified and written back. `staged`, when given, gets
    {"sha256": ...} for the bytes THIS call wrote, so the caller need not re-read."""
    st = _stage_clean(path)
    if st is None:
        return False
    tmp, before, after = st
    if not _publish(tmp, path, before):
        return False
    if staged is not None:
        staged["sha256"] = after
    return True


def _sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def migrate(state_path, dest):
    """Cleans every owned notebook, recording each in the state file before it is published, so one
    lost state write cannot strand the set; a stop between the renames is reconciled by _resume."""
    try:
        with open(state_path, "r", encoding = "utf-8") as f:
            lines = f.read().splitlines()
    except OSError:
        return 0

    out = list(lines)  # malformed lines and untouched records survive verbatim
    changed = 0
    stopped = False
    for i, line in enumerate(lines):
        parts = line.split("  ", 1)  # "<sha256>  <relpath>"
        if len(parts) != 2:
            continue
        rec, rel = parts
        path = os.path.join(dest, rel)
        if not (rel.endswith(".ipynb") and os.path.isfile(path)):
            continue
        try:
            digest = _sha256(path)
        except OSError:
            continue
        if digest != rec:  # not ours, the user has edited it, or a publish was cut off
            if _resume(path, rec):
                changed += 1
            continue
        staged = _stage_clean(path)
        if staged is None:
            continue
        tmp, before, after = staged
        out[i] = "%s  %s" % (after, rel)
        if not _write_state(state_path, out):
            _unlink(tmp)
            out[i] = line
            stopped = True
            break
        if _publish(tmp, path, before):
            changed += 1
        else:
            # nothing landed, so take the record back off
            out[i] = line
            _write_state(state_path, out)

    if changed:
        print(f"[unsloth-nb] cleaned {changed} notebook(s) (Colab intro + widget outputs)")
    if stopped:
        print(
            "[unsloth-nb] could not record the cleaned notebooks; leaving the rest for "
            "the next start"
        )
    return 0


def main(argv):
    ap = argparse.ArgumentParser(description = "Strip the Colab-only intro sentence.")
    ap.add_argument("--state", help = "sync state file (enables migration mode)")
    ap.add_argument("--dest", help = "notebooks dir (with --state)")
    ap.add_argument("paths", nargs = "*", help = "notebooks to strip in place")
    args = ap.parse_args(argv)

    if args.state:
        if not args.dest:
            ap.error("--state requires --dest")
        return migrate(args.state, args.dest)

    changed = sum(1 for p in args.paths if strip_notebook(p))
    if changed:
        print(f"[unsloth-nb] cleaned {changed} notebook(s) (Colab intro + widget outputs)")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))