let sizeOptions;
let shapeOptions;
let redSlider1, greenSlider1, blueSlider1;
let redInput1, greenInput1, blueInput1;
let redSlider2, greenSlider2, blueSlider2;
let redInput2, greenInput2, blueInput2;
let shapes = []; // List to store the shapes
let lastDrag = []; // Stack to keep track of drag actions for undo functionality
let previewSwatch1, previewSwatch2;
let margin, canvasSize, controlX, toolsBackground;
let innerMargin = 16; // Inner margin for tools within the white box
let currentShape = 'square'; // Default shape
let sectionSpacing = 32; // Vertical spacing between sections
let uploadedImage;
let traceMode = false;
let traceOpacity = 0.5;
let lastSnapX = -1;
let lastSnapY = -1;

function setup() {
  pixelDensity(8.0); // Higher resolution setup
  margin = 50; // Margin from the edges of the window
  canvasSize = 1024; // Fixed canvas size, changed from 768 to 1024
  let cnv = createCanvas(canvasSize, canvasSize);
  cnv.position(margin, margin); // Position canvas with margin
  colorMode(RGB, 255);

  positionElements();

  // Set initial slider values
  redSlider1.value(0);   redInput1.value('0');
  greenSlider1.value(0); greenInput1.value('0');
  blueSlider1.value(0);  blueInput1.value('0');
  redSlider2.value(180);   redInput2.value('180');
  greenSlider2.value(180); greenInput2.value('180');
  blueSlider2.value(180);  blueInput2.value('180');

  updatePreviewSwatches();
}

