#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
from html2image import Html2Image
hti = Html2Image()

driver_type = 'chromedriver'
driver=False


def html2im(html,css=None):

    html = html
    #css = css
    #hti.screenshot(html_str=html, save_as='red_page.png')
    hti.screenshot(html_str=html, css_str=css, save_as='red_page.png')
    return 'red_page.png'
    
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()
    out=None
    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")
            #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")
            for i,ea in enumerate(check_b):
                try:    
                    ea = int(ea)
                    ea = (ea*2)-2
                    print (ea)
                    messages[ea].screenshot(f'{i}-{uid}-tmp.png')
                    messages[ea+1].screenshot(f'{i+1}-{uid}-tmp.png')
                    im_roll = get_concat_v_cut(f'{i}-{uid}-tmp.png',f'{i+1}-{uid}-tmp.png')
                    im_roll.save(f'comb{i}-{uid}-tmp.png')
                    out_box.append(f'comb{i}-{uid}-tmp.png')
                except Exception:
                    mes="<center>Some blocks returned an error"
            if out_box:
                if len(out_box)>1:
                    im_roll = get_concat_v_cut(f'{out_box[0]}',f'{out_box[1]}')
                    im_roll.save(f'comb-{uid}-tmp.png')                
                    for i in range(2,len(out_box)):
                        im_roll = get_concat_v_cut(f'comb-{uid}-tmp.png',f'{out_box[i]}')
                        im_roll.save(f'comb-{uid}-tmp.png')
                    out = f'comb-{uid}-tmp.png'
                else:
                    tmp_im = Image.open(out_box[0])
                    tmp_im.save(f'comb-{uid}-tmp.png')
                    out = f'comb-{uid}-tmp.png'
                if check_h:
                    head_paste = get_concat_v_cut(f'head-{uid}-tmp.png',f'comb-{uid}-tmp.png')
                    head_paste.save(f'comb-{uid}-tmp.png')
                    out = f'comb-{uid}-tmp.png'
            screenshot = obj.screenshot(f'{uid}-tmp.png')
        except WebDriverException as e:
            return Image.new('RGB', (1, 1)), f'<center>Please enter a valid URL of a website/host.',out_box,out
        finally:
            if driver:
                driver.quit()
        return Image.open(f'{uid}-tmp.png'), mes,out_box,out
    else:
        return None, '<center>Please enter a valid URL of a website/host.',out_box,out

with gr.Blocks() as app:
    
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="URL (must be Public)",lines=1)
            css_inp=gr.Textbox(label="CSS")
            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,11,12,13,14,15,16,17,18,19,20], 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():
            out=gr.Image()
            outgal=gr.Gallery()
        outim = gr.Image()
    btn.click(html2im,[inp,css_inp],out)
    #btn.click(run_script,[inp,height,width,check_b,check_h],[outim,message,outgal,out])
app.queue(default_concurrency_limit=5).launch()