# Copyright (C) 2022 Justin Schuh <code@justinschuh.com>
#
# This file may be distributed under the terms of the GNU GPLv3 license.
#
# Credit to original source:
# https://klipper.discourse.group/t/saving-and-adjusting-per-build-surface-z-offsets/696 

[gcode_macro _apply_bed_surface_offset]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% if surfaces.active %}
    SET_SURFACE_ACTIVE SURFACE={surfaces.active}
  {% endif %}

[gcode_macro _init_surfaces]
gcode:
  {% set km = printer["gcode_macro _km_globals"] %}
  {% if "bed_surfaces" in printer.save_variables.variables %}
    {% set old_surfaces = printer.save_variables.variables.bed_surfaces %}
  {% else %}
    {% set old_surfaces = { 'active' : '', 'available' : {} } %}
  {% endif %}
  {% set settings = printer.configfile.settings %}
  {% set new_probe_z = (settings.probe | default(settings.bltouch) |
                        default(settings.smart_effector) | default({})
                        ).z_offset|default(0.0)|float %}
  {% set new_endstop_z = (settings.stepper_z | default({})).position_endstop |
                        default(0.0)|float %}
  {% if 'endstop_z' not in old_surfaces %}
    {% set dummy = old_surfaces.__setitem__('endstop_z', new_endstop_z) %}
  {% endif %}
  {% if 'probe_z' not in old_surfaces %}
    {% set dummy = old_surfaces.__setitem__('probe_z', new_probe_z) %}
  {% endif %}
  {% set surfaces = { 'active' : '', 'available' : {},
                      'endstop_z' : old_surfaces.endstop_z,
                      'probe_z' : old_surfaces.probe_z } %}
  {% for s in km.bed_surfaces %}
    {% set s = s.split()|join(' ')|lower %}
    {% if s|length > km.bed_surface_max_name_length or
          s|list|select("in", " \r\n\"\'")|list %}
      {action_raise_error('Invalid surface name "%s". Name must be %d or fewer '
        'characters and must not include space or quotation characters'
        | format(s, km.bed_surface_max_name_length))}
    {% endif %}
    {% if s in old_surfaces.available %}
      {% set dummy = surfaces.available.__setitem__(s,
                       old_surfaces.available[s]) %}
    {% else %}
      {% set dummy = surfaces.available.__setitem__(s, {'offset' : 0.0}) %}
    {% endif %}
  {% endfor %}
  {% if old_surfaces.active in surfaces.available %}
    {% set dummy = surfaces.__setitem__('active', old_surfaces.active) %}
  {% elif km.bed_surfaces %}
    {% set dummy = surfaces.__setitem__('active', km.bed_surfaces[0]|lower) %}
  {% endif %}
  SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
  _APPLY_BED_SURFACE_OFFSET
  {% if new_probe_z != surfaces.probe_z or
        new_endstop_z != surfaces.endstop_z %}
    { action_respond_info(
      'Z probe offset or endstop position changed. Run ADJUST_SURFACE_OFFSETS '
      'to adjust the offset for all saved surfaces by the change differential, '
      'or run ADJUST_SURFACE_OFFSETS IGNORE=1 to hide this message without '
      'making changes.') }
  {% endif %}

[gcode_macro adjust_surface_offsets]
description: Adjusts surface offsets to account for changes in the Z endstop position or probe Z offset.
  Usage: ADJUST_SURFACE_OFFSETS [IGNORE]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% set settings = printer.configfile.settings %}
  {% set new_probe_z = (settings.probe | default(settings.bltouch) |
                        default(settings.smart_effector) | default({})
                        ).z_offset|default(0.0)|float %}
  {% set new_endstop_z = (settings.stepper_z | default({})).position_endstop |
                          default(0.0)|float %}
  {% set diff = (surfaces.probe_z - new_probe_z +
                 surfaces.endstop_z - new_endstop_z)|round(6) %}
  {% if not params.IGNORE|default(0)|int %}
    {% for s in surfaces.available %}
      {% set offset = (surfaces.available[s].offset - diff)|round(6) %}
      {% set dummy = surfaces.available.__setitem__(s, {'offset' : offset}) %}
    {% endfor %}
    { action_respond_info("All bed surfaces now adjusted by %1.4f"|
                          format(diff))}
  {% elif diff != 0 %}
    { action_respond_info("Status cleared without adjustment") }
  {% endif %}
  {% set dummy = surfaces.__setitem__('endstop_z', new_endstop_z| round(6)) %}
  {% set dummy = surfaces.__setitem__('probe_z', new_probe_z|round(6)) %}
  SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"