function positionElements() {
  // Remove any existing tool elements
  removeElements();

  controlX = canvasSize + 2 * margin; // X position for controls, matching left and top canvas margin
  let sectionSpacing = 25; // Reduced spacing between sections (approximately 90% of previous 28)

  // Create white background for tools
  toolsBackground = createDiv('');
  toolsBackground.position(controlX, margin);
  toolsBackground.size(400, canvasSize - 16); // Increased width to 400px to accommodate grid size on one line
  toolsBackground.style('background-color', 'white');
  toolsBackground.style('padding', '8px');
  toolsBackground.style('border-radius', '4px');
  toolsBackground.style('box-shadow', '0 0 8px rgba(0,0,0,0.1)');
  toolsBackground.style('border', '1px solid black');

  let currentY = margin + innerMargin;

  // Add italic text at the top
  let italicText = createP('tools can be changed mid drawing');
  italicText.position(controlX + innerMargin, currentY);
  italicText.style('margin', '0');
  italicText.style('font-style', 'italic');
  italicText.style('color', '#B3B3B3');
  italicText.style('font-family', 'Arial, sans-serif');
  italicText.style('font-size', '12px');

  currentY += 30;

  // Add text for shape selection
  let shapeText = createP('select shape:');
  shapeText.position(controlX + innerMargin, currentY);
  shapeText.style('margin', '0');
  shapeText.style('font-weight', 'bold');
  shapeText.style('font-family', 'Arial, sans-serif');
  shapeText.style('font-size', '14px');

  currentY += 24;

  // Create shape options
  shapeOptions = createDiv('');
  shapeOptions.position(controlX + innerMargin, currentY);
  shapeOptions.style('display', 'flex');
  shapeOptions.style('flex-wrap', 'wrap');
  shapeOptions.style('gap', '8px');
  shapeOptions.style('width', '368px'); // Increased width

  createShapeOption('square', '■');
  createShapeOption('hrect2', '▬');
  createShapeOption('vrect2', '▮');
  createShapeOption('hrect4', '▬');
  createShapeOption('vrect4', '▮');
  createShapeOption('circle', '●');

  currentY += 90 + sectionSpacing; // Added section spacing

  // Add text for grid size
  let gridSizeText = createP('grid size:');
  gridSizeText.position(controlX + innerMargin, currentY);
  gridSizeText.style('margin', '0');
  gridSizeText.style('font-weight', 'bold');
  gridSizeText.style('font-family', 'Arial, sans-serif');
  gridSizeText.style('font-size', '14px');

  currentY += 24;

  // Setup radio buttons for pixel sizes on one line
  sizeOptions = createRadio();
  sizeOptions.option('2', '2');
  sizeOptions.option('4', '4');
  sizeOptions.option('8', '8');
  sizeOptions.option('16', '16');
  sizeOptions.option('32', '32');
  sizeOptions.option('64', '64');
  sizeOptions.option('128', '128');
  sizeOptions.option('256', '256');
  sizeOptions.style('width', '368px'); // Increased width
  sizeOptions.style('display', 'flex');
  sizeOptions.style('justify-content', 'space-between');
  sizeOptions.position(controlX + innerMargin, currentY);
  sizeOptions.selected('16');
  sizeOptions.style('font-family', 'Arial, sans-serif');
  sizeOptions.style('font-size', '12px');

  // Style radio buttons
  let radioButtons = selectAll('input[type="radio"]');
  radioButtons.forEach((button, index) => {
    button.style('margin-right', '4px');
  });

  currentY += 40 + sectionSpacing; // Reduced height for single line + section spacing

  // Add text for Color 1
  let color1Text = createP('color 1:');
  color1Text.position(controlX + innerMargin, currentY);
  color1Text.style('margin', '0');
  color1Text.style('font-weight', 'bold');
  color1Text.style('font-family', 'Arial, sans-serif');
  color1Text.style('font-size', '14px');

  currentY += 24;

  // Create color sliders and inputs for the first color set
  [redSlider1, greenSlider1, blueSlider1, redInput1, greenInput1, blueInput1] = createColorControls(controlX + innerMargin, currentY, '1');

  // Create preview swatch for Color 1
  previewSwatch1 = createDiv('');
  previewSwatch1.position(controlX + innerMargin + 280, currentY); // Adjusted position
  previewSwatch1.style('width', '40px');
  previewSwatch1.style('height', '40px');
  previewSwatch1.style('border', '1px solid black');

  // Add black, white, and gray buttons for Color 1
  createColorPresetButtons(controlX + innerMargin + 328, currentY, '1'); // Adjusted position

  currentY += 110 + sectionSpacing; // Added section spacing

  // Add text for Color 2
  let color2Text = createP('color 2:');
  color2Text.position(controlX + innerMargin, currentY);
  color2Text.style('margin', '0');
  color2Text.style('font-weight', 'bold');
  color2Text.style('font-family', 'Arial, sans-serif');
  color2Text.style('font-size', '14px');

  currentY += 24;

  // Create color sliders and inputs for the second color set
  [redSlider2, greenSlider2, blueSlider2, redInput2, greenInput2, blueInput2] = createColorControls(controlX + innerMargin, currentY, '2');

  // Create preview swatch for Color 2
  previewSwatch2 = createDiv('');
  previewSwatch2.position(controlX + innerMargin + 280, currentY); // Adjusted position
  previewSwatch2.style('width', '40px');
  previewSwatch2.style('height', '40px');
  previewSwatch2.style('border', '1px solid black');

  // Add black, white, and gray buttons for Color 2
  createColorPresetButtons(controlX + innerMargin + 328, currentY, '2'); // Adjusted position

  currentY += 110 + sectionSpacing; // Added section spacing

  // Add text for Image upload
  let imageUploadText = createP('image upload:');
  imageUploadText.position(controlX + innerMargin, currentY);
  imageUploadText.style('margin', '0');
  imageUploadText.style('font-weight', 'bold');
  imageUploadText.style('font-family', 'Arial, sans-serif');
  imageUploadText.style('font-size', '14px');

  currentY += 24;

  // Add file input for image upload
  let fileInput = createFileInput(handleFile);
  fileInput.position(controlX + innerMargin, currentY);
  fileInput.style('font-family', 'Arial, sans-serif');
  fileInput.style('font-size', '12px');

  currentY += 40;

  // Add trace mode toggle and opacity slider on the same line
  let traceModeCheckbox = createCheckbox('trace mode', false);
  traceModeCheckbox.position(controlX + innerMargin, currentY);
  traceModeCheckbox.style('font-family', 'Arial, sans-serif');
  traceModeCheckbox.style('font-size', '12px');
  traceModeCheckbox.changed(() => {
    traceMode = traceModeCheckbox.checked();
  });

  let traceOpacitySlider = createSlider(0, 1, 0.5, 0.01);
  traceOpacitySlider.position(controlX + innerMargin + 120, currentY);
  traceOpacitySlider.style('width', '100px');
  traceOpacitySlider.input(() => {
    traceOpacity = traceOpacitySlider.value();
  });

  let traceOpacityLabel = createSpan('opacity');
  traceOpacityLabel.position(controlX + innerMargin + 230, currentY);
  traceOpacityLabel.style('font-family', 'Arial, sans-serif');
  traceOpacityLabel.style('font-size', '12px');

  currentY += 40 + sectionSpacing; // Added section spacing

  // Create undo button and download buttons on the same line
  let undoButton = createButton('undo last move');
  undoButton.position(controlX + innerMargin, currentY);
  undoButton.mousePressed(undoLastMove);
  undoButton.style('font-family', 'Arial, sans-serif');
  undoButton.style('font-size', '12px');

  let downloadButton = createButton('download');
  downloadButton.position(controlX + innerMargin + 120, currentY);
  downloadButton.mousePressed(downloadCanvas);
  downloadButton.style('font-family', 'Arial, sans-serif');
  downloadButton.style('font-size', '12px');
  
  // New SVG download button
  let downloadSVGButton = createButton('download svg');
  downloadSVGButton.position(controlX + innerMargin + 200, currentY);
  downloadSVGButton.mousePressed(downloadSVG);
  downloadSVGButton.style('font-family', 'Arial, sans-serif');
  downloadSVGButton.style('font-size', '12px');

  updateShapeButtons();
}

