# import asyncio
from pyppeteer import launch
from pyppeteer.errors import ElementHandleError
import gradio as gr

# Firefall start ---------------------------------------------------------------------------------------------

import os, json, requests

# enter the client id from imss here
client_id = 'MDSR_Firefall'
client_secret = 's8e-8CGebu-kO3Vt_ICCNzQU8sCVYCHqcuFq' #enter the client secret from imss here
permanent_auth_code = 'eyJhbGciOiJSUzI1NiIsIng1dSI6Imltc19uYTEtc3RnMS1rZXktcGFjLTEuY2VyIiwia2lkIjoiaW1zX25hMS1zdGcxLWtleS1wYWMtMSIsIml0dCI6InBhYyJ9.eyJpZCI6Ik1EU1JfRmlyZWZhbGxfc3RnIiwidHlwZSI6ImF1dGhvcml6YXRpb25fY29kZSIsImNsaWVudF9pZCI6Ik1EU1JfRmlyZWZhbGwiLCJ1c2VyX2lkIjoiTURTUl9GaXJlZmFsbEBBZG9iZUlEIiwiYXMiOiJpbXMtbmExLXN0ZzEiLCJvdG8iOmZhbHNlLCJjcmVhdGVkX2F0IjoiMTY4MTE0NTIxNDk1MCIsInNjb3BlIjoic3lzdGVtIn0.Yoz7IPhmIBV2uNKl1CJJ9rJ0HmvDBQFbh0AihlHdsOa1E3yBs7WB9ilTCUVodifg8gh1yw4QRllV1NKS2RYeiGxQU7rXAF7SEnH_X_Tqdl735PBnBFL8sW_x76dzmT6MZIzynz8Ywu57qztvFnHoLMfJ7HsNt7rkOqF3IZByOinxyJzRTwMfygHSKjoQx6A4S7LbuQWjlqDbM9RaeCcakMEqGvSKqkLQvtMg40ZQYSNELoFtbATfwuVrHWOglAQS4A2FR24ziop137imu4HrTr-syDYki8VWV27WuGGo632_K2vJwqbaYjZvyrtsuBLH3fGGgXgyM5EA_Jk_lcMFog' #imss -> service tokens -> permanent auth token
ims_url = 'https://ims-na1-stg1.adobelogin.com/ims/token/v2'
firefall_client_id = "MDSR_Firefall"
ims_org_id = client_id
api_key = client_id
azure_url = 'https://firefall-stage.adobe.io/v1/completions'

def get_openai_response(azure_url, ims_org_id, api_key, temp_auth_token, json_data):
    headers = {
        'x-gw-ims-org-id': ims_org_id,
        'x-api-key': api_key,
        'Authorization': f'Bearer {temp_auth_token}',
        'Content-Type': 'application/json',
    }
    response = requests.post(azure_url, headers=headers, json=json_data)
    return json.loads(response.text)

def get_temp_auth_token(ims_url, client_id, client_secret, permanent_auth_code):
    params = {
        'client_id': client_id,
        'client_secret': client_secret,
        'code': permanent_auth_code,
        'grant_type': 'authorization_code',
    }
    response = requests.post(ims_url, params=params)
    return json.loads(response.text)


response = get_temp_auth_token(ims_url, client_id, client_secret, permanent_auth_code)
# print(response)
temp_auth_token = response['access_token']

async def gptResponse(franklinHTML, originalStyles, css_file):
    query = f'I will provide you :-\n1. HTML code,\n2. \"Computed\" Styles (obtained from \"Computed\" tab) of all the elements inside this HTML.\n\nBased on these two information, can you generate \"CSS\" code for the given HTML so that it styles the HTML according to the provided \"Computed\" styles.\n\n'
    query += f'1. HTML code:-\n\n'
    query += f'```{franklinHTML}```\n\n'
    query += f'2. \"Computed\" Styles (obtained from \"Computed\" tab) of all the elements inside this HTML:-\n\n'
    query += f'{originalStyles}\n\n'
    query += f'Generate \"CSS\" code for the given HTML block so that it styles the HTML block according to the provided \"Computed\" styles. Don\'t included the styles inline, instead generate a separate file for it (external css). Only generate the \"CSS\" code and no explanation.'
    query += f'\n\nThe generated CSS code must be around 1000 tokens atleast'
    
    with open("query.txt", 'w') as f:
        f.write(query)

    json_data =  {
                    "dialogue":{
                        "question": query
                    },
                    "llm_metadata": {
                        "llm_type": "azure_chat_openai",
                        "model_name": "gpt-4-32k",
                        "temperature": 0.2,
                        # "model_name": "gpt-4",
                        # "temperature": 0.0,
                        # "max_tokens": 8071,
                        # "top_p": 1.0,
                        # "frequency_penalty": 0,
                        # "presence_penalty": 0,
                        # "n": 1,
                        # "llm_type": "azure_chat_openai"
                    }
                }

    openai_response = get_openai_response(azure_url, ims_org_id, api_key, temp_auth_token, json_data)
    # print("OpenAI Response:", openai_response)
    res = openai_response['generations'][0][0]['text']
    with open(css_file, 'w') as f:
        f.write(res)
    # return ("OpenAI Response: " + openai_response['generations'][0][0]['text'])

