<?php

declare(strict_types=1);

namespace Tests\Support\Helper;

use Codeception\Module\WebDriver;
use Facebook\WebDriver\Remote\RemoteWebElement;
use Facebook\WebDriver\WebDriverBy;

class ShadowAwareWebDriver extends WebDriver
{
    public function findElement(string $selector): RemoteWebElement
    {
        $items = $this->findElements($selector);
        return $items[array_key_first($items)];
    }

    /**
     * @param string $selector
     * @return RemoteWebElement[]
     */
    public function findElements(string $selector): array
    {
        if (!str_contains($selector, '::shadow')) {
            return $this->_findElements($selector);
        }
        $selectors = array_map(
            fn(string $selector) => WebDriverBy::cssSelector(trim($selector)),
            explode('::shadow', $selector)
        );
        $firstSelector = array_shift($selectors);
        $topLevelRoots = $this->_findElements($firstSelector);

        $result = [];

        foreach ($topLevelRoots as $root) {
            foreach ($this->resolveShadowSelector($root, $selectors) as $element) {
                $result[] = $element;
            }
        }

        return $result;
    }

    /**
     * @param WebDriverBy[] $selectors
     * @return RemoteWebElement[]
     */
    protected function resolveShadowSelector(RemoteWebElement $root, array $selectors): array
    {
        if (empty($selectors)) {
            return [];
        }
        $selector = array_shift($selectors);

        $subs = $root->getShadowRoot()->findElements($selector);

        if (empty($selectors)) {
            return $subs;
        }

        $result = [];
        foreach($subs as $sub) {
            foreach ($this->resolveShadowSelector($sub, $selectors) as $fromSub) {
                $result[] = $fromSub;
            }
        }
        return $result;
    }
}

$shadowDom = $this->findElement('parent-web-component::shadow child-web-component')->getShadowRoot();
$shadowDom->findElement(WebDriverBy::cssSelector('#my-selector'));