# patchworks workflow: convert -> prepare -> segment (GPU jobs) -> merge.
#   pixi run dry | go | slurm

import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(workflow.basedir) / "scripts"))
# patchworks from this clone, even while pixi reinstalls it (scripts/_src.py).
from _src import use_clone_src

use_clone_src()
from _pw import validate_config, with_defaults

# Settings a config leaves out take their defaults (scripts/_pw.py DEFAULTS).
config.update(with_defaults(config))
# Refuse a bad config now, dry runs included, not in prepare after hours of
# converting. Only here, in the main process: every cluster job reads this
# file again, and some checks import Cellpose.
# run_multi checked every config before starting (and says so): checking
# again here would load Cellpose in each of its Snakemakes at once, which
# all share the controller job's memory.
if getattr(workflow, "is_main_process", True) and not os.environ.get(
    "PATCHWORKS_CONFIG_CHECKED"
):
    validate_config(config)

include: "rules/common.smk"
include: "rules/convert.smk"
include: "rules/segment.smk"
include: "rules/merge.smk"


localrules:
    fetch_model,


rule all:
    input:
        f"{RUN}/labels.done",


# Mail with the failing step's log tail (SLURM's own mail has no traceback).
onerror:
    from patchworks._notify import failing_step, log_tail, send

    _rule, _log = failing_step(log)
    if _log is None:
        _cands = [p for p in (CONVERTLOG, OCCUPANCYLOG, PREPARELOG, MERGELOG)
                  if Path(p).exists()]
        _cands += [str(p) for p in Path(LOGS).glob("segment/*.log")]
        _log = max(_cands, key=lambda p: Path(p).stat().st_mtime, default=None)
    _what = f"rule {_rule}" if _rule else "the workflow"
    send(
        NOTIFY_EMAIL,
        f"[patchworks] FAILED: {_rule or 'workflow'} - {LABEL_NAME}",
        f"The patchworks workflow failed in {_what}.\n\n"
        f"  work_dir      : {WORK}\n"
        f"  config        : {LABEL_NAME}\n"
        f"  failing step  : {_rule or 'unknown'}\n"
        f"  step log      : {_log}\n"
        f"  snakemake log : {log}\n\n"
        + (f"Last lines of {_log}:\n\n{log_tail(_log)}\n"
           if _log else "No step log was written.\n"),
    )


onsuccess:
    if "finish" in NOTIFY_EVENTS:
        from patchworks._notify import done_message, send

        # What was asked for: the labels (the default), or only the
        # conversion (the launcher's "Convert only", run_multi's first
        # phase) -- the mail says which.
        _settings = getattr(workflow, "dag_settings", None)
        send(
            NOTIFY_EMAIL,
            *done_message(
                getattr(_settings, "targets", None) or [],
                work=WORK,
                image=IMAGE,
                occupancy=OCCUPANCY,
                label_name=LABEL_NAME,
            ),
        )
