import os
import io
import time
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from PIL import Image

SCREENSHOT_DIR = os.path.join('static', 'screenshots')
os.makedirs(SCREENSHOT_DIR, exist_ok=True)

FORMAT_EXT = {
    'png': 'png',
    'jpg': 'jpg',
    'jpeg': 'jpg',
    'pdf': 'pdf',
    'tiff': 'tiff',
    'tif': 'tiff',
}

MAX_FULLPAGE_HEIGHT = 32000


def create_driver(width, height, zoom):
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--hide-scrollbars')
    chrome_options.add_argument(f'--window-size={width},{height}')

    zoom_factor = zoom / 100.0
    if zoom_factor != 1.0:
        chrome_options.add_argument(f'--force-device-scale-factor={zoom_factor}')

    driver = webdriver.Chrome(options=chrome_options)
    driver.set_window_size(width, height)
    return driver


def capture_screenshot_image(driver, url, width, height, fullsize, zoom):
    driver.get(url)

    WebDriverWait(driver, 15).until(
        lambda d: d.execute_script('return document.readyState') == 'complete'
    )
    time.sleep(1)

    if fullsize:
        total_height = driver.execute_script(
            'return Math.max('
            '  document.body.scrollHeight,'
            '  document.body.offsetHeight,'
            '  document.documentElement.clientHeight,'
            '  document.documentElement.scrollHeight,'
            '  document.documentElement.offsetHeight'
            ')'
        )
        total_height = max(total_height, 100)
        total_height = min(total_height, MAX_FULLPAGE_HEIGHT)
        driver.set_window_size(width, total_height)
        time.sleep(0.5)

        total_height = driver.execute_script(
            'return Math.max('
            '  document.body.scrollHeight,'
            '  document.body.offsetHeight,'
            '  document.documentElement.clientHeight,'
            '  document.documentElement.scrollHeight,'
            '  document.documentElement.offsetHeight'
            ')'
        )
        total_height = min(total_height, MAX_FULLPAGE_HEIGHT)
        driver.set_window_size(width, total_height)
        time.sleep(0.5)

    png_data = driver.get_screenshot_as_png()
    img = Image.open(io.BytesIO(png_data))
    return img


def apply_scale(img, scale_pct):
    if scale_pct == 100:
        return img
    factor = scale_pct / 100.0
    new_w = max(1, int(img.width * factor))
    new_h = max(1, int(img.height * factor))
    return img.resize((new_w, new_h), Image.LANCZOS)


def save_in_format(img, filepath, fmt):
    if fmt == 'png':
        img.save(filepath, 'PNG', optimize=True)

    elif fmt in ('jpg', 'jpeg'):
        if img.mode in ('RGBA', 'LA', 'P'):
            background = Image.new('RGB', img.size, (255, 255, 255))
            if img.mode == 'P':
                img = img.convert('RGBA')
            background.paste(img, mask=img.split()[-1] if 'A' in img.mode else None)
            img = background
        elif img.mode != 'RGB':
            img = img.convert('RGB')
        img.save(filepath, 'JPEG', quality=92, optimize=True)

    elif fmt in ('tiff', 'tif'):
        img.save(filepath, 'TIFF')

    elif fmt == 'pdf':
        pdf_img = img.convert('RGB') if img.mode != 'RGB' else img
        pdf_img.save(filepath, 'PDF', resolution=72.0)

    else:
        img.save(filepath, 'PNG', optimize=True)