// wfc_solver_worker.js
let grid,
  pCols,
  pRows,
  patternsW,
  patternIdToIndexW,
  adjacencyW,
  currentPatternSizeW_solver;
let queue, queuedCoordinatesSet, backtrackStack;
let currentAttemptFailed;
let maxTriesInWorker; // Renamed from maxTries to avoid confusion
let useRandomSelectionW;
let MAX_BACKTRACK_DEPTH_W;

// --- Worker Stats (subset for WFC run) ---
let solverRunStats = {};
function resetSolverRunStats() {
  solverRunStats = {
    maxBacktrackStackDepthReached: 0,
    totalBacktracks: 0,
    totalContradictions: 0,
    forcedBacktracksByDepth: 0,
    cellsCollapsed: 0,
    propagationSteps: 0,
    // startTime/endTime/durationMs will be managed by main thread based on messages
  };
}
let solverMemoryStats = {
  estimatedGridSnapshotBytes: 0, // Will be calculated once grid is initialized
  backtrackStackPeakBytes: 0,
};

// --- Helper functions for solver (deepCopy, getPatternNeighbors, etc.) ---
function deepCopyGrid(gridToCopy) {
  if (!gridToCopy) return null;
  return gridToCopy.map((row) => row.map((cellOptions) => [...cellOptions]));
}
function deepCopyQueue(queueToCopy) {
  if (!queueToCopy) return null;
  return queueToCopy.map((item) => (Array.isArray(item) ? [...item] : item));
}
function getPatternNeighbors(c, r) {
  return [
    [c - 1, r, 'left'],
    [c + 1, r, 'right'],
    [c, r - 1, 'up'],
    [c, r + 1, 'down'],
  ].filter(([nc, nr]) => nc >= 0 && nr >= 0 && nc < pCols && nr < pRows);
}

function initializeGridInWorker() {
  const allPatternIndices = Array.from(
    { length: patternsW.length },
    (_, i) => i
  );
  grid = Array.from({ length: pRows }, () =>
    Array.from({ length: pCols }, () => [...allPatternIndices])
  );
  queue = [];
  queuedCoordinatesSet = new Set();
  maxTriesInWorker = pCols * pRows * 30 * currentPatternSizeW_solver; // Original logic for max tries
  currentAttemptFailed = false;
  backtrackStack = [];
  resetSolverRunStats(); // Reset stats for this run

  // Estimate memory for one grid snapshot (for backtrack stack calculations)
  let oneGridBytes = 16; // Outer array
  if (grid && grid.length > 0 && grid[0] && grid[0].length > 0) {
    oneGridBytes += grid.length * 16; // Each row array
    const cell = grid[0][0];
    const bytesPerIndex =
      patternsW.length > 65535 ? 4 : patternsW.length > 255 ? 2 : 1;
    const cellBytes = 16 + cell.length * bytesPerIndex;
    oneGridBytes += grid.length * grid[0].length * cellBytes;
  }
  solverMemoryStats.estimatedGridSnapshotBytes = oneGridBytes;
  solverMemoryStats.backtrackStackPeakBytes = 0; // Reset for this run

  // Send initial grid state (all tiles are "dirty" with all options)
  const initialUpdates = [];
  for (let r = 0; r < pRows; r++) {
    for (let c = 0; c < pCols; c++) {
      initialUpdates.push({ c, r, cellOptions: [...grid[r][c]] });
    }
  }
  if (initialUpdates.length > 0) {
    self.postMessage({
      type: 'TILE_UPDATE',
      payload: { updates: initialUpdates },
    });
  }
  self.postMessage({
    type: 'STATS_UPDATE',
    payload: { memory: solverMemoryStats },
  });
}

function pickCellWithLowestEntropy() {
  /* ... Same as original ... */
  if (!grid) return null;
  let bestEntropy = Infinity;
  let bestCells = [];
  for (let r = 0; r < pRows; r++) {
    for (let c = 0; c < pCols; c++) {
      if (!grid[r] || !grid[r][c]) continue;
      const numOptions = grid[r][c].length;
      if (numOptions > 1) {
        const currentEntropy = numOptions + Math.random() * 0.1; // Add jitter
        if (currentEntropy < bestEntropy) {
          bestEntropy = currentEntropy;
          bestCells = [[c, r]];
        } else if (numOptions === Math.floor(bestEntropy)) {
          // Check only integer part for equality
          bestCells.push([c, r]);
        }
      }
    }
  }
  return bestCells.length > 0
    ? bestCells[Math.floor(Math.random() * bestCells.length)]
    : null;
}

