from contextlib import contextmanager
from pathlib import Path
from typing import Any, Dict, Iterator, List, Union, Optional, Callable, ContextManager
import os
import tempfile

from test_driver.logger import rootlog
from test_driver.machine import Machine, NixStartScript, retry
from test_driver.vlan import VLan
from test_driver.polling_condition import PollingCondition


def get_tmp_dir() -> Path:
    """Returns a temporary directory that is defined by TMPDIR, TEMP, TMP or CWD
    Raises an exception in case the retrieved temporary directory is not writeable
    See https://docs.python.org/3/library/tempfile.html#tempfile.gettempdir
    """
    tmp_dir = Path(tempfile.gettempdir())
    tmp_dir.mkdir(mode=0o700, exist_ok=True)
    if not tmp_dir.is_dir():
        raise NotADirectoryError(
            "The directory defined by TMPDIR, TEMP, TMP or CWD: {0} is not a directory".format(
                tmp_dir
            )
        )
    if not os.access(tmp_dir, os.W_OK):
        raise PermissionError(
            "The directory defined by TMPDIR, TEMP, TMP, or CWD: {0} is not writeable".format(
                tmp_dir
            )
        )
    return tmp_dir


class Driver:
    """A handle to the driver that sets up the environment
    and runs the tests"""

    tests: str
    vlans: List[VLan]
    machines: List[Machine]
    polling_conditions: List[PollingCondition]

    def __init__(
        self,
        start_scripts: List[str],
        vlans: List[int],
        tests: str,
        out_dir: Path,
        keep_vm_state: bool = False,
    ):
        self.tests = tests
        self.out_dir = out_dir

        tmp_dir = get_tmp_dir()

        with rootlog.nested("start all VLans"):
            self.vlans = [VLan(nr, tmp_dir) for nr in vlans]

        def cmd(scripts: List[str]) -> Iterator[NixStartScript]:
            for s in scripts:
                yield NixStartScript(s)

        self.polling_conditions = []

        self.machines = [
            Machine(
                start_command=cmd,
                keep_vm_state=keep_vm_state,
                name=cmd.machine_name,
                tmp_dir=tmp_dir,
                callbacks=[self.check_polling_conditions],
                out_dir=self.out_dir,
            )
            for cmd in cmd(start_scripts)
        ]

    def __enter__(self) -> "Driver":
        return self

    def __exit__(self, *_: Any) -> None:
        with rootlog.nested("cleanup"):
            for machine in self.machines:
                machine.release()

    def subtest(self, name: str) -> Iterator[None]:
        """Group logs under a given test name"""
        with rootlog.nested(name):
            try:
                yield
                return True
            except Exception as e:
                rootlog.error(f'Test "{name}" failed with error: "{e}"')
                raise e

    def test_symbols(self) -> Dict[str, Any]:
        @contextmanager