# Firefall end -----------------------------------------------------------------------------------------------

async def highlight_element(page, selector):
    try:
        element = await page.querySelector(selector)
        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, selector, text_file):
    try:
        await page.waitForSelector(selector)

        originalStyles = await page.evaluate('''(selector) => {
            const contentTypeIndex = {};
            const childs = {};
            let finalString = "";
            function elementIndentifier(element) {

                const tagName = element.tagName.toLowerCase();

                // Check for different element types
                if (tagName === 'p' || tagName[0] === 'h' || tagName === 'span' || tagName === 'div' || tagName === 'ul' || tagName === 'ol' || tagName === 'article') { // divs can be used for wrapping text as well
                    // Text-based elements
                    return `Text Content = "${element.textContent}"`;
                } else if (tagName === 'img' || tagName === 'audio' || tagName === 'video') {
                    // Image, audio or video elements
                    let splitUrl = element.src.split('/').slice(-1)[0];
                    if(element['data-src']){
                        splitUrl = element['data-src'].split('/').slice(-1)[0];
                    }

                    // Search for the file extension
                    const extensions = ['.jpg', '.png', '.webp', '.svg'];
                    let fileName = splitUrl;
                    extensions.forEach(extension => {
                        if (splitUrl.includes(extension)) {
                            if(extension === '.svg'){
                                fileName = splitUrl.split(extension)[0] + '.png';
                            }else{
                                fileName = splitUrl.split(extension)[0] + extension;
                            }
                        }
                    });

                    return `src = "${fileName}"`;
                } else if (tagName === 'a' || tagName === 'link') {
                    // Link elements
                    return `href = "${element.href}"`;
                } else if (tagName === 'input' || tagName === 'textarea' || tagName === 'select') {
                    // Form elements
                    return `value = "${element.value}"`;
                } else if (tagName === 'table') {
                    // Table-related elements
                    // Example to retrieve cell content
                    const cellContent = Array.from(element.querySelectorAll('td, th')).map(cell => cell.innerHTML);
                    return `rows = "${element.rows}", "cellContent = "${cellContent.join(', ')}"`;
                } else if (tagName === 'input' && (element.type === 'checkbox' || element.type === 'radio')) {
                    // Checkbox and Radio Button elements
                    return `checked = "${element.checked}"`;
                }

            }

            function contentType(element) {

                const tagName = element.tagName.toLowerCase();

                // Check for different element types
                if (tagName === 'p' || tagName[0] === 'h' || tagName === 'span' || tagName === 'div' || tagName === 'ul' || tagName === 'ol') {
                    return "text";
                } else if (tagName === 'img') {
                    return "image";
                } else if (tagName === 'audio') {
                    return "audio";
                } else if (tagName === 'video') {
                    return "video";
                } else if (tagName === 'a' || tagName === 'link') {
                    return "link";
                } else if (tagName === 'input' || tagName === 'textarea' || tagName === 'select') {
                    return "form";
                } else if (tagName === 'table') {
                    return "table";
                } else if (tagName === 'input' && (element.type === 'checkbox' || element.type === 'radio')) {
                    // Checkbox and Radio Button elements
                    return "button";
                }
            }

            function nameElems(element) {

                for (let i = element.children.length - 1; i >= 0; i--) {
                    nameElems(element.children[i]);
                }

                const s1 = element.children.length ? "container" : contentType(element);
                contentTypeIndex[s1] ? contentTypeIndex[s1]++ : contentTypeIndex[s1] = 1;
                const elemName = `${s1}${contentTypeIndex[s1]}`;
    
                element.name = elemName;

            }


            function postorderTraversal(element) {

                // Recursively traverse child nodes in postorder

                for (let i = element.children.length - 1; i >= 0; i--) {
                    postorderTraversal(element.children[i]);
                }

    
                // Process the current element

                const s1 = element.children.length ? "container" : contentType(element);

                if (element.children.length) {
                    for (let i = element.children.length - 1; i >= 0; i--) {
                        childs[element.name] ? childs[element.name].push(element.children[i].name) : childs[element.name] = [element.children[i].name];
                    }
                    finalString = finalString + `${childs[element.name]} ${childs[element.name] > 1 ? 'are' : 'is'} nested inside a container. Lets name this container as ${element.name}. The computed styles of ${element.name} are:-\n`;
                } else{
                    finalString += `Lets name ${s1} with ${elementIndentifier(element)} as ${element.name}. The computed styles of ${element.name} are:-\n`;
                }

                const computedStyles = getComputedStyle(element);
                desired_properties = [
                    "background-color", "box-sizing", "clear", "color", "display", 
                    "flex-direction", "float", "font-family", "font-size", "font-weight", 
                    "height", "line-height", "margin-bottom", "margin-left", "margin-right", 
                    "margin-top", "padding-bottom", "padding-left", "padding-right", 
                    "padding-top", "text-align",  "width"
                ];

                finalString += `\"\"`
                for (let i = 0; i < desired_properties.length; i++) {
                    const prop = desired_properties[i]; 
                    finalString += `\t${prop}: ${computedStyles.getPropertyValue(prop)}`;
                    if(prop !== "width"){
                        finalString += `\n`;
                    }
                }
                finalString += `\"\"\n\n`

            }

            const rootElement = document.querySelector(selector);
            nameElems(rootElement);
            postorderTraversal(rootElement);

            return finalString;
        }''', selector)


        with open(text_file, 'w') as f:
            f.write(originalStyles)

        return originalStyles
    
    except ElementHandleError:
        pass
    return None

