from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import WebDriverException, TimeoutException

import threading

from PIL import Image
from io import BytesIO
import re

import time
import os


class LLM:
    def __init__(self, demo) -> None:
        self.demo = demo
        self.value = None

        self.input_tag = (By.XPATH, "//textarea[@data-testid='textbox']")


    def call_function(self):
        pass


    def get_response_by_label(self):
        response_element = WebDriverWait(self.driver, 30).until(
            EC.visibility_of_all_elements_located(self.reponse_tag)
        )

        while True:
            current_len = len(response_element[-1].get_attribute("aria-label").replace("bot's message: ", ''))
            time.sleep(1)
            previous_len = len(response_element[-1].get_attribute("aria-label").replace("bot's message: ", ''))

            print(current_len, previous_len)
            if current_len == previous_len and current_len > 1:
                break
            else:
                print('Not yet...')

        return response_element[-1].get_attribute("aria-label").replace("bot's message: ", '')


    def generate(self, text):
        self.options = webdriver.ChromeOptions()
        self.options.add_argument("--disable-gpu")
        self.options.add_argument('--no-sandbox')
        self.options.add_argument('--disable-dev-shm-usage')
        if not self.demo: self.options.add_argument('--headless')
        self.driver = webdriver.Chrome(options=self.options)

        try:
            self.driver.get(self.url)

            input_element = WebDriverWait(self.driver, 60).until(
                EC.element_to_be_clickable(self.input_tag)
            )
            input_element.clear()
            input_element.send_keys(text)

            self.call_function()

        except WebDriverException as e:
            if self.demo: print(e)
        finally:
            print('Driver exited...')
            if self.driver:
                self.driver.quit()
            return self.value


class MultimodalLLM(LLM):
    def __init__(self, demo) -> None:
        super().__init__(demo)

        self.upload_tag = (By.XPATH, "//input[@type='file']")
        self.image_tag = (By.XPATH, "//div[contains(@class, 'image-frame')]")

    def call_function(self):
        pass


    def generate(self, image_path, text):
        self.options = webdriver.ChromeOptions()
        self.options.add_argument("--disable-gpu")
        self.options.add_argument('--no-sandbox')
        self.options.add_argument('--disable-dev-shm-usage')
        if not self.demo: self.options.add_argument('--headless')
        self.driver = webdriver.Chrome(options=self.options)

        try:
            self.driver.get(self.url)
            upload_element = WebDriverWait(self.driver, 60).until(
                EC.presence_of_element_located(self.upload_tag)
            )
            print(f"connect to {self.url}")
        
            image_path = os.path.abspath(image_path)
            upload_element.send_keys(image_path)

            WebDriverWait(self.driver, 10).until(
                EC.visibility_of_element_located(self.image_tag)
            )

            input_element = WebDriverWait(self.driver, 10).until(
                EC.element_to_be_clickable(self.input_tag)
            )
            input_element.clear()
            input_element.send_keys(text)
            print(f"send question: {text}")

            self.call_function()


        except WebDriverException as e:
            if self.demo: print(e)
        finally:
            print('Driver exited...')
            if self.driver:
                self.driver.quit()
            return self.value


class GraphAnalyzer(MultimodalLLM):
    def __init__(self, demo=True, url="https://mplug-tinychart-3b.hf.space") -> None:
        super().__init__(demo)
        self.url = url

        self.submit_tag = (By.XPATH, "//button[text()='Send']")
        self.reponse_tag = (By.XPATH, "//*[@id='chatbot']/div[2]/div[2]/div/div[2]/div/button")
        self.alert_tag = (By.XPATH, "/html/body/gradio-app/div/div/div[2]/div/div/div[2]/div[2]")


    def call_function(self):
        submit_element = self.driver.find_element(*self.submit_tag)
        submit_element.click()

        try:
            alert_element = WebDriverWait(self.driver, 5).until(
                EC.visibility_of_all_elements_located(self.alert_tag)
            )
            alert = alert_element[0].text
            print('Alert:', alert)
            self.value = {'status': 'alert', 'text': alert}
            #retry_time = re.search(r"\d{1,2}:\d{2}:\d{2}", alert)

            return
        except TimeoutException: pass

        response = self.get_response_by_label()
        print('Reponse:', response)
        self.value = {'status': 'success', 'text': response}


class ImageDescriptor(MultimodalLLM):
    def __init__(self, demo=True, url="https://vikhyatk-moondream2.hf.space") -> None:
        super().__init__(demo)
        self.url = url

        self.submit_tag = (By.XPATH, "//button[text()='Submit']")
        self.reponse_tag = (By.XPATH, "//*[@id='component-9']/div[2]/div/span/p")
        self.alert_tag = (By.XPATH, "/html/body/gradio-app/div/div/div[2]/div/div/div[2]/div[2]")
       
    def call_function(self):
        submit_element = self.driver.find_element(*self.submit_tag)
        submit_element.click()
        """
        try:
            alert_element = WebDriverWait(self.driver, 5).until(
                EC.visibility_of_all_elements_located(self.alert_tag)
            )
            alert = alert_element[0].text
            print('Alert:', alert)
            self.value = {'status': 'alert', 'text': alert}
            #retry_time = re.search(r"\d{1,2}:\d{2}:\d{2}", alert)

            return
        except TimeoutException: pass
        """
        response_element = WebDriverWait(self.driver, 30).until(
            EC.visibility_of_all_elements_located(self.reponse_tag)
        )
        
        while True:
            current_len = len(response_element[0].text) - 1
            time.sleep(1)
            previous_len = len(response_element[0].text) - 1

            print(current_len, previous_len)
            if current_len == previous_len and current_len > 1:
                break
            else:
                print('Not yet...')

        response = response_element[0].text
        print('Reponse:', response)
        self.value = {'status': 'success', 'text': response}


class DataSummarizer(LLM):
    def __init__(self, demo=True, url="https://yousef4-mistral-7b-v0-3-chat.hf.space") -> None:
        super().__init__(demo)
        self.url = url
        self.submit_tag = (By.XPATH, "//button[text()='Submit']")
        self.reponse_tag = (By.XPATH, "//*[@data-testid='bot']")

    def call_function(self):
        submit_element = self.driver.find_element(*self.submit_tag)
        submit_element.click()

        response = self.get_response_by_label()
        print('Reponse:', response)
        self.value = {'status': 'success', 'text': response}


class MultiWorkers:
    def __init__(self, demo=True) -> None:
        self.x = GraphAnalyzer(demo)
        self.y = ImageDescriptor(demo)
        self.z = DataSummarizer(demo)


if __name__ == '__main__':
    x = GraphAnalyzer()
    x.generate('./temp/graph_1.png', 'Give me insight about the graph')
    #y = ImageDescriptor()
    #y.generate("C:/Users/ACER/Downloads/ComfyUI_07194_.png", "Describe the image")

    #z = DataSummarizer(demo=True)
    #zz = z.generate("John Lennon")
    #print('>>', zz)
    pass