function createColorControls(x, startY, suffix) {
  let sliders = [];
  let inputs = [];

  // Setup sliders and inputs for RGB controls
  sliders[0] = createSlider(0, 255, suffix === '1' ? 0 : 180);
  sliders[0].position(x, startY);
  sliders[0].style('width', '200px');
  inputs[0] = createInput(sliders[0].value().toString()).position(x + 210, startY).size(40);

  sliders[1] = createSlider(0, 255, suffix === '1' ? 0 : 180);
  sliders[1].position(x, startY + 30);
  sliders[1].style('width', '200px');
  inputs[1] = createInput(sliders[1].value().toString()).position(x + 210, startY + 30).size(40);

  sliders[2] = createSlider(0, 255, suffix === '1' ? 0 : 180);
  sliders[2].position(x, startY + 60);
  sliders[2].style('width', '200px');
  inputs[2] = createInput(sliders[2].value().toString()).position(x + 210, startY + 60).size(40);

  // Setup event listeners to synchronize sliders and inputs
  sliders.forEach((slider, index) => {
    slider.input(() => {
      inputs[index].value(slider.value().toString());
      updatePreviewSwatches();
    });
    inputs[index].input(() => {
      slider.value(int(inputs[index].value()));
      updatePreviewSwatches();
    });
  });

  // Set font for inputs
  inputs.forEach(input => {
    input.style('font-family', 'Arial, sans-serif');
  });

  return [sliders[0], sliders[1], sliders[2], inputs[0], inputs[1], inputs[2]];
}

function createColorPresetButtons(x, y, suffix) {
  let blackButton = createButton('B');
  blackButton.position(x, y);
  blackButton.mousePressed(() => setColorPreset(0, 0, 0, suffix));
  blackButton.style('width', '20px');
  blackButton.style('height', '20px');
  blackButton.style('padding', '0');
  blackButton.style('margin-right', '5px');

  let whiteButton = createButton('W');
  whiteButton.position(x + 25, y);
  whiteButton.mousePressed(() => setColorPreset(255, 255, 255, suffix));
  whiteButton.style('width', '20px');
  whiteButton.style('height', '20px');
  whiteButton.style('padding', '0');

  let grayButton = createButton('G');
  grayButton.position(x, y + 25);
  grayButton.mousePressed(() => setColorPreset(180, 180, 180, suffix));
  grayButton.style('width', '20px');
  grayButton.style('height', '20px');
  grayButton.style('padding', '0');
}

