import sys, os

src = os.path.join(os.path.dirname(__file__), '../../frontend/js/labs/opticsbench.js')

with open(src, 'r', encoding='utf-8') as f:
    content = f.read()

# ─────────────────────────────────────────────────────────────────────────────
# PATCH 1: Insert _drawRaysAnimated + _drawArrowLabels before _bindEvents
# in ThinLensSim
# ─────────────────────────────────────────────────────────────────────────────

MARKER_BIND = '  _bindEvents() {\n    const cv = this.canvas;\n    const getPos = (e) => {'
idx = content.find(MARKER_BIND)
if idx == -1:
    print('ERROR: _bindEvents marker not found'); sys.exit(1)

NEW_METHODS_1 = (
    '  /* === _drawRaysAnimated: principal rays with per-ray progress === */\n'
    '  _drawRaysAnimated(ctx, lx, ay, d, h, f, dPrime, hPrime) {\n'
    '    const T = this._rayAnimT;\n'
    '    if (T[0] >= 1 && T[1] >= 1 && T[2] >= 1) { this._drawRays(ctx, lx, ay, d, h, f, dPrime, hPrime); return; }\n'
    '    const objX = lx - d, objY = ay - h;\n'
    '    const hasImage = dPrime !== null && isFinite(dPrime);\n'
    '    const isVirtual = hasImage && dPrime < 0;\n'
    '    const COLORS = [\'#06D6E0\', \'#7BF5A4\', \'#FFD166\'];\n'
    '    ctx.lineWidth = 1.8;\n'
    '    const lerp = (a, b, t) => a + (b - a) * Math.min(1, Math.max(0, t));\n'
    '    const drawPts = (color, pts, t) => {\n'
    '      if (t <= 0 || pts.length < 2) return;\n'
    '      const totalLen = pts.reduce((s, p, i) => i === 0 ? 0 : s + Math.hypot(p[0]-pts[i-1][0], p[1]-pts[i-1][1]), 0);\n'
    '      const target = totalLen * t;\n'
    '      const draw = () => {\n'
    '        ctx.strokeStyle = color; ctx.setLineDash([]);\n'
    '        ctx.beginPath(); ctx.moveTo(pts[0][0], pts[0][1]);\n'
    '        let drawn = 0;\n'
    '        for (let i = 1; i < pts.length; i++) {\n'
    '          const segLen = Math.hypot(pts[i][0]-pts[i-1][0], pts[i][1]-pts[i-1][1]);\n'
    '          if (drawn + segLen <= target) { ctx.lineTo(pts[i][0], pts[i][1]); drawn += segLen; }\n'
    '          else { const fr = segLen > 0 ? (target - drawn) / segLen : 0; ctx.lineTo(lerp(pts[i-1][0], pts[i][0], fr), lerp(pts[i-1][1], pts[i][1], fr)); break; }\n'
    '        }\n'
    '        ctx.stroke();\n'
    '      };\n'
    '      if (window.LabFX) LabFX.glow.drawGlow(ctx, draw, { color, intensity: 10 });\n'
    '      else draw();\n'
    '    };\n'
    '    const FAR = lx + 360;\n'
    '    const imgX = hasImage ? lx + dPrime : null, imgY = hasImage ? ay - hPrime : null;\n'
    '    // Ray 1: parallel to axis -> through F\'\n'
    '    if (T[0] > 0) {\n'
    '      let pts;\n'
    '      if (!hasImage) { pts = [[objX, objY], [lx, objY], [FAR, objY]]; }\n'
    '      else if (!isVirtual) { pts = [[objX, objY], [lx, objY], [imgX, imgY]]; }\n'
    '      else { const s = (objY - ay) / f; pts = [[objX, objY], [lx, objY], [FAR, objY + s*(FAR-lx)]]; }\n'
    '      drawPts(COLORS[0], pts, T[0]);\n'
    '    }\n'
    '    // Ray 2: through optical center (straight)\n'
    '    if (T[1] > 0) {\n'
    '      const s = (objY - ay) / (objX - lx);\n'
    '      drawPts(COLORS[1], [[objX, objY], [FAR, ay + s*(FAR-lx)]], T[1]);\n'
    '    }\n'
    '    // Ray 3: through front focus F -> parallel after lens\n'
    '    if (T[2] > 0) {\n'
    '      const fx = lx - f, s = (objY - ay) / (objX - fx);\n'
    '      const hitY = objY + s * (lx - objX);\n'
    '      const endX = hasImage && !isVirtual ? Math.max(imgX + 60, FAR) : FAR;\n'
    '      drawPts(COLORS[2], [[objX, objY], [lx, hitY], [endX, hitY]], T[2]);\n'
    '    }\n'
    '  }\n'
    '\n'
    '  /* === Arrow labels: h_o, h_i, magnification Gamma === */\n'
    '  _drawArrowLabels(ctx, lx, ay, d, h, dPrime, hPrime) {\n'
    '    const objX = lx - d;\n'
    '    ctx.font = \'11px Manrope, system-ui, sans-serif\'; ctx.textBaseline = \'middle\';\n'
    '    ctx.fillStyle = \'rgba(155,93,229,0.85)\'; ctx.textAlign = \'right\';\n'
    '    ctx.fillText(\'ho=\' + h.toFixed(0), objX - 6, ay - h / 2);\n'
    '    if (dPrime !== null && isFinite(dPrime)) {\n'
    '      const imgX = lx + dPrime, isVirtual = dPrime < 0;\n'
    '      const M = -dPrime / d;\n'
    '      const Gstr = isFinite(M) ? (M >= 0 ? \'+\' : \'\') + M.toFixed(2) : \'---\';\n'
    '      const imgColor = isVirtual ? \'rgba(255,133,162,0.85)\' : \'rgba(6,214,224,0.85)\';\n'
    '      ctx.fillStyle = imgColor; ctx.textAlign = \'left\';\n'
    '      ctx.fillText("hi=" + Math.abs(hPrime).toFixed(0), imgX + 6, ay - hPrime / 2);\n'
    '      ctx.fillStyle = \'#FFD166\'; ctx.textAlign = \'center\';\n'
    '      ctx.fillText(\'G=\' + Gstr, (lx + imgX) / 2, ay + 60);\n'
    '    }\n'
    '  }\n'
    '\n'
)