[gcode_macro set_surface_offset]
description: Sets the offset for a surface and moves the toolhead (if homed).
  Usage: SET_SURFACE_OFFSET [OFFSET=<offset>] [SURFACE=<surface>]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% set SURFACE = params.SURFACE|default(surfaces.active)|lower %}
  {% if SURFACE not in surfaces.available %}
    { action_raise_error("Bed surface %s does not exist." | format(SURFACE)) }
  {% endif %}
  {% set active = surfaces.available[SURFACE] %}
  # If no offset is provided just print out the current offset.
  {% set OFFSET = params.OFFSET|default(active.offset)|float %}
  {% if OFFSET != active.offset %}
    {% set dummy = surfaces.available[SURFACE].__setitem__("offset", OFFSET) %}
    SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% if SURFACE == surfaces.active %}
      _km_set_gcode_offset_base Z="{OFFSET}" MOVE={
        1 if printer.toolhead.homed_axes == 'xyz' else 0}
    {% endif %}
  {% endif %}
  { action_respond_info("Bed surface: %s Offset: %.3f"
                        | format(SURFACE, OFFSET)) }
  # Dummy argument block for Mainsail
  {% set dummy = None if True else "
  {% set dummy = params.SURFACE|default(active surface) %}
  {% set dummy = params.OFFSET|default(none)|float %}
  " %} # End argument block for Mainsail

[gcode_macro set_surface_active]
description: Sets the active bed surface and moves the toolhead (if homed). If no SURFACE argument is present the available surfaces are listed and the active one is preceded by a "*".
  Usage: SET_SURFACE_ACTIVE [SURFACE=<surface>]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% if "SURFACE" in params %}
    {% set SURFACE = params.SURFACE|lower %}
    {% if SURFACE not in surfaces.available %}
      { action_raise_error("Bed surface %s does not exist." | format(SURFACE)) }
    {% endif %}
    {% if SURFACE != surfaces.active %}
      {% set dummy = surfaces.__setitem__("active", SURFACE) %}
      SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% endif %}
    {% if surfaces.available[SURFACE].offset !=
          printer.gcode_move.homing_origin.z %}
      _km_set_gcode_offset_base Z="{surfaces.available[SURFACE].offset
        }" MOVE={1 if printer.toolhead.homed_axes == 'xyz' else 0}
    {% endif %}
    {action_respond_info("Active bed surface: %s; offset: %.3f"
                        | format(SURFACE, surfaces.available[SURFACE].offset))}
  {% else %}
    {% set output = [] %}
    {% for s in surfaces.available|list|sort %}
      {% set dummy = output.append("%s&nbsp;%s - offset: %.3f" 
                      | format("*" if s == surfaces.active else "&nbsp;",
                                s, surfaces.available[s].offset)) %}
    {% endfor %}
    {action_respond_info(output|join('\n'))}
  {% endif %}
  # Dummy argument block for Mainsail
  {% set dummy = None if True else "
  {% set dummy = params.SURFACE|default(none) %}
  " %} # End argument block for Mainsail

[gcode_macro set_gcode_offset]
description: Wraps SET_GCODE_OFFSET to update the current bed sheet offset.
  Usage: SET_GCODE_OFFSET [X=<pos>|X_ADJUST=<adjust>]
                          [Y=<pos>|Y_ADJUST=<adjust>]
                          [Z=<pos>|Z_ADJUST=<adjust>]
                          [MOVE=1 [MOVE_SPEED=<speed>]]
rename_existing: _KM_SET_GCODE_OFFSET_BASE
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% if surfaces.active and
        not printer["gcode_macro _km_save_state"].is_ephemeral %}
    {% set Z = params.Z|default(0.0)|float|round(6) %}
    {% set Z_ADJUST = params.Z_ADJUST|default(0.0)|float %}
    {% if 'Z' in params and
                 Z != surfaces.available[surfaces.active].offset %}
      {% set dummy = surfaces.available[surfaces.active].__setitem__("offset",
                       Z) %}
      SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% elif Z_ADJUST != 0.0 %}
      {% set dummy = surfaces.available[surfaces.active].__setitem__(
        "offset", (Z_ADJUST + printer.gcode_move.homing_origin.z)|round(6)) %}
      SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% endif %}
  {% endif %}
  _km_set_gcode_offset_base{% for k in params%}{' '~k~'="'~params[k]~'"'
    }{% endfor %}