app.py

from __future__ import annotations

import os
import random
import time
import gradio as gr
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from PIL import Image
from io import BytesIO
import base64

import trafilatura

from huggingface_hub import whoami

from languages import ISO_CODE_TO_LANGUAGE_NAME

OFFLINE = os.environ.get("OFFLINE", False)

def pil_image_to_base64(image):
    # Save the image to a BytesIO buffer
    buffer = BytesIO()
    image.save(buffer, format="PNG")  # You can change the format if needed
    buffer.seek(0)

    # Encode the bytes into a base64 string
    img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

    # Format the base64 string for use in an HTML image tag
    html_img_tag_src = f"data:image/png;base64,{img_base64}"
    return html_img_tag_src

def fetch_screenshot_and_text_from_url(url):
    screen_width = 1080
    height = 350
    text = ""
    
    if OFFLINE:
        screenshot = Image.new('RGB', (350, height))
        text = f"Some dummy text for {url} (offline mode enabled)"

    else:
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')

        try:
            driver = webdriver.Chrome(options=options)
            #driver.set_window_size(1080, 720)  # Adjust the window size here
            driver.get(url)

            driver.implicitly_wait(10)

            # Wait for the page to fully load; you may adjust the sleep time or implement a wait condition
            # time.sleep(2)

            # fetch html from web page
            html_str = driver.page_source

            # Execute JS to find the full height of the rendered page
            scroll_height = driver.execute_script("return document.body.scrollHeight")

            # Resize the window to full page height
            driver.set_window_size(screen_width, max(scroll_height + 200, 900))

            raw_screenshot = driver.get_screenshot_as_png()

            screenshot = Image.open(BytesIO(raw_screenshot))

            # extract text
            text = trafilatura.extract(html_str)

        except WebDriverException as e:
            screenshot = Image.new('RGB', (1, 1))
        finally:
            if driver:
                driver.quit()


    # embed base65 encoded image as <img> tag into html string
    screenshot_html_str = f"""<div style="width: 100%; height: {height}px; overflow-y: scroll;"><img src="{pil_image_to_base64(screenshot)}" /></div>"""    

    # return gr.update(value=html_str, visible=True), text, gr.update(visible=True)
    return screenshot_html_str, text

    
with gr.Blocks(fill_height=True) as demo:

    gr.Markdown(
    """
    # Seed Crawl Annotator
    """)
        
    profile_state = gr.State([])
    gr.LoginButton()


    with gr.Column(visible=False) as wrapper_col:
        def handle_login(profile: gr.OAuthProfile | None) -> dict:
            if profile:
                gr.Info(f"Logged in as {profile.username}")
                return {
                    profile_state: f"{profile.username}", 
                    wrapper_col: gr.update(visible=True),
                }
            else:
                gr.Warning(f"You need to login to use this app.")
                return {
                    profile_state: None,
                    wrapper_col: gr.update(visible=False),
                }
            
        demo.load(handle_login, inputs=None, outputs=[profile_state, wrapper_col])

        url_field = gr.Textbox(label="Website URL", placeholder="Enter a URL you want to annotate", interactive=True)

        with gr.Row():
            set_random_btn = gr.Button("Set Random URL", variant="secondary", interactive=True)

            load_btn = gr.Button("Annotate URL", variant="primary", interactive=True)

        with gr.Row():
            extracted_text = gr.Textbox(label="Extracted text", max_lines=15, lines=15, visible=False, placeholder="Click on `Load URL` to fetch Web page's text content.")
            
            screenshot_scrollable = gr.HTML(visible=False)

        with gr.Column(visible=False) as output_col:
            with gr.Row():
                language_codes = gr.Dropdown(
                        [("unknown", "unknown")] + [(f"{code}: {name}", code) for code, name in ISO_CODE_TO_LANGUAGE_NAME.items()], 
                        label="Language codes",
                        multiselect=True,
                        # allow_custom_value=True,
                )
                categories = gr.CheckboxGroup(["News", "Culture/History", "Government", "Political Parties", "Other"], label="Categories")

            with gr.Row():
                do_crawl_btn = gr.Button("✅ Do Crawl", elem_classes="success")
                dont_crawl_btn = gr.Button("❌ Don't Crawl", elem_classes="error")
                # random_subpage_btn = gr.Button("🔁 Load Another Subpage", variant="secondary")


        def set_random_url():
            candidate_urls = [
                "http://example.com",
                "https://wikipedia.org/",
                "https://occiglot.eu",
                "https://ostendorff.org",
                "https://fr.wikipedia.org/",
                "https://amazon.com/"
            ]
            selected_url = random.choice(candidate_urls)
            return selected_url    
            
        set_random_btn.click(fn=set_random_url, outputs=url_field)


        def load_url(url):
            screenshot_html_str, text = fetch_screenshot_and_text_from_url(url)

            if not screenshot_html_str or not text:
                gr.Error("Could not fetch data for url")
            else:

                return {
                    screenshot_scrollable: gr.update(value=screenshot_html_str, visible=True), 
                    extracted_text:  gr.update(value=text, visible=True),
                    output_col: gr.update(visible=True),
                    language_codes: "unknown", # Reset by set to invalid value # gr.update(None, label=url),
                    categories:  gr.update(value=None),
                }

        load_btn.click(fn=load_url, inputs=url_field, outputs=[screenshot_scrollable, extracted_text, output_col, language_codes, categories], api_name="load_url")

        def do_crawl(profile_state, url, language_codes, categories, do_crawl=True):

            if profile_state:
                html_str = f"<b>Thanks {profile_state}, we have saved your feedback!</b>"
                gr.Info("Thanks for your feedback")
            else:
                gr.Error("Feedback could not be saved")
                html_str = f"<b>Feedback could not be saved.</b> You are not authenticated."

            return {
                url_field: "",
                output_col: gr.update(visible=False),
                extracted_text: gr.update(value=None, visible=False),
                screenshot_scrollable: gr.update(value="", visible=False),
            }

        # def do_crawl(profile_state, url, language_codes, categories):
        #     return do_crawl_or_not(profile_state, url, language_codes, categories, do_crawl=True)
            
        # def dont_crawl(profile_state, url, language_codes, categories):
        #     return do_crawl_or_not(profile_state, url, language_codes, categories, do_crawl=False)
            
        
        do_crawl_btn.click(
            fn=do_crawl, 
            inputs=[profile_state, url_field, language_codes, categories], 
            outputs=[
                url_field, 
                output_col, 
                extracted_text, 
                screenshot_scrollable
            ], 
            api_name="do_crawl",
        )
        dont_crawl_btn.click(
            fn=do_crawl, 
            inputs=[profile_state, url_field, language_codes, categories], 
            outputs=[
                url_field, 
                output_col, 
                extracted_text, 
                screenshot_scrollable
            ], 
            api_name="do_crawl",
        )

        # dont_crawl_btn.click(fn=dont_crawl, inputs=[profile_state, url, language_codes, categories], outputs=[url, output_col, extracted_text, screenshot_scrollable], api_name="dont_crawl")

        # def random_subpage(url):
        #     new_url = "http://example.com"

        #     return [new_url, *fetch_screenshot_and_text_from_url(new_url)]

        # random_subpage_btn.click(fn=random_subpage, inputs=url, outputs=[url, screenshot_scrollable, extracted_text, output_col], api_name="load_random_subpage")



if __name__ == "__main__":
    demo.launch()