heygen_gui.py

import os
import random
import re
import sys
import threading
import time

import pandas as pd
from PyQt6.QtCore import QObject, pyqtSignal, pyqtSlot
from PyQt6.QtWidgets import (QApplication, QFileDialog, QGroupBox, QHBoxLayout,
                             QLabel, QLineEdit, QMainWindow, QMessageBox,
                             QPushButton, QTextEdit, QVBoxLayout, QWidget)
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait


# Worker Signals to update UI from thread
class WorkerSignals(QObject):
    log = pyqtSignal(str)
    error = pyqtSignal(str)
    finished = pyqtSignal()

class HeyGenApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("HeyGen Automation Tool")
        self.setGeometry(100, 100, 600, 700)

        # Control Flags
        self.stop_event = threading.Event()
        self.pause_event = threading.Event()
        self.pause_event.set() # Start unpaused
        self.driver = None
        self.signals = WorkerSignals()
        self.signals.log.connect(self.log_message)
        self.signals.error.connect(self.log_error)
        self.signals.finished.connect(self.on_automation_finished)

        self.extracted_ids = []   # list of tuples (name, video_title)
        self.data_lock = threading.Lock()


        self.init_ui()

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # 1. Configuration Checkbox
        group_config = QGroupBox("Configuration")
        layout_config = QHBoxLayout()
        
        self.path_input = QLineEdit()
        self.path_input.setPlaceholderText("Select CSV file...")
        layout_config.addWidget(QLabel("CSV File:"))
        layout_config.addWidget(self.path_input)
        
        btn_browse = QPushButton("Browse")
        btn_browse.clicked.connect(self.browse_file)
        layout_config.addWidget(btn_browse)
        
        group_config.setLayout(layout_config)
        layout.addWidget(group_config)

        # 2. Browser Controls
        group_browser = QGroupBox("1. Browser & Login")
        layout_browser = QVBoxLayout()
        
        self.btn_open_browser = QPushButton("Open Browser")
        self.btn_open_browser.clicked.connect(self.open_browser)
        self.btn_open_browser.setStyleSheet("background-color: #e1f5fe; padding: 5px;")
        layout_browser.addWidget(self.btn_open_browser)
        
        self.btn_confirm_login = QPushButton("I Have Logged In (Press OK)")
        self.btn_confirm_login.clicked.connect(self.confirm_login)
        self.btn_confirm_login.setEnabled(False)
        self.btn_confirm_login.setStyleSheet("background-color: #dcedc8; padding: 5px;")
        layout_browser.addWidget(self.btn_confirm_login)
        
        group_browser.setLayout(layout_browser)
        layout.addWidget(group_browser)

        # 3. Automation Controls
        group_run = QGroupBox("2. Automation Controls")
        layout_run = QHBoxLayout()
        
        self.btn_start = QPushButton("Start Automation")
        self.btn_start.clicked.connect(self.start_automation)
        self.btn_start.setEnabled(False)
        self.btn_start.setStyleSheet("background-color: #c8e6c9; padding: 5px;")
        layout_run.addWidget(self.btn_start)
        
        self.btn_pause = QPushButton("Pause")
        self.btn_pause.clicked.connect(self.toggle_pause)
        self.btn_pause.setEnabled(False)
        self.btn_pause.setStyleSheet("background-color: #fff9c4; padding: 5px;")
        layout_run.addWidget(self.btn_pause)
        
        self.btn_stop = QPushButton("Stop")
        self.btn_stop.clicked.connect(self.stop_automation)
        self.btn_stop.setEnabled(False)
        self.btn_stop.setStyleSheet("background-color: #ffcdd2; padding: 5px;")
        layout_run.addWidget(self.btn_stop)
        

        # Extract IDS button
        self.btn_extract = QPushButton("Extract IDs")
        self.btn_extract.clicked.connect(self.export_ids_to_csv)
        self.btn_extract.setEnabled(True)
        self.btn_extract.setStyleSheet("background-color: #bbdefb; padding: 5px;")
        layout_run.addWidget(self.btn_extract)

        group_run.setLayout(layout_run)
        layout.addWidget(group_run)

        # 4. Logs
        group_log = QGroupBox("Logs")
        layout_log = QVBoxLayout()
        self.text_log = QTextEdit()
        self.text_log.setReadOnly(True)
        layout_log.addWidget(self.text_log)
        group_log.setLayout(layout_log)
        layout.addWidget(group_log)

    def log_message(self, msg):
        self.text_log.append(msg)
        sb = self.text_log.verticalScrollBar()
        sb.setValue(sb.maximum())

    def log_error(self, msg):
        error_msg = f"<span style='color: red;'>❌ {msg.upper()}</span>"
        self.text_log.append(error_msg)
        sb = self.text_log.verticalScrollBar()
        sb.setValue(sb.maximum())

    def browse_file(self):
        fname, _ = QFileDialog.getOpenFileName(self, "Open CSV", "", "CSV Files (*.csv)")
        if fname:
            self.path_input.setText(fname)
            self.log_message(f"Selected CSV: {fname}")

    def export_ids_to_csv(self):
        if not self.extracted_ids:
            self.log_message("No extracted IDs to export.")
            return
        
        try:
            csv_dir = os.path.dirname(self.path_input.text())
            csv_path = os.path.join(csv_dir, "IDs.csv")
            
            df_new = pd.DataFrame(
                self.extracted_ids,
                columns=["Name", "URL", "ID"]
            )

            # Create or append to IDs.csv
            if os.path.exists(csv_path):
                df_existing = pd.read_csv(csv_path)
                df_final = pd.concat([df_existing, df_new], ignore_index=True)
            else:
                df_final = df_new

            df_final.to_csv(csv_path, index=False)

            self.log_message(
                f"Exported {len(self.extracted_ids)} extracted IDs to {csv_path}"
            )

        except PermissionError:
            self.log_message(f"Failed: Please close IDs.csv if it's open in another program.")
        except Exception as e:
            self.log_message(f"Failed to export IDs: {e}")


    def open_browser(self):
        if not self.path_input.text():
            QMessageBox.warning(self, "Error", "Please select a CSV file first.")
            return

        try:
            self.log_message("Opening browser...")
            options = Options()
            profile_path = os.path.join(os.getcwd(), "selenium_profile")
            options.add_argument(f"user-data-dir={profile_path}")
            options.add_argument("--start-maximized")
            
            self.driver = webdriver.Chrome(options=options)
            self.wait = WebDriverWait(self.driver, 30)
            self.actions = ActionChains(self.driver)
            
            self.driver.get("https://app.heygen.com")
            
            self.log_message("Browser opened. Please login manually if needed.")
            
            self.btn_open_browser.setEnabled(False)
            self.btn_confirm_login.setEnabled(True)
            
        except Exception as e:
            self.log_message(f"Error opening browser: {e}")
            QMessageBox.critical(self, "Error", f"Failed to open browser: {e}")

    def confirm_login(self):
        self.log_message("Login confirmed. Ready to start.")
        self.btn_confirm_login.setEnabled(False)
        self.btn_start.setEnabled(True)

    def toggle_pause(self):
        if self.pause_event.is_set():
            self.pause_event.clear()
            self.btn_pause.setText("Resume")
            self.log_message("Paused.")
        else:
            self.pause_event.set()
            self.btn_pause.setText("Pause")
            self.log_message("Resumed.")

    def stop_automation(self):
        self.stop_event.set()
        self.log_message("Stopping...")
        self.btn_start.setEnabled(True)
        self.btn_stop.setEnabled(False)
        self.btn_pause.setEnabled(False)

    def start_automation(self):
        self.stop_event.clear()
        self.pause_event.set()
        
        self.btn_start.setEnabled(False)
        self.btn_stop.setEnabled(True)
        self.btn_pause.setEnabled(True)
        
        thread = threading.Thread(target=self.run_logic)
        thread.daemon = True
        thread.start()
    
    @pyqtSlot()
    def on_automation_finished(self):
        self.btn_stop.setEnabled(False)
        self.btn_pause.setEnabled(False)
        self.btn_start.setEnabled(True)
        self.export_ids_to_csv()
        self.log_message("✅ Automation finished.")

    def run_logic(self):
        try:
            df = pd.read_csv(self.path_input.text())
            # Drop rows with missing URL or Name
            df = df.dropna(subset=["URL", "Name"])
            links = df["URL"].tolist()
            names = df["Name"].tolist()
            
            self.signals.log.emit(f"Found {len(links)} videos to process.")

            for i, (link, name) in enumerate(zip(links, names)):
                # Check Stop
                if self.stop_event.is_set():
                    self.signals.log.emit("Process stopped by user.")
                    break

                # Check Pause
                while not self.pause_event.is_set():
                    if self.stop_event.is_set(): break
                    time.sleep(1)
                
                self.signals.log.emit(f"▶️ Starting ({i+1}/{len(links)}): {name}")
                
                try:
                    self.process_video(link, name)
                    self.signals.log.emit(f"✅ Completed: {name}")
                except Exception as e:
                    self.signals.error.emit(f"❌ERROR PROCESSING {name}: {e}")
                    
                # Rate Limiting / Sleep
                if i < len(links) - 1:
                    sleep_time = random.randint(30, 60)
                    for _ in range(sleep_time):
                        if self.stop_event.is_set(): break
                        time.sleep(1)

            self.signals.finished.emit()
            
        except Exception as e:
            self.signals.error.emit(f"❌CRITICAL ERROR: {e}")
            self.signals.finished.emit()

    def process_video(self, link, name):
         # 1. Click CREATE button
        try:
            WebDriverWait(self.driver, 30).until(
                EC.presence_of_element_located(
                    (By.XPATH, "//span[normalize-space()='Create']")
                )
            )

            create_btn = WebDriverWait(self.driver, 30).until(
                lambda d: next(
                    (btn for btn in d.find_elements(
                        By.XPATH,
                        "//span[normalize-space()='Create']/ancestor::button"
                    )
                    if btn.is_displayed() and btn.is_enabled()),
                    None
                )
            )

            if not create_btn:
                raise Exception("❌COULD NOT FIND VISIBLE CREATE BUTTON")

            self.actions.click(create_btn).perform()
            time.sleep(0.5)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR CLICKING CREATE BUTTON: {e}")
            raise

        # 2. Click "Translate a video"
        try:
            translate_btn = WebDriverWait(self.driver, 20).until(
                EC.element_to_be_clickable((By.XPATH, "//span[text()='Translate a video']"))
            )
            self.driver.execute_script(
                "arguments[0].scrollIntoView({block:'center'});",
                translate_btn
            )
            time.sleep(0.5)
            self.actions.click(translate_btn).perform()
            time.sleep(1)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR CLICKING TRANSLATE BUTTON: {e}")
            raise

        # 3. Click "Video dubbing"
        try:
            video_dubbing_btn = WebDriverWait(self.driver, 10).until(
                EC.element_to_be_clickable((By.XPATH, "//span[text()='Video dubbing']"))
            )
            self.driver.execute_script(
                "arguments[0].scrollIntoView({block:'center'});",
                video_dubbing_btn
            )
            time.sleep(0.5)
            self.actions.click(video_dubbing_btn).perform()
            time.sleep(1)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR CLICKING VIDEO DUBBING BUTTON: {e}")
            raise

        # 4. Paste Drive link
        try:
            input_box = self.wait.until(EC.presence_of_element_located(
                (By.XPATH, "//input[@placeholder=\"Paste your video's URL here\"]")
            ))
            self.driver.execute_script(
                "arguments[0].scrollIntoView({block:'center'});",
                input_box
            )
            time.sleep(0.3)
            input_box.clear()
            input_box.send_keys(link)
            time.sleep(0.5)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR PASTING LINK: {e}")
            raise

        # 5. Click ✓ validation button
        try:
            check_button = self.wait.until(EC.element_to_be_clickable(
                (By.XPATH, "//input[@placeholder=\"Paste your video's URL here\"]/following::button[1]")
            ))
            self.actions.click(check_button).perform()
            time.sleep(2)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR CLICKING VALIDATION BUTTON: {e}")
            raise
        
        # 6. Click NEXT
        try:
            next_btn = self.wait.until(EC.element_to_be_clickable(
                (By.XPATH, "//button[contains(text(),'Next')]")
            ))
            self.driver.execute_script(
                "arguments[0].scrollIntoView({block:'center'});",
                next_btn
            )
            time.sleep(0.3)
            self.actions.click(next_btn).perform()
            time.sleep(0.5)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR CLICKING NEXT BUTTON: {e}")
            raise

        # Wait for language options
        try:
            lang_container = self.wait.until(EC.element_to_be_clickable(
                (By.XPATH, "//div[contains(@class,'tw-h-full tw-flex-1 tw-overflow-hidden')]")
            ))
            self.actions.click(lang_container).perform()
            time.sleep(0.5)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR INTERACTING WITH LANGUAGE CONTAINER: {e}")
            raise

        # 7. Choose English
        try:
            english_btn = self.wait.until(EC.element_to_be_clickable(
                (By.XPATH, "//div[contains(text(),'English')]")
            ))
            self.driver.execute_script(
                "arguments[0].scrollIntoView({block:'center'});",
                english_btn
            )
            time.sleep(0.3)
            self.actions.click(english_btn).perform()
            time.sleep(random.randint(2, 5))
        except Exception as e:
            self.signals.error.emit(f"❌ERROR SELECTING ENGLISH: {e}")
            raise
        
        # 8. Click TRANSLATE
        try:
            translate_final = self.wait.until(EC.element_to_be_clickable(
                (By.XPATH, "//button[text()='Translate']")
            ))
            self.driver.execute_script(
                "arguments[0].scrollIntoView({block:'center'});",
                translate_final
            )
            time.sleep(0.3)
            self.actions.click(translate_final).perform()
            time.sleep(5)
        except Exception as e:
            self.signals.error.emit(f"❌ERROR CLICKING TRANSLATE BUTTON: {e}")
            raise
        
        try:
            time.sleep(2)
            file_id = re.search(r"/d/([^/]+)", link).group(1)
            
            # 9. Hover/Find video
            video_title = self.wait.until(EC.element_to_be_clickable(
                (By.XPATH, f"//*[contains(text(),'{file_id}')]")
            ))
            
            video_title.click()

            # 10. Rename
            active = self.driver.switch_to.active_element
            active.send_keys(Keys.CONTROL, "a")
            active.send_keys(Keys.BACKSPACE)
            active.send_keys(name)
            active.send_keys(Keys.ENTER)
            time.sleep(0.5)
            
            # 11. The card is already hovered/selected after rename, find and click the more button directly
            more_btn = WebDriverWait(self.driver, 10).until(
                EC.visibility_of_element_located(
                    (By.XPATH, "//button[@aria-haspopup='dialog']//iconpark-icon[@name='more-level']/ancestor::button")
                )
            )
            self.driver.execute_script("arguments[0].click();", more_btn)
            time.sleep(0.2)
            
            # 12. Wait for dialog and click "Copy ID"
            copy_id_btn = WebDriverWait(self.driver, 10).until(
                EC.element_to_be_clickable(
                    (By.XPATH, "//span[text()='Copy ID']/ancestor::div[contains(@class,'tw-cursor-pointer')]")
                )
            )
            self.actions.click(copy_id_btn).perform()
            time.sleep(0.2)
            
            # 13. Get the copied ID from clipboard
            copied_id = self.driver.execute_script("return navigator.clipboard.readText();")
            
            # If clipboard access fails, try alternative method
            if not copied_id:
                copied_id = file_id
            
            # Store in extracted_ids (name, url, id)
            with self.data_lock:
                self.extracted_ids.append((name, link, copied_id))
            
        except Exception as e:
            self.signals.error.emit(f"❌ERROR IN FINAL STEPS: {e}")
            raise


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = HeyGenApp()
    window.show()
    sys.exit(app.exec())