import os
import time
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.chrome.service import Service
from PIL import Image

# Hämta sökvägen till katalogen där skriptet körs
ROOT_DIRECTORY = os.getcwd()

def take_screenshot(url, output_path, width, height, delay, suffix):
    """Tar en skärmdump av en webbsida efter en given fördröjning."""

    options = Options()
    options.add_argument('--headless=new')

    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)
    driver.get(url)

    # Vänta tills sidan har laddats (justera vid behov)
    time.sleep(delay)

    # Ställ in viewport-storlek med hänsyn till Retina-skärm
    driver.set_window_size(width * 2, height * 2)  # Multiplicera med 2

    # Ta skärmdump av hela sidan
    temp_screenshot_path = os.path.join(os.path.dirname(output_path), f"{os.path.basename(output_path)}_{suffix}.temp")
    driver.save_screenshot(temp_screenshot_path)
    driver.quit()

    # Beskär och skala ner skärmdumpen
    try:
        im = Image.open(temp_screenshot_path)
        im = im.crop((0, 0, width * 2, height * 2))
        im = im.resize((width, height), Image.LANCZOS)  # Skala ner till hälften
        im.save(output_path + f'_{suffix}.jpg', 'JPEG')  # Spara som JPG
        os.remove(temp_screenshot_path)
    except Exception as e:
        print(f"Error processing screenshot: {e}")

def main():
    """Går igenom katalogstrukturen och tar skärmdumpar."""

    for root, dirs, _ in os.walk(ROOT_DIRECTORY):
        for dir_name in dirs:
            if 'x' in dir_name and '_' in dir_name:
                dir_path = os.path.join(root, dir_name)
                if os.path.isdir(dir_path):
                    # Extrahera bredd och höjd från katalognamnet
                    width, height = map(int, dir_name.split('_')[0].split('x'))

                    # Skapa sökväg till index.html
                    index_html_path = os.path.join(dir_path, 'index.html')
                    if not os.path.exists(index_html_path):
                        continue

                    # Skapa sökväg för skärmdumpen i "fallbacks"-katalogen
                    fallbacks_dir = os.path.join(root, 'fallbacks')
                    os.makedirs(fallbacks_dir, exist_ok=True)

                    # Extrahera delarna av sökvägen för att skapa filnamnet
                    parts = root.split(os.sep)[-2:]  # De två sista delarna av sökvägen
                    base_file_name = f"{dir_name.split('_')[0]}_{parts[0]}_{parts[1]}_{'_'.join(dir_name.split('_')[1:])}"

                    # Skapa URL:en till index.html
                    url = f'file://{index_html_path}'

                    # Ta första skärmdumpen efter 2 sekunder
                    take_screenshot(url, os.path.join(fallbacks_dir, base_file_name), width, height, 2, '1')

                    # Ta andra skärmdumpen efter 6 sekunder
                    take_screenshot(url, os.path.join(fallbacks_dir, base_file_name), width, height, 6, '2')

                    print(f'Skärmdumpar sparade: {base_file_name}_1.jpg och {base_file_name}_2.jpg')

if __name__ == '__main__':
    main()