function collapseCell(c, r) {
  /* ... Same as original, but use solverRunStats and post TILE_UPDATE ... */
  const initialOptionsForCellIndices = grid[r][c];
  if (initialOptionsForCellIndices.length === 0) {
    currentAttemptFailed = true;
    solverRunStats.totalContradictions++;
    return;
  }

  if (backtrackStack.length >= MAX_BACKTRACK_DEPTH_W) {
    // Post status about max depth reached
    self.postMessage({
      type: 'STATUS_UPDATE',
      payload: {
        message: `Max backtrack depth (${MAX_BACKTRACK_DEPTH_W}) reached. Forcing backtrack.`,
      },
    });
    currentAttemptFailed = true;
    solverRunStats.forcedBacktracksByDepth++;
    solverRunStats.totalContradictions++;
    return;
  }

  let optionsToSystematicallyTry = [...initialOptionsForCellIndices];
  // Randomize if checkbox was checked
  if (useRandomSelectionW && optionsToSystematicallyTry.length > 1) {
    for (let i = optionsToSystematicallyTry.length - 1; i > 0; i--) {
      const j = Math.floor(Math.random() * (i + 1));
      [optionsToSystematicallyTry[i], optionsToSystematicallyTry[j]] = [
        optionsToSystematicallyTry[j],
        optionsToSystematicallyTry[i],
      ];
    }
  } else if (!useRandomSelectionW && optionsToSystematicallyTry.length > 1) {
    // Weighted selection
    let totalWeight = 0;
    const weights = optionsToSystematicallyTry.map((idx) => {
      const pat = patternsW[idx]; // patternsW is array of pattern objects
      const w = pat && pat.frequency > 0 ? pat.frequency : 1;
      totalWeight += w;
      return w;
    });
    if (totalWeight === 0) {
      // Should not happen if frequencies are positive
      // Fallback to random if all weights are zero
    } else {
      let rnd = Math.random() * totalWeight;
      let chosenIndexInSystematicList;
      for (let i = 0; i < optionsToSystematicallyTry.length; i++) {
        rnd -= weights[i];
        if (rnd <= 0) {
          chosenIndexInSystematicList = i;
          break;
        }
      }
      if (typeof chosenIndexInSystematicList === 'undefined')
        chosenIndexInSystematicList = optionsToSystematicallyTry.length - 1; // Should not be needed

      // Move chosen to front
      const actualChosenItem = optionsToSystematicallyTry.splice(
        chosenIndexInSystematicList,
        1
      )[0];
      optionsToSystematicallyTry.unshift(actualChosenItem);
    }
  }

  let chosenPatternIndex = optionsToSystematicallyTry[0]; // Try the first one (either random or weighted first)

  if (typeof chosenPatternIndex === 'undefined') {
    currentAttemptFailed = true;
    solverRunStats.totalContradictions++;
    return;
  }

  const gridSnapshot = deepCopyGrid(grid);
  const queueSnapshot = deepCopyQueue(queue);
  const queuedCoordinatesSetSnapshot = new Set(queuedCoordinatesSet);

  backtrackStack.push({
    gridBeforeChoice: gridSnapshot,
    queueBeforeChoice: queueSnapshot,
    queuedCoordinatesSetBeforeChoice: queuedCoordinatesSetSnapshot,
    cellCoords: [c, r],
    optionsAtCell: optionsToSystematicallyTry,
    lastTriedOptionIndexInList: 0,
  });
  solverRunStats.maxBacktrackStackDepthReached = Math.max(
    solverRunStats.maxBacktrackStackDepthReached,
    backtrackStack.length
  );
  solverMemoryStats.backtrackStackPeakBytes = Math.max(
    solverMemoryStats.backtrackStackPeakBytes,
    backtrackStack.length * solverMemoryStats.estimatedGridSnapshotBytes
  );
  self.postMessage({
    type: 'STATS_UPDATE',
    payload: {
      wfcRun: {
        maxBacktrackStackDepthReached:
          solverRunStats.maxBacktrackStackDepthReached,
      },
      memory: {
        backtrackStackPeakBytes: solverMemoryStats.backtrackStackPeakBytes,
      },
    },
  });

  grid[r][c] = [chosenPatternIndex];
  solverRunStats.cellsCollapsed++;
  self.postMessage({
    type: 'TILE_UPDATE',
    payload: { updates: [{ c, r, cellOptions: grid[r][c] }] },
  });

  const coordKey = `${c},${r}`;
  if (!queuedCoordinatesSet.has(coordKey)) {
    queue.push([c, r]);
    queuedCoordinatesSet.add(coordKey);
  }
}

