import pyautogui
from PIL import Image, ImageDraw, ImageFont

import pyautogui
from PIL import Image, ImageDraw, ImageFont

def capture_screenshot(grid_size=(26,26)):
    screenshot = pyautogui.screenshot()
    width, height = screenshot.size
    
    # Create a new image with extra white space
    new_width = width + 140  # 100 pixels for left margin
    new_height = height + 95  # 95 pixels for bottom margin, 20 for top
    new_image = Image.new('RGB', (new_width, new_height), color='white')
    new_image.paste(screenshot, (100, 20))  # 20 pixels from top, 100 from left
    
    draw = ImageDraw.Draw(new_image)
    cell_width = width // grid_size[0]
    cell_height = height // grid_size[1]

    # Draw black border (adjusted to fit the screenshot exactly)
    draw.rectangle([98, 18, width+101, height+21], outline='black', width=2)

    # Draw grid lines
    for i in range(grid_size[0] + 1):
        x = i * cell_width + 100
        draw.line([(x, 20), (x, height + 20)], fill='red', width=1)
    for i in range(grid_size[1] + 1):
        y = i * cell_height + 20
        draw.line([(100, y), (width + 100, y)], fill='red', width=1)

    # Label grid cells
    font = ImageFont.truetype("arialbd.ttf", 52)
    small_font = ImageFont.truetype("arialbd.ttf", 24)  # Increased to 24px
    for col in range(grid_size[0]):
        label = chr(65 + col)
        x = col * cell_width + 100  # Align with grid lines
        draw.text((x, new_height - 45), label, fill='red', font=font, anchor='mt')
    for row in range(grid_size[1]):
        label = str(row + 1)
        y = row * cell_height + 20
        draw.text((60, y), label, fill='red', font=font, anchor='rm')
    
    # Add '26' label
    draw.text((60, height + 20), '26', fill='red', font=font, anchor='rm')

    # Add small coordinates at every grid intersection point
    for col in range(grid_size[0] + 1):
        for row in range(grid_size[1] + 1):
            x = col * cell_width + 100
            y = row * cell_height + 20
            coord = f"{chr(65 + col)}{row + 1}"
            
            # Calculate text size
            text_bbox = draw.textbbox((x, y), coord, font=small_font, anchor='mm')
            text_width = text_bbox[2] - text_bbox[0]
            text_height = text_bbox[3] - text_bbox[1]
            
            # Draw pink background
            padding = 2
            draw.rectangle([x - text_width//2 - padding, y - text_height//2 - padding,
                            x + text_width//2 + padding, y + text_height//2 + padding],
                           fill='pink')
            
            # Draw coordinate text
            draw.text((x, y), coord, fill='black', font=small_font, anchor='mm')

    screenshot_path = 'grid_screenshot.png'
    new_image.save(screenshot_path)
    return screenshot_path