import asyncio
from pyppeteer import launch
from pyppeteer.errors import ElementHandleError

async def highlight_element(page, attribute):
    try:
        element = await page.querySelector(attribute)
        if element:
            await page.evaluate('element => element.style.backgroundColor = "yellow"', element)
            styles = await page.evaluate('(element) => { const computedStyles = window.getComputedStyle(element); return Array.from(computedStyles).map(prop => `${prop}: ${computedStyles.getPropertyValue(prop)}`); }', element)
            return styles
    except ElementHandleError:
        pass
    return None

async def getStyles(page, attribute, text_file):
    try:
        await page.waitForSelector(attribute)

        container_styles = await page.evaluate('''(attribute) => {
            const container = document.querySelector(attribute);
            const computedStyles = getComputedStyle(container);
            const styles = {};
            for (let i = 0; i < computedStyles.length; i++) {
                const prop = computedStyles[i];
                styles[prop] = computedStyles.getPropertyValue(prop);
            }
            return styles;
        }''', attribute)

        styles = await page.evaluate('''(attribute) => {
            const container = document.querySelector(attribute);
            const elements = container.querySelectorAll("*");

            const extractStyles = (element, tagName, index, totalIndex) => {
                const computedStyles = getComputedStyle(element);
                const styles = {};
                for (let i = 0; i < computedStyles.length; i++) {
                    const prop = computedStyles[i];
                    styles[prop] = computedStyles.getPropertyValue(prop);
                }
                return styles;
            };

            const result = {};
            elements.forEach((element, index) => {
                const tagName = element.tagName.toLowerCase();
                const tagIndex = [...container.querySelectorAll(tagName)].indexOf(element);
                const totalIndex = index;
                const variableName = `${tagName}-${tagIndex}-${totalIndex}`;
                result[variableName] = extractStyles(element, tagName, tagIndex, totalIndex);
            });

            return result;
        }''', attribute)

        with open(text_file, 'w') as f:
            f.write(f'Styles of container {attribute}:\n')
            for prop, value in container_styles.items():
                f.write(f'{prop}: {value}\n')
            f.write('\n')

            for variable_name, element_styles in styles.items():
                f.write(f'Here goes the styles of {variable_name}:\n')
                for prop, value in element_styles.items():
                    f.write(f'{prop}: {value}\n')
                f.write('\n')

    except ElementHandleError:
        pass
    return None

async def highlightAndStyles(url, attribute):
    browser = await launch()
    page = await browser.newPage()
    await page.goto(url)
    
    text_file = "styles.txt"
    image_file = "screenshot.png"
    await getStyles(page, attribute, text_file)
    styles = await highlight_element(page, attribute)
    
    await page.screenshot({'path': image_file, 'fullPage': True})
    await browser.close()

url = input("Enter the URL of the web page: ")
attribute = input("Enter the class or ID attribute to highlight (e.g., .container or #element): ")

asyncio.get_event_loop().run_until_complete(highlightAndStyles(url, attribute))
print("Screenshot and styles saved successfully!")