new_content = content[:idx] + NEW_METHODS_1 + content[idx:]
content = new_content

# ─────────────────────────────────────────────────────────────────────────────
# PATCH 2: Add R slider + parabolic/spherical toggle to MirrorSim constructor
# ─────────────────────────────────────────────────────────────────────────────
# Find MirrorSim constructor and add _R and _parabolic fields
OLD_MIRROR_CTOR = '    this._photonPaths   = [];\n\n    this._prevType = \'concave\';\n    this._transT   = 1.0;\n    this._transRaf = null;\n\n    this._drag   = null;\n    this._hoverX = -999;\n    this._hoverY = -999;\n\n    this.onUpdate  = null;\n    this.onAnimate = null;\n\n    this._bindEvents();\n    new ResizeObserver(() => { this.fit(); this.draw(); }).observe(canvas.parentElement);\n  }'

NEW_MIRROR_CTOR = (
    '    this._photonPaths   = [];\n\n'
    '    this._prevType = \'concave\';\n'
    '    this._transT   = 1.0;\n'
    '    this._transRaf = null;\n\n'
    '    this._drag   = null;\n'
    '    this._hoverX = -999;\n'
    '    this._hoverY = -999;\n\n'
    '    this.onUpdate  = null;\n'
    '    this.onAnimate = null;\n\n'
    '    /* Feature 2: R slider + spherical aberration toggle */\n'
    '    this._R           = 240;    // radius of curvature (positive=concave, negative=convex)\n'
    '    this._useR        = false;  // true = R-slider mode; false = classic type+f mode\n'
    '    this._parabolic   = false;  // false = spherical mirror; true = perfect parabolic\n'
    '\n'
    '    this._bindEvents();\n'
    '    new ResizeObserver(() => { this.fit(); this.draw(); }).observe(canvas.parentElement);\n'
    '  }'
)

if OLD_MIRROR_CTOR not in content:
    print('ERROR: MirrorSim ctor block not found'); sys.exit(1)

content = content.replace(OLD_MIRROR_CTOR, NEW_MIRROR_CTOR, 1)

# ─────────────────────────────────────────────────────────────────────────────
# PATCH 3: Add setMirrorR, setMirrorParabolic, _drawAberrationFan
#          before MirrorSim._bindEvents
# ─────────────────────────────────────────────────────────────────────────────
# Find the _bindEvents in MirrorSim (it's further down, after chk method)
MIRROR_BIND_MARKER = '  _bindEvents() {\n    const cv = this.canvas;\n    const getPos = e => {'

