document.addEventListener('DOMContentLoaded', async () => {
  // Get screenshot from storage
  const result = await chrome.storage.local.get(['tempScreenshot']);
  if (!result.tempScreenshot) {
    console.error('No screenshot found in storage');
    return;
  }

  // Create image element
  const img = new Image();
  img.onload = () => {
    // Set canvas size to match image
    const canvas = document.getElementById('canvas');
    const ctx = canvas.getContext('2d');
    canvas.width = img.width;
    canvas.height = img.height;

    // Draw image
    ctx.drawImage(img, 0, 0);

    // Drawing state
    let isDrawing = false;
    let currentTool = 'brush';
    let startX = 0;
    let startY = 0;
    let lastText = null;

    // Tool button setup
    document.querySelectorAll('.tool-btn').forEach(btn => {
      btn.addEventListener('click', () => {
        // Remove active class from all buttons
        document.querySelectorAll('.tool-btn').forEach(b => b.classList.remove('active'));
        // Add active class to clicked button
        btn.classList.add('active');
        currentTool = btn.dataset.tool;
      });
    });

    // Mouse event handlers
    canvas.addEventListener('mousedown', startDrawing);
    canvas.addEventListener('mousemove', draw);
    canvas.addEventListener('mouseup', stopDrawing);
    canvas.addEventListener('mouseout', stopDrawing);

    function startDrawing(e) {
      isDrawing = true;
      const rect = canvas.getBoundingClientRect();
      startX = e.clientX - rect.left;
      startY = e.clientY - rect.top;

      if (currentTool === 'text') {
        const text = prompt('Enter text:', '');
        if (text) {
          ctx.fillStyle = document.getElementById('color-picker').value;
          ctx.font = '20px Arial';
          ctx.fillText(text, startX, startY);
        }
      }
    }

    function draw(e) {
      if (!isDrawing) return;

      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;

      ctx.strokeStyle = document.getElementById('color-picker').value;
      ctx.lineWidth = document.getElementById('stroke-width').value;
      ctx.lineCap = 'round';

      switch (currentTool) {
        case 'brush':
          ctx.beginPath();
          ctx.moveTo(startX, startY);
          ctx.lineTo(x, y);
          ctx.stroke();
          startX = x;
          startY = y;
          break;

        case 'rectangle':
          // Clear previous preview
          ctx.putImageData(lastImageData, 0, 0);
          // Draw new rectangle
          ctx.beginPath();
          ctx.rect(startX, startY, x - startX, y - startY);
          ctx.stroke();
          break;

        case 'arrow':
          // Clear previous preview
          ctx.putImageData(lastImageData, 0, 0);
          // Draw arrow line
          ctx.beginPath();
          ctx.moveTo(startX, startY);
          ctx.lineTo(x, y);
          ctx.stroke();
          // Draw arrow head
          const angle = Math.atan2(y - startY, x - startX);
          const headLength = 20;
          ctx.beginPath();
          ctx.moveTo(x, y);
          ctx.lineTo(x - headLength * Math.cos(angle - Math.PI / 6), y - headLength * Math.sin(angle - Math.PI / 6));
          ctx.moveTo(x, y);
          ctx.lineTo(x - headLength * Math.cos(angle + Math.PI / 6), y - headLength * Math.sin(angle + Math.PI / 6));
          ctx.stroke();
          break;
      }
    }

    function stopDrawing() {
      isDrawing = false;
      // Save the current canvas state
      lastImageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
    }

    // Save initial canvas state
    let lastImageData = ctx.getImageData(0, 0, canvas.width, canvas.height);

    // Set up done button
    document.getElementById('done-btn').addEventListener('click', () => {
      // Convert canvas to data URL
      const dataUrl = canvas.toDataURL();
      
      // Save to storage
      chrome.storage.local.set({ screenshot: dataUrl }, () => {
        window.close();
      });
    });

    // Set up cancel button
    document.getElementById('cancel-btn').addEventListener('click', () => {
      window.close();
    });

    // Select brush tool by default
    document.querySelector('[data-tool="brush"]').click();
  };

  // Load image
  img.src = result.tempScreenshot;
});