function propagateStep() {
  /* ... Same as original, but use solverRunStats and post TILE_UPDATE ... */
  solverRunStats.propagationSteps++;
  if (queue.length === 0) return { success: true };

  const [c, r] = queue.shift();
  queuedCoordinatesSet.delete(`${c},${r}`);

  if (!grid || !grid[r] || !grid[r][c]) {
    currentAttemptFailed = true;
    solverRunStats.totalContradictions++;
    return { success: false };
  }
  const currentPatternIndicesInCell = grid[r][c];
  if (currentPatternIndicesInCell.length === 0) {
    currentAttemptFailed = true;
    solverRunStats.totalContradictions++;
    return { success: false };
  }

  const tileUpdatesForPropagation = [];

  for (const [nc, nr, dirToNeighbor] of getPatternNeighbors(c, r)) {
    if (!grid[nr] || !grid[nr][nc]) {
      currentAttemptFailed = true;
      solverRunStats.totalContradictions++;
      return { success: false };
    }
    const originalNeighborOptionIndices = grid[nr][nc];
    if (originalNeighborOptionIndices.length === 0) {
      currentAttemptFailed = true;
      solverRunStats.totalContradictions++;
      return { success: false }; // Already a contradiction
    }

    let cumulativeAllowedNeighborPatternIndices = new Set();
    currentPatternIndicesInCell.forEach((p_idx) => {
      const p_id_str = patternsW[p_idx].id; // patternsW is array of pattern objects
      if (!adjacencyW[p_id_str]) return; // Should not happen with valid adjacency

      let compatibleIndices;
      if (dirToNeighbor === 'right')
        compatibleIndices = adjacencyW[p_id_str].right;
      else if (dirToNeighbor === 'left')
        compatibleIndices = adjacencyW[p_id_str].left;
      else if (dirToNeighbor === 'up')
        compatibleIndices = adjacencyW[p_id_str].up;
      else if (dirToNeighbor === 'down')
        compatibleIndices = adjacencyW[p_id_str].down;
      else return;

      compatibleIndices.forEach((q_idx) =>
        cumulativeAllowedNeighborPatternIndices.add(q_idx)
      );
    });

    const newNeighborOptionIndices = originalNeighborOptionIndices.filter(
      (idx) => cumulativeAllowedNeighborPatternIndices.has(idx)
    );

    if (newNeighborOptionIndices.length === 0) {
      currentAttemptFailed = true;
      solverRunStats.totalContradictions++;
      return { success: false };
    }

    if (
      newNeighborOptionIndices.length < originalNeighborOptionIndices.length
    ) {
      grid[nr][nc] = newNeighborOptionIndices;
      tileUpdatesForPropagation.push({
        c: nc,
        r: nr,
        cellOptions: grid[nr][nc],
      });

      const neighborCoordKey = `${nc},${nr}`;
      if (!queuedCoordinatesSet.has(neighborCoordKey)) {
        queue.push([nc, nr]);
        queuedCoordinatesSet.add(neighborCoordKey);
      }
    }
  }
  if (tileUpdatesForPropagation.length > 0) {
    self.postMessage({
      type: 'TILE_UPDATE',
      payload: { updates: tileUpdatesForPropagation },
    });
  }
  return { success: true };
}

function allCollapsed() {
  /* ... Same as original ... */
  if (!grid) return false;
  for (let r_idx = 0; r_idx < pRows; r_idx++) {
    for (let c_idx = 0; c_idx < pCols; c_idx++) {
      if (
        !grid[r_idx] ||
        !grid[r_idx][c_idx] ||
        grid[r_idx][c_idx].length > 1
      ) {
        return false;
      }
    }
  }
  return true;
}

let solverPhase = 'collapse'; // 'collapse' or 'propagate'
let runLoopTimeoutId = null;