# -------------------------------------------------------------------------------------------------------------

from bs4 import BeautifulSoup
import shutil
from urllib.parse import urlparse
# import urllib
# from urllib.request import Request, urlretrieve
# from urllib.error import HTTPError
from cairosvg import svg2png

def download_images(url, selector):

    response = requests.get(url, stream=True, headers={'User-agent': 'Mozilla/6.0'})
    soup = BeautifulSoup(response.text, 'html.parser')
    
    img_tags = []
    
    # can now and operate multiple class or ids but this assumes that selector will be for a div or block in originalDOM will be a div
    img_tags += soup.select_one(f'div{selector}').find_all('img') 
    
    # modifiedSelector = selector[1:]
    # if selector[0] == '.':
    #     img_tags += soup.find(class_ = modifiedSelector).find_all('img')
    # else:
    #     img_tags += soup.find(id = modifiedSelector).find_all('img')
    
    for img in img_tags:

        img_url = img['src']
        try:
            img_url = img['data-src']
        except KeyError:
            img_url = img['src']
        
        # Split URL by the last backslash
        split_url = img_url.rsplit('/', 1)
        file_name_with_query = split_url[1]

        # Search for the file extension
        extensions = ['.jpg', '.png', '.webp', '.svg']
        img_name = file_name_with_query
        for extension in extensions:
            if extension in file_name_with_query:
                img_name = file_name_with_query.split(extension)[0] + extension
                break
        
        try:
            r = requests.get(img_url, stream=True, headers={'User-agent': 'Mozilla/6.0'})
            if r.status_code == 200:
                with open(img_name, 'wb') as f:
                    r.raw.decode_content = True
                    shutil.copyfileobj(r.raw, f)
        except:
            parsed_url = urlparse(url)
            scheme = parsed_url.scheme
            netloc = parsed_url.netloc
            homePage_url = f"{scheme}://{netloc}"

            new_url = homePage_url + img['src']

            try:
                r = requests.get(new_url, stream=True, headers={'User-agent': 'Mozilla/6.0'})
                if r.status_code == 200:
                    with open(img_name, 'wb') as f:
                        r.raw.decode_content = True
                        shutil.copyfileobj(r.raw, f)
            except:
                newer_url = img['currentSrc']
                try:
                    r = requests.get(newer_url, stream=True, headers={'User-agent': 'Mozilla/6.0'})
                    if r.status_code == 200:
                        with open(img_name, 'wb') as f:
                            r.raw.decode_content = True
                            shutil.copyfileobj(r.raw, f)
                except:
                    print(Exception)

        if '.svg' in img_name:
            with open(img_name, 'r') as f:
                svg_code = f.read()
            img_name = img_name[:img_name.index('.svg')] + '.png'
            svg2png(bytestring=svg_code,write_to=img_name)
        