function setColorPreset(r, g, b, suffix) {
  if (suffix === '1') {
    redSlider1.value(r);   redInput1.value(r);
    greenSlider1.value(g); greenInput1.value(g);
    blueSlider1.value(b);  blueInput1.value(b);
  } else {
    redSlider2.value(r);   redInput2.value(r);
    greenSlider2.value(g); greenInput2.value(g);
    blueSlider2.value(b);  blueInput2.value(b);
  }
  updatePreviewSwatches();
}

function updatePreviewSwatches() {
  let color1 = `rgb(${redInput1.value()}, ${greenInput1.value()}, ${blueInput1.value()})`;
  let color2 = `rgb(${redInput2.value()}, ${greenInput2.value()}, ${blueInput2.value()})`;
  previewSwatch1.style('background-color', color1);
  previewSwatch2.style('background-color', color2);
}

function createShapeOption(shape, label) {
  let option = createButton('');
  option.parent(shapeOptions);
  option.attribute('data-shape', shape);
  option.style('width', '50px');
  option.style('height', '50px');
  option.style('display', 'flex');
  option.style('justify-content', 'center');
  option.style('align-items', 'center');
  option.style('background-color', 'white');
  option.style('border', '1px solid black');
  option.style('cursor', 'pointer');
  
  let shapeDiv = createDiv(label);
  shapeDiv.parent(option);
  shapeDiv.style('font-size', '24px');
  shapeDiv.style('display', 'flex');
  shapeDiv.style('justify-content', 'center');
  shapeDiv.style('align-items', 'center');
  
  if (shape === 'hrect2') {
    shapeDiv.style('width', '30px');
    shapeDiv.style('height', '15px');
    shapeDiv.style('background-color', 'black');
  } else if (shape === 'vrect2') {
    shapeDiv.style('width', '15px');
    shapeDiv.style('height', '30px');
    shapeDiv.style('background-color', 'black');
  } else if (shape === 'hrect4') {
    shapeDiv.style('width', '40px');
    shapeDiv.style('height', '10px');
    shapeDiv.style('background-color', 'black');
  } else if (shape === 'vrect4') {
    shapeDiv.style('width', '10px');
    shapeDiv.style('height', '40px');
    shapeDiv.style('background-color', 'black');
  } else {
    shapeDiv.html(label);
  }
  
  option.mousePressed(() => {
    currentShape = shape;
    updateShapeButtons();
  });
}

function updateShapeButtons() {
  shapeOptions.elt.childNodes.forEach(button => {
    let buttonShape = button.getAttribute('data-shape');
    button.style.backgroundColor = buttonShape === currentShape ? '#ddd' : 'white';
  });
}

function draw() {
  background(255); // Clear the canvas
  noStroke();

  // If we're in trace mode and have an uploaded image, draw it
  if (traceMode && uploadedImage) {
    push();
    tint(255, traceOpacity * 255); // Apply opacity
    image(uploadedImage, 0, 0, width, height);
    pop();
  }

  // Draw all shapes in the list
  shapes.forEach(shape => {
    fill(shape.color);
    drawShape(shape);
  });
  
  // Add a thin stroke around the canvas
  stroke(0);
  strokeWeight(1);
  noFill();
  rect(0, 0, width, height);
}

function drawShape(shape) {
  switch(shape.shape) {
    case 'square':
      rect(shape.x, shape.y, shape.size, shape.size);
      break;
    case 'hrect2':
      rect(shape.x, shape.y, shape.size * 2, shape.size);
      break;
    case 'vrect2':
      rect(shape.x, shape.y, shape.size, shape.size * 2);
      break;
    case 'hrect4':
      rect(shape.x, shape.y, shape.size * 4, shape.size);
      break;
    case 'vrect4':
      rect(shape.x, shape.y, shape.size, shape.size * 4);
      break;
    case 'circle':
      ellipse(shape.x + shape.size/2, shape.y + shape.size/2, shape.size, shape.size);
      break;
  }
}