function runStepInWorker() {
  if (runLoopTimeoutId) clearTimeout(runLoopTimeoutId); // Clear previous timeout if any

  if (solverPhase === 'stop') {
    // A way to stop the loop if main thread terminates worker
    return;
  }

  if (currentAttemptFailed) {
    currentAttemptFailed = false;
    let backtrackedSuccessfully = false;
    solverRunStats.totalBacktracks++;
    self.postMessage({
      type: 'STATS_UPDATE',
      payload: { wfcRun: { totalBacktracks: solverRunStats.totalBacktracks } },
    });

    // THIS IS THE CORRECTED BACKTRACKING LOOP
    while (backtrackStack.length > 0) {
        const lastDecision = backtrackStack.pop();
        const {
            gridBeforeChoice,
            queueBeforeChoice,
            queuedCoordinatesSetBeforeChoice,
            cellCoords,
            optionsAtCell,
            lastTriedOptionIndexInList,
        } = lastDecision;

        const [cc, cr] = cellCoords;
        let nextOptionIndexToTryInList = lastTriedOptionIndexInList + 1;

        if (nextOptionIndexToTryInList < optionsAtCell.length) {
            grid = deepCopyGrid(gridBeforeChoice); // Restore state
            queue = deepCopyQueue(queueBeforeChoice);
            queuedCoordinatesSet = new Set(queuedCoordinatesSetBeforeChoice);

            const newChosenPatternIndex = optionsAtCell[nextOptionIndexToTryInList];
            grid[cr][cc] = [newChosenPatternIndex];
            // Post update for the cell being retried
            self.postMessage({
                type: 'TILE_UPDATE',
                payload: { updates: [{ c: cc, r: cr, cellOptions: grid[cr][cc] }] },
            });

            const coordKeyBacktrack = `${cc},${cr}`;
            if (!queuedCoordinatesSet.has(coordKeyBacktrack)) {
                queue.push([cc, cr]);
                queuedCoordinatesSet.add(coordKeyBacktrack);
            }

            // THE FIX: Push the updated decision back onto the stack
            if (backtrackStack.length < MAX_BACKTRACK_DEPTH_W) {
                backtrackStack.push({
                    ...lastDecision,
                    lastTriedOptionIndexInList: nextOptionIndexToTryInList,
                });
                solverRunStats.maxBacktrackStackDepthReached = Math.max(
                    solverRunStats.maxBacktrackStackDepthReached,
                    backtrackStack.length
                );
                solverMemoryStats.backtrackStackPeakBytes = Math.max(
                    solverMemoryStats.backtrackStackPeakBytes,
                    backtrackStack.length * solverMemoryStats.estimatedGridSnapshotBytes
                );
                self.postMessage({
                    type: 'STATS_UPDATE',
                    payload: {
                        wfcRun: {
                            maxBacktrackStackDepthReached:
                            solverRunStats.maxBacktrackStackDepthReached,
                        },
                        memory: {
                            backtrackStackPeakBytes:
                            solverMemoryStats.backtrackStackPeakBytes,
                        },
                    },
                });
            }

            self.postMessage({
                type: 'STATUS_UPDATE',
                payload: {
                message: `Backtracked. Retrying option ${
                    nextOptionIndexToTryInList + 1
                }/${optionsAtCell.length} at (${cc},${cr}). Stack: ${
                    backtrackStack.length
                }`,
                },
            });
            solverPhase = 'propagate';
            backtrackedSuccessfully = true;

            const allTilesUpdate = [];
            for (let r_idx = 0; r_idx < pRows; r_idx++) {
                for (let c_idx = 0; c_idx < pCols; c_idx++) {
                    if (grid[r_idx] && grid[r_idx][c_idx]) {
                    allTilesUpdate.push({
                        c: c_idx,
                        r: r_idx,
                        cellOptions: [...grid[r_idx][c_idx]],
                    });
                    }
                }
            }
            if (allTilesUpdate.length > 0) {
                self.postMessage({
                    type: 'TILE_UPDATE',
                    payload: { updates: allTilesUpdate },
                });
            }
            break; // Exit backtrack loop
        }
    }


    if (!backtrackedSuccessfully) {
      self.postMessage({
        type: 'REQUEST_RESTART_FROM_SOLVER',
        payload: {
          message:
            'Backtrack stack exhausted -> WFC failed in worker. Requesting restart.',
          wfcStatus: 'Failed (Backtrack Exhausted in Worker)',
          currentRunStats: { wfcRun: solverRunStats }, // Send current stats before restart
        },
      });
      solverPhase = 'stop'; // Stop this worker's loop
      return;
    }
  }

if (allCollapsed()) {
    solverPhase = 'stop';
    // Add the final grid to the payload. This is the definitive final state.
    self.postMessage({
      type: 'WFC_COMPLETE',
      payload: { 
          finalStats: { wfcRun: solverRunStats },
          finalGrid: grid 
      },
    });
    return;
}

  if (maxTriesInWorker-- <= 0 && solverPhase !== 'propagate') {
    self.postMessage({
      type: 'REQUEST_RESTART_FROM_SOLVER',
      payload: {
        message: 'Max tries reached in worker -> Requesting restart.',
        wfcStatus: 'Failed (Timeout in Worker)',
        currentRunStats: { wfcRun: solverRunStats },
      },
    });
    solverPhase = 'stop';
    return;
  }

  if (solverPhase === 'collapse') {
    const cellToCollapse = pickCellWithLowestEntropy();
    if (!cellToCollapse) {
      if (!allCollapsed()) {
        // Should be a contradiction if no cell to collapse but not all done
        currentAttemptFailed = true; // Trigger backtrack
        solverRunStats.totalContradictions++;
        self.postMessage({
          type: 'STATUS_UPDATE',
          payload: {
            message:
              'Error: No cell to collapse, but not all collapsed. Contradiction.',
          },
        });
      } else {
        // Should have been caught by allCollapsed() earlier
        solverPhase = 'stop';
        self.postMessage({
          type: 'WFC_COMPLETE',
          payload: { finalStats: { wfcRun: solverRunStats } },
        });
        return;
      }
    } else {
      collapseCell(cellToCollapse[0], cellToCollapse[1]);
      if (!currentAttemptFailed) {
        solverPhase = 'propagate';
      }
    }
  } else if (solverPhase === 'propagate') {
    const MAX_PROP_PER_BATCH = Math.max(30, Math.floor((pCols * pRows) / 20)); // Batch propagations
    let processedInBatch = 0;
    let propResult = { success: true };

    while (
      processedInBatch++ < MAX_PROP_PER_BATCH &&
      queue.length > 0 &&
      !currentAttemptFailed
    ) {
      propResult = propagateStep();
      if (!propResult.success) break; // currentAttemptFailed would be true
    }

    if (queue.length === 0 || currentAttemptFailed || !propResult.success) {
      solverPhase = 'collapse'; // Go back to collapse if queue empty or contradiction
    }
  }

  // Post periodic stat updates
  self.postMessage({
    type: 'STATS_UPDATE',
    payload: { wfcRun: solverRunStats },
  });

  if (solverPhase !== 'stop') {
    runLoopTimeoutId = setTimeout(runStepInWorker, 0); // Continue the loop non-blockingly
  }
}

