import sys

path = "/home/slin/enduro-trails/prototype/static/app.js"
with open(path, 'r') as f:
    content = f.read()

# Add onTrailsCheckbox function after onTerrainCheckbox function
js_code = """
function onTrailsCheckbox() {
  const map = window._map;
  if (!map) return;
  
  const trackChecked = document.getElementById('trails-track-cb').checked;
  const pathChecked = document.getElementById('trails-path-cb').checked;
  
  // Save state
  localStorage.setItem('trails-track', trackChecked ? '1' : '0');
  localStorage.setItem('trails-path', pathChecked ? '1' : '0');
  
  // Toggle layer visibility
  if (map.getLayer('trails-track')) {
    map.setLayoutProperty('trails-track', 'visibility', trackChecked ? 'visible' : 'none');
  }
  if (map.getLayer('trails-path-bridleway')) {
    map.setLayoutProperty('trails-path-bridleway', 'visibility', pathChecked ? 'visible' : 'none');
  }
}

function restoreTrailsState() {
  const trackState = localStorage.getItem('trails-track');
  const pathState = localStorage.getItem('trails-path');
  
  // Default: both checked (visible)
  const trackOn = trackState === null || trackState === '1';
  const pathOn = pathState === null || pathState === '1';
  
  const trackCb = document.getElementById('trails-track-cb');
  const pathCb = document.getElementById('trails-path-cb');
  
  if (trackCb) trackCb.checked = trackOn;
  if (pathCb) pathCb.checked = pathOn;
  
  const map = window._map;
  if (map) {
    if (map.getLayer('trails-track')) {
      map.setLayoutProperty('trails-track', 'visibility', trackOn ? 'visible' : 'none');
    }
    if (map.getLayer('trails-path-bridleway')) {
      map.setLayoutProperty('trails-path-bridleway', 'visibility', pathOn ? 'visible' : 'none');
    }
  }
}
"""

# Insert after onTerrainCheckbox function (find the closing brace after applyTerrainLayer calls)
# Look for "function applyTerrainLayer" and insert before it
marker = "function applyTerrainLayer("
idx = content.find(marker)
if idx == -1:
    print("ERROR: marker not found")
    sys.exit(1)

content = content[:idx] + js_code + "\n" + content[idx:]

# Also add restoreTrailsState() call in restoreTerrainState or initTerrain
# Find restoreTerrainState call and add restoreTrailsState after it
restore_marker = "restoreTerrainState();"
# Find the last occurrence (in initTerrain)
last_idx = content.rfind(restore_marker)
if last_idx != -1:
    insert_pos = last_idx + len(restore_marker)
    content = content[:insert_pos] + "\n    restoreTrailsState();" + content[insert_pos:]

# Also after style.load restoreTerrainState
first_idx = content.find(restore_marker)
if first_idx != -1 and first_idx != last_idx:
    insert_pos = first_idx + len(restore_marker)
    content = content[:insert_pos] + "\n      restoreTrailsState();" + content[insert_pos:]

with open(path, 'w') as f:
    f.write(content)

print("OK - JS patched")