function mouseDragged() {
  if (mouseX < 0 || mouseX > width || mouseY < 0 || mouseY > height) return;

  let currentSize = int(sizeOptions.value());
  let color1 = color(redInput1.value(), greenInput1.value(), blueInput1.value());
  let color2 = color(redInput2.value(), greenInput2.value(), blueInput2.value());

  let snapX, snapY, checkerIndex;

  switch(currentShape) {
    case 'square':
    case 'circle':
      snapX = floor(mouseX / currentSize) * currentSize;
      snapY = floor(mouseY / currentSize) * currentSize;
      checkerIndex = (floor(snapX / currentSize) + floor(snapY / currentSize)) % 2;
      break;
    case 'hrect2':
      snapX = floor(mouseX / (currentSize * 2)) * (currentSize * 2);
      snapY = floor(mouseY / currentSize) * currentSize;
      checkerIndex = (floor(snapX / (currentSize * 2)) + floor(snapY / currentSize)) % 2;
      break;
    case 'vrect2':
      snapX = floor(mouseX / currentSize) * currentSize;
      snapY = floor(mouseY / (currentSize * 2)) * (currentSize * 2);
      checkerIndex = (floor(snapX / currentSize) + floor(snapY / (currentSize * 2))) % 2;
      break;
    case 'hrect4':
      snapX = floor(mouseX / (currentSize * 4)) * (currentSize * 4);
      snapY = floor(mouseY / currentSize) * currentSize;
      checkerIndex = (floor(snapX / (currentSize * 4)) + floor(snapY / currentSize)) % 2;
      break;
    case 'vrect4':
      snapX = floor(mouseX / currentSize) * currentSize;
      snapY = floor(mouseY / (currentSize * 4)) * (currentSize * 4);
      checkerIndex = (floor(snapX / currentSize) + floor(snapY / (currentSize * 4))) % 2;
      break;
  }

  // Only add a shape if we've moved to a new grid position
  if (snapX !== lastSnapX || snapY !== lastSnapY) {
    let chosenColor = checkerIndex === 0 ? color1 : color2;

    shapes.push({
      shape: currentShape,
      color: chosenColor,
      size: currentSize,
      x: snapX,
      y: snapY
    });

    lastDrag.push(shapes.length - 1);
    lastSnapX = snapX;
    lastSnapY = snapY;
  }
}

function mouseReleased() {
  if (lastDrag.length > 0) {
    let sessionShapes = lastDrag.slice();
    lastDrag = [];
    lastDrag.push(sessionShapes);
  }
  // Reset position tracking for next drag
  lastSnapX = -1;
  lastSnapY = -1;
}

function undoLastMove() {
  if (lastDrag.length > 0) {
    let lastSession = lastDrag.pop();
    while (lastSession.length) {
      let index = lastSession.pop();
      shapes.splice(index, 1);
    }
  }
}

function keyPressed() {
  if (key === 'd' || key === 'D') {
    downloadCanvas();
  }
}

function downloadCanvas() {
  let tempCanvas = createGraphics(width, height);
  tempCanvas.background(255);
  tempCanvas.noStroke();

  shapes.forEach(shape => {
    tempCanvas.fill(shape.color);
    drawShapeOnCanvas(tempCanvas, shape);
  });

  saveCanvas(tempCanvas, 'drawing', 'png');
}

function colorToHex(c) {
  let r = hex(c.levels[0], 2);
  let g = hex(c.levels[1], 2);
  let b = hex(c.levels[2], 2);
  return `#${r}${g}${b}`;
}