self.onmessage = function (e) {
  const {
    type,
    pCols: pc,
    pRows: pr,
    patterns: p,
    patternIdToIndex: pIdToIdx,
    adjacency: adj,
    currentPatternSize: cPS,
    maxBacktrackDepth,
    useRandomSelection,
  } = e.data;

  if (type === 'START_SOLVE') {
    pCols = pc;
    pRows = pr;
    // patternsW are received with ArrayBuffers for cells. These are now owned by this worker.
    // We need to reconstruct them into Uint8Arrays if they are not already.
    // Assuming main thread sent them as objects with .cells being ArrayBuffer.
    patternsW = p.map((pat) => ({ ...pat, cells: new Uint8Array(pat.cells) }));
    patternIdToIndexW = pIdToIdx;
    adjacencyW = adj;
    currentPatternSizeW_solver = cPS;
    MAX_BACKTRACK_DEPTH_W = maxBacktrackDepth;
    useRandomSelectionW = useRandomSelection;

    solverPhase = 'collapse'; // Initial phase

    try {
      initializeGridInWorker();
      self.postMessage({
        type: 'STATUS_UPDATE',
        payload: {
          message: `Solver worker started. Grid: ${pCols}x${pRows}. Patterns: ${patternsW.length}`,
        },
      });
      runStepInWorker(); // Start the loop
    } catch (err) {
      console.error('Error in solver worker START_SOLVE:', err);
      self.postMessage({
        type: null,
        error: err.message + (err.stack ? '\n' + err.stack : ''),
      });
      solverPhase = 'stop';
    }
  }
};

// Catch unhandled errors in the worker
self.onerror = function (message, source, lineno, colno, error) {
  console.error(
    'Unhandled error in Solver Worker:',
    message,
    source,
    lineno,
    colno,
    error
  );
  self.postMessage({
    type: null,
    error: `Unhandled: ${message} ${error ? error.stack : ''}`,
  });
  solverPhase = 'stop'; // Stop processing on unhandled error
  return true; // Prevents default error handling
};