#         # try:
#         #     urlretrieve(img['src'], img_name)
#         #     # print(f"Downloaded {img_name}")
#         # except HTTPError:
#         #     urlretrieve(img_url, img_name)
#         # # except Exception as e:
#         # #     print(f"Failed to download {img_name}: {str(e)}")

# # # Usage example
# # download_images('https://www.elixirsolutions.com/', element_class='columncontainer')

# -------------------------------------------------------------------------------------------------------------

async def regeneratedBlock(html_file, css_file, rebuilt_file):
    with open(html_file, 'r') as file:
        content = file.read()
    
    new_content = f'<link rel="stylesheet" href="{css_file}">\n' + content
    
    with open(html_file, 'w') as file:
        file.write(new_content)
    
    browser = await launch(handleSIGINT=False, handleSIGTERM=False, handleSIGHUP=False)
    page = await browser.newPage()    
    await page.goto(f'file://{os.getcwd()}/{html_file}')

    await page.setViewport({"width": 3072, "height": 1920})
    await page.screenshot({'path': rebuilt_file})
    await browser.close()

# --------------------------------------------------------------------------------------------------------------

async def highlightAndStyles(url, selector, franklinHTML):
    download_images(url, selector)
    browser = await launch(handleSIGINT=False, handleSIGTERM=False, handleSIGHUP=False)
    page = await browser.newPage()
    await page.goto(url)

    text_file = "originalStyles.txt"
    image_file = "originalScreenshot.png"
    css_file = "franklinStyles.css"
    html_file = "franklinHTML.html"
    rebuilt_file = "regeneratedScreenshot.png"

    with open(html_file, 'w') as f:
        f.write(franklinHTML)
    
    originalStyles = await getStyles(page, selector, text_file)
    await highlight_element(page, selector)

    await page.setViewport({"width": 3072, "height": 1920})
    await page.screenshot({'path': image_file, 'fullPage': True})
    await browser.close()

    await gptResponse(franklinHTML, originalStyles, css_file)

    await regeneratedBlock(html_file, css_file, rebuilt_file)
    return image_file, rebuilt_file, text_file, css_file

# --------------------------------------------------------------------------------------------------------------

with open('elixirsolutions.html', 'r') as f:
    content1 = f.read()
with open('flickrblog.html', 'r') as f:
    content2 = f.read()
with open('snoopdogg.html', 'r') as f:
    content3 = f.read()
with open('gabnow.html', 'r') as f:
    content4 = f.read()
with open('childrensradiofoundation.html', 'r') as f:
    content5 = f.read()
with open('unicef.html', 'r') as f:
    content6 = f.read()
with open('countly.html', 'r') as f:
    content7 = f.read()
with open('simpleset.html', 'r') as f:
    content8 = f.read()

demo = gr.Interface(
    highlightAndStyles,
    [
        gr.Text(),
        gr.Text(),
        gr.Text()
    ],
    [
        gr.Image(),
        gr.Image(),
        gr.File(),
        gr.File()
    ],
    examples=[
        ["https://www.elixirsolutions.com/", ".columncontainer", f'{content1}'],
        ["https://blog.flickr.net/en", ".featured-post-wrap", f'{content2}'],
        ["https://snoopdogg.com/", ".ae-outer-wrapper", f'{content3}'],
        ["https://www.gabnow.org/", ".woocommerce.columns-4", f'{content4}'],
        ["https://childrensradiofoundation.org/", ".col-lg-8.col-sm-12", f'{content5}'],
        ["https://www.unicef.org.uk/", ".c-selected-blocks", f'{content6}'],
        ["https://countly.com/", ".grid-3-cards-ver-home", f'{content7}'],
        ["https://simpleset.net/", ".studies", f'{content8}'],
    ],
    title="Franklin CSS",
    description="",
)
demo.launch(share=False)




# default signal handlers caused the program to only run on main thread of main interpreter. KEyboard INterrupt 
# caused program to go out of main thread. Switching off those default handlers made it work.