idx2 = content.find(MIRROR_BIND_MARKER)
if idx2 == -1:
    print('ERROR: chk marker not found'); sys.exit(1)

NEW_METHODS_2 = (
    '  /* === Feature 2: R-slider mode for MirrorSim === */\n'
    '  setMirrorR(R) {\n'
    '    this._useR = true;\n'
    '    this._R = +R;\n'
    '    // Derive type and f from R\n'
    '    const absR = Math.abs(this._R);\n'
    '    if (absR < 5) { this.type = \'flat\'; }\n'
    '    else if (this._R > 0) { this.type = \'concave\'; this.f = absR / 2; }\n'
    '    else { this.type = \'convex\'; this.f = absR / 2; }\n'
    '    this.draw(); this._emit();\n'
    '  }\n'
    '\n'
    '  setMirrorParabolic(on) {\n'
    '    this._parabolic = !!on;\n'
    '    this.draw();\n'
    '  }\n'
    '\n'
    '  /* Draw 5 parallel rays showing spherical vs parabolic aberration */\n'
    '  _drawAberrationFan(ctx, mx, ay, f) {\n'
    '    if (!isFinite(f) || Math.abs(f) < 5) return;\n'
    '    const mH = Math.min(this.H * 0.38, 140);\n'
    '    const heights = [-0.85, -0.45, 0, 0.45, 0.85];\n'
    '    const COLORS = [\'#FF6B6B\', \'#FFD166\', \'#7BF5A4\', \'#06D6E0\', \'#B8A4FF\'];\n'
    '    ctx.save(); ctx.lineWidth = 1.4;\n'
    '    heights.forEach((fr, i) => {\n'
    '      const rayH = fr * mH;\n'
    '      // For parabolic mirror: all parallel rays focus exactly at f\n'
    '      // For spherical: marginal rays (fr != 0) focus closer by h^2/(2R) approx\n'
    '      const fEff = this._parabolic ? f : f - (rayH * rayH) / (2 * Math.abs(f) * 2);\n'
    '      const startX = mx - this.d - 40;\n'
    '      const hitY   = ay - rayH; // hits mirror at height rayH\n'
    '      // Incident ray: horizontal from left to mirror\n'
    '      ctx.strokeStyle = COLORS[i]; ctx.globalAlpha = 0.75;\n'
    '      ctx.setLineDash([]);\n'
    '      ctx.beginPath(); ctx.moveTo(startX, ay - rayH); ctx.lineTo(mx, hitY); ctx.stroke();\n'
    '      // Reflected ray: goes toward focal point fEff\n'
    '      const focX = mx - fEff;\n'
    '      if (focX > 0 && focX < this.W) {\n'
    '        ctx.beginPath(); ctx.moveTo(mx, hitY); ctx.lineTo(focX, ay);\n'
    '        // extend a bit past focus\n'
    '        const dx = focX - mx, dy = ay - hitY, len = Math.hypot(dx, dy);\n'
    '        if (len > 1) ctx.lineTo(focX + dx/len*50, ay + dy/len*50);\n'
    '        ctx.stroke();\n'
    '      }\n'
    '    });\n'
    '    ctx.globalAlpha = 1;\n'
    '    // label\n'
    '    const label = this._parabolic ? \'Параболическое (идеальный фокус)\' : \'Сферическое (аберрация)\';\n'
    '    const col   = this._parabolic ? \'#7BF5A4\' : \'#FF6B6B\';\n'
    '    const bx = 12, by = this.H - 36;\n'
    '    ctx.fillStyle = \'rgba(13,13,26,0.85)\';\n'
    '    ctx.beginPath(); ctx.roundRect(bx, by, 250, 24, 6); ctx.fill();\n'
    '    ctx.font = \'bold 11px Manrope, system-ui, sans-serif\';\n'
    '    ctx.textAlign = \'left\'; ctx.textBaseline = \'middle\';\n'
    '    ctx.fillStyle = col;\n'
    '    ctx.fillText(label, bx + 8, by + 12);\n'
    '    ctx.restore();\n'
    '  }\n'
    '\n'
)

new_content2 = content[:idx2] + NEW_METHODS_2 + content[idx2:]
content = new_content2

with open(src, 'w', encoding='utf-8') as f:
    f.write(content)

print('OK lines:', content.count('\n'))