from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import Select
from selenium.webdriver.common.keys import Keys
import time
import os

start_url = "https://idp.company.com/idp/startSSO.ping?PartnerSpId=https://github.company.com"
username = "none"
password = os.getenv("GH_TOKEN")
avatar_url = "https://avatars.github.company.com/u/9214?&size=32"
download_path = "avatar.jpg"

chrome_options = Options()
chrome_options.add_argument("--headless")
chrome_options.add_argument("--disable-gpu")
chrome_options.add_argument("--no-sandbox")
chrome_options.add_argument("--disable-dev-shm-usage")
chrome_options.add_argument("--disable-extensions")
chrome_options.add_argument("--disable-infobars")
chrome_options.add_argument("--disable-notifications")
chrome_options.add_argument("--window-size=1920,1080")

driver = webdriver.Chrome(options=chrome_options)

try:
    driver.get(start_url)
    time.sleep(3)

    location_dropdown = Select(driver.find_element(By.ID, "location"))
    location_dropdown.select_by_value("HO")
    time.sleep(1)

    sign_on_button = driver.find_element(By.CSS_SELECTOR, "a.ping-button.normal.allow[title='Sign On']")
    sign_on_button.click()
    time.sleep(5)

    driver.get(avatar_url)
    time.sleep(3)

    avatar_element = driver.find_element(By.TAG_NAME, "img")
    with open(download_path, "wb") as file:
        file.write(avatar_element.screenshot_as_png)

    print(f"Avatar downloaded successfully to {download_path}.")

finally:
    driver.quit()