using System;
using System.Collections.Generic;
using System.Linq;
using OpenQA.Selenium;

namespace OpenQA.Selenium.Extensions
{
    public static class DriverExtensions
    {
        public const int DefaultTimeout = 10;

        public static void Wait(this IWebDriver driver, int seconds = DefaultTimeout)
        {
            if (seconds <= 60)
                seconds *= 1000;

            System.Threading.Thread.Sleep(seconds);
        }

        public static IWebElement FindElement(this IWebDriver driver, By by, Func<IWebElement, bool> predicate)
        {
            return driver.FindElements(by, predicate).First();
        }

        public static IEnumerable<IWebElement> FindElements(this IWebDriver driver, By by, Func<IWebElement, bool> predicate)
        {
            return driver.FindElements(by).Where(predicate);
        }

        public static IWebElement WaitForElement(this IWebDriver driver, By by, Func<IWebElement, bool> predicate = null, int seconds = DefaultTimeout)
        {
            return driver.WaitForElements(by, predicate, seconds).First();
        }

        public static IEnumerable<IWebElement> WaitForElements(this IWebDriver driver, By by, Func<IWebElement, bool> predicate = null, int seconds = DefaultTimeout)
        {
            IEnumerable<IWebElement> els;
            var retry = 0;

            do
            {
                retry++;
                driver.Wait(1);

                els = driver.FindElements(by);
                if (predicate != null)
                    els = els.Where(predicate);

            } while (els.Count() == 0 && retry < seconds);

            return els;
        }

        public static FrameSwitcher SwitchToFrame(this IWebDriver driver, By by)
        {
            return new FrameSwitcher(driver, by);
        }

        public static IJavaScriptExecutor GetJavaScriptExecutor(this IWebDriver driver)
        {
            return driver as IJavaScriptExecutor;
        }
    }

    public class FrameSwitcher : IDisposable
    {
        private readonly IWebDriver _driver;

        private readonly string _window;

        public FrameSwitcher(IWebDriver driver, By by)
        {
            _driver = driver;
            _window = _driver.GetWindowHandle();

            var frameId = _driver.FindElement(by).GetAttribute("id");
            _driver.SwitchTo().Frame(frameId);
        }

        public void Dispose()
        {
            _driver.SwitchTo().Window(_window);
        }
    }
}