import os
import base64
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 NoSuchElementException, TimeoutException

# Definindo o diretório onde as imagens serão salvas
output_dir = 'captcha_images'
os.makedirs(output_dir, exist_ok=True)

driver = webdriver.Chrome()

try:
    # URL da página
    url = "https://pje.trt2.jus.br/consultaprocessual/captcha/detalhe-processo/0000005-09.2014.5.02.0080/1"

    for i in range(2000):
        driver.get(url)
        
        try:
            # Esperar até que o captcha esteja presente e visível
            captcha_element = WebDriverWait(driver, 10).until(
                EC.visibility_of_element_located((By.XPATH, '//*[@id="imagemCaptcha"]'))
            )

            # Pegar a URL da imagem do captcha (Base64)
            captcha_base64 = captcha_element.get_attribute('src').split(',')[1]  # Remover o prefixo 'data:image/png;base64,'

            # Decodificar o Base64 e salvar como um arquivo de imagem
            file_path = os.path.join(output_dir, f"captcha_{i+1}.png")
            with open(file_path, "wb") as file:
                file.write(base64.b64decode(captcha_base64))
            
            print(f"Imagem captcha {i+1} baixada com sucesso.")
        
        except TimeoutException:
            print("O elemento captcha não foi encontrado no tempo limite especificado.")
        except NoSuchElementException:
            print("O elemento captcha não foi encontrado na página.")

finally:
    print("Fechando o navegador.")
    driver.quit()