#from __future__ import annotations
from selenium import webdriver
import gradio as gr
import uuid
import re
from PIL import Image
from io import BytesIO
from pathlib import Path
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
#from selenium_webdriver import WebElement

driver_type = 'chromedriver'
driver=False

def get_concat_h_cut(in1, in2):
    im1=Image.open(in1)
    im2=Image.open(in2)
    dst = Image.new('RGB', (im1.width + im2.width,
                            min(im1.height, im2.height)))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (im1.width, 0))
    return dst


def get_concat_v_cut(in1, in2):
    im1=Image.open(in1)
    im2=Image.open(in2)    
    dst = Image.new(
        'RGB', (min(im1.width, im2.width), im1.height + im2.height))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (0, im1.height))
    return dst



def run_script(url: str, height: int, width: int, check_b,check_h):
    mes_box=[]
    out_box=[]
    uid=uuid.uuid4()
    is_url=True
    if is_url:
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        mes='<center>operation success'
        try:
            driver = webdriver.Chrome(options=options)
            driver.get(url)
            html=driver.page_source
            #print (driver.html)
            driver.implicitly_wait(30)
            driver.set_window_size(int(width), int(height))
            page = driver.find_element(By.TAG_NAME, "html")
            print (dir(page))
            print (page.text)
            driver.execute_script("arguments[0].style.background = 'blue';",page)
            main_head = driver.find_element(By.CLASS_NAME, "main-content-header")
            head_shot = main_head.screenshot(f'head-{uid}-tmp.png')
            
            obj = driver.find_element(By.CLASS_NAME, "main")            
            messages = driver.find_elements(By.CLASS_NAME, "message")
            print (len(messages))
            mes_blocks=len(messages)%2
            for ea in check_b:
                try:    
                    ea = int(ea)
                    ea = (ea*2)-2
                    print (ea)
                    messages[ea].screenshot(f'{ea}-{uid}-tmp.png')
                    messages[ea+1].screenshot(f'{ea+1}-{uid}-tmp.png')
                    if check_h:
                        head_paste = get_concat_v_cut(f'head-{uid}-tmp.png',f'{ea}-{uid}-tmp.png')
                        print(head_paste)
                        out_box.append(get_concat_v_cut(head_paste,f'{ea+1}-{uid}-tmp.png'))
                    if not check_h:    
                        out_box.append(get_concat_v_cut(f'{ea}-{uid}-tmp.png',f'{ea+1}-{uid}-tmp.png'))
                    print(f'out_box:: {out_box}')
                    #mes_box.clear()
                except Exception:
                    mes="<center>Some blocks returned an error"
            '''
            for i,mes in enumerate(messages):
                print (mes.text)
                if not (i%2):
                    messages[i].screenshot(f'{i}-{uid}-tmp.png')
                    messages[i+1].screenshot(f'{i+1}-{uid}-tmp.png')
                    mes_box.append(())
                    out_box.append(get_concat_v_cut(f'{i}-{uid}-tmp.png',f'{i+1}-{uid}-tmp.png'))
                    print(f'out_box:: {out_box}')
                    #mes_box.clear()
            '''    
            #obj = driver.find_element(By.ID, "app-container")            
            #obj = driver.find_element(By.ID, "conversation-feed")            
            #obj = driver.find_element_by_id("in_html")
            ######################
            #inp = driver.find_element("id","input")
            '''
            find_element(By.ID, "id")
            find_element(By.NAME, "name")
            find_element(By.XPATH, "xpath")
            find_element(By.LINK_TEXT, "link text")
            find_element(By.PARTIAL_LINK_TEXT, "partial link text")
            find_element(By.TAG_NAME, "tag name")
            find_element(By.CLASS_NAME, "class name")
            find_element(By.CSS_SELECTOR, "css selector")       
            '''
            #inp.send_keys("test")
            ######################
            cookie_jar = []
            #cookie_jar.append(driver.get_cookies())
            #print(cookie_jar)
            screenshot = obj.screenshot(f'{uid}-tmp.png')
            #screenshot = obj.get_screenshot_as_png()
        except WebDriverException as e:
            return Image.new('RGB', (1, 1)), f'<center>{e}',out_box
        finally:
            if driver:
                driver.quit()

        #return [Image.open(BytesIO(screenshot)), 'operation success.',cookie_jar,html]
        return Image.open(f'{uid}-tmp.png'), mes,out_box
    else:
        return None, '<center>Please enter a valid URL of a website/host.',out_box

with gr.Blocks() as app:
    
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="URL",lines=1)
            btn= gr.Button()
            
        with gr.Column():
            check_h=gr.Checkbox(label="Show Header", value=True)
            check_b=gr.CheckboxGroup(label="Chatblocks", choices=[1,2,3,4,5,6,7,8,9,10], value=[1])
            with gr.Row():
                height=gr.Number(label="Height", value=4096)
                width=gr.Number(label="Width",value=800)
    message=gr.HTML('<center>Enter URL')
    with gr.Row():
        with gr.Column():
            outgal=gr.Gallery()
        outim = gr.Image()

    btn.click(run_script,[inp,height,width,check_b,check_h],[outim,message,outgal])
app.launch()