function downloadSVG() {
  // Create SVG string with proper namespace
  let svgString = `<svg width="${width}" height="${height}" viewBox="0 0 ${width} ${height}" xmlns="http://www.w3.org/2000/svg">\n`;
  
  // Add white background
  svgString += `  <rect x="0" y="0" width="${width}" height="${height}" fill="white"/>\n`;
  
  // Deduplicate shapes - keep only the last (topmost) shape at each position
  let shapeMap = new Map();
  shapes.forEach(shape => {
    let key = `${shape.x},${shape.y},${shape.shape},${shape.size}`;
    shapeMap.set(key, shape);
  });
  
  // Export only the unique topmost shapes
  shapeMap.forEach(shape => {
    let hexColor = colorToHex(shape.color);
    
    let x = shape.x;
    let y = shape.y;
    let s = shape.size;
    
    switch(shape.shape) {
      case 'square':
        svgString += `  <rect x="${x}" y="${y}" width="${s}" height="${s}" fill="${hexColor}" stroke="none"/>\n`;
        break;
      case 'hrect2':
        svgString += `  <rect x="${x}" y="${y}" width="${s * 2}" height="${s}" fill="${hexColor}" stroke="none"/>\n`;
        break;
      case 'vrect2':
        svgString += `  <rect x="${x}" y="${y}" width="${s}" height="${s * 2}" fill="${hexColor}" stroke="none"/>\n`;
        break;
      case 'hrect4':
        svgString += `  <rect x="${x}" y="${y}" width="${s * 4}" height="${s}" fill="${hexColor}" stroke="none"/>\n`;
        break;
      case 'vrect4':
        svgString += `  <rect x="${x}" y="${y}" width="${s}" height="${s * 4}" fill="${hexColor}" stroke="none"/>\n`;
        break;
      case 'circle':
        let cx = x + s / 2;
        let cy = y + s / 2;
        let r = s / 2;
        svgString += `  <circle cx="${cx}" cy="${cy}" r="${r}" fill="${hexColor}" stroke="none"/>\n`;
        break;
    }
  });

  svgString += `</svg>`;

  // Create a blob and download it
  let blob = new Blob([svgString], { type: 'image/svg+xml' });
  let url = URL.createObjectURL(blob);
  let link = document.createElement('a');
  link.href = url;
  link.download = 'drawing.svg';
  document.body.appendChild(link);
  link.click();
  document.body.removeChild(link);
  URL.revokeObjectURL(url);
}

function drawShapeOnCanvas(canvas, shape) {
  switch(shape.shape) {
    case 'square':
      canvas.rect(shape.x, shape.y, shape.size, shape.size);
      break;
    case 'hrect2':
      canvas.rect(shape.x, shape.y, shape.size * 2, shape.size);
      break;
    case 'vrect2':
      canvas.rect(shape.x, shape.y, shape.size, shape.size * 2);
      break;
    case 'hrect4':
      canvas.rect(shape.x, shape.y, shape.size * 4, shape.size);
      break;
    case 'vrect4':
      canvas.rect(shape.x, shape.y, shape.size, shape.size * 4);
      break;
    case 'circle':
      canvas.ellipse(shape.x + shape.size/2, shape.y + shape.size/2, shape.size, shape.size);
      break;
  }
}

function handleFile(file) {
  if (file.type === 'image') {
    loadImage(file.data, img => {
      uploadedImage = createGraphics(width, height);
      let scale = Math.min(width / img.width, height / img.height);
      let newWidth = img.width * scale;
      let newHeight = img.height * scale;
      uploadedImage.image(img, (width - newWidth) / 2, (height - newHeight) / 2, newWidth, newHeight);
    });
  } else {
    console.log('Not an image file!');
  }
}

function windowResized() {
  canvasSize = calculateCanvasSize();
  resizeCanvas(canvasSize, canvasSize);
  let cnv = select('canvas');
  cnv.position(margin, margin);
  positionElements();
  updatePreviewSwatches();
}

function calculateCanvasSize() {
  return min(windowWidth - 500, windowHeight - 2 * margin, 1024); // Max size changed from 768 to 1024
}