import org.openqa.selenium.*;
import org.openqa.selenium.chrome.ChromeDriver;
import org.testng.ITestResult;
import org.testng.annotations.*;
import java.io.File;
import java.nio.file.Files;
import java.time.Duration;

public class ScreenshotTest {

    WebDriver driver;

    @BeforeMethod
    public void setup() {
        driver = new ChromeDriver();
        driver.manage().timeouts().implicitlyWait(Duration.ofSeconds(10));
        driver.manage().window().setSize(new Dimension(1920, 1080));
    }

    @Test
    public void testFullPageScreenshot() throws Exception {
        driver.get("https://www.selenium.dev/");

        File screenshot = ((TakesScreenshot) driver)
            .getScreenshotAs(OutputType.FILE);

        File dest = new File("screenshots/full-page.png");
        dest.getParentFile().mkdirs();
        Files.copy(screenshot.toPath(), dest.toPath());
        System.out.println("Screenshot saved: " + dest.getAbsolutePath());
    }

    @Test
    public void testElementScreenshot() throws Exception {
        driver.get("https://www.selenium.dev/selenium/web/web-form.html");

        // Selenium 4: Element-level screenshot
        WebElement form = driver.findElement(By.tagName("form"));
        File elementShot = form.getScreenshotAs(OutputType.FILE);

        File dest = new File("screenshots/element.png");
        dest.getParentFile().mkdirs();
        Files.copy(elementShot.toPath(), dest.toPath());
        System.out.println("Element screenshot saved");
    }

    @Test
    public void testScreenshotAsBase64() {
        driver.get("https://www.selenium.dev/");

        String base64 = ((TakesScreenshot) driver)
            .getScreenshotAs(OutputType.BASE64);

        System.out.println("Base64 length: " + base64.length());
        // Useful for embedding in reports
    }

    // Auto-screenshot on failure
    @AfterMethod
    public void captureOnFailure(ITestResult result) throws Exception {
        if (result.getStatus() == ITestResult.FAILURE) {
            File screenshot = ((TakesScreenshot) driver)
                .getScreenshotAs(OutputType.FILE);
            String name = result.getName();
            File dest = new File("screenshots/failures/" + name + ".png");
            dest.getParentFile().mkdirs();
            Files.copy(screenshot.toPath(), dest.toPath());
            System.out.println("Failure screenshot: " + dest.getPath());
        }
        if (driver != null) driver.quit();
    }
}

from selenium import webdriver
from selenium.webdriver.common.by import By
import pytest, os

@pytest.fixture
def driver():
    d = webdriver.Chrome()
    d.implicitly_wait(10)
    d.set_window_size(1920, 1080)
    yield d
    d.quit()


def test_full_page_screenshot(driver):
    driver.get("https://www.selenium.dev/")

    os.makedirs("screenshots", exist_ok=True)
    driver.save_screenshot("screenshots/full-page.png")
    print("Screenshot saved")


def test_element_screenshot(driver):
    """Selenium 4: Element-level screenshot"""
    driver.get("https://www.selenium.dev/selenium/web/web-form.html")

    form = driver.find_element(By.TAG_NAME, "form")
    os.makedirs("screenshots", exist_ok=True)
    form.screenshot("screenshots/element.png")
    print("Element screenshot saved")


def test_screenshot_as_base64(driver):
    driver.get("https://www.selenium.dev/")

    base64_data = driver.get_screenshot_as_base64()
    print(f"Base64 length: {len(base64_data)}")


# Auto-screenshot on failure via conftest.py
# @pytest.hookimpl(tryfirst=True, hookwrapper=True)
# def pytest_runtest_makereport(item, call):
#     outcome = yield
#     report = outcome.get_result()
#     if report.when == "call" and report.failed:
#         driver = item.funcargs.get("driver")
#         if driver:
#             os.makedirs("screenshots/failures", exist_ok=True)
#             path = f"screenshots/failures/{item.name}.png"
#             driver.save_screenshot(path)
#             print(f"Failure screenshot: {path}")