# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from collections import namedtuple
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Union
import itertools
import os
import re

import lit.Test

import stl.test.tests

_envlst_cache = dict()
_preprocessed_file_cache = dict()
_expected_result_entry_cache = dict()


@dataclass
class _TmpEnvEntry:
    env: Dict[str, str] = field(default_factory=dict)


@dataclass(frozen=True)
class EnvEntry:
    def __init__(self, tmp_env: _TmpEnvEntry):
        object.__setattr__(self, "_env_keys", tuple(tmp_env.env.keys()))
        object.__setattr__(self, "_env_vals", tuple(tmp_env.env.values()))

    def getEnvVal(self, key: str, default: Optional[str] = None) \
            -> Optional[str]:
        if key not in self._env_keys:
            return default

        # TRANSITION: All of this is to avoid having to install frozendict.
        # Reconsider this at a future date.
        return self._env_vals[self._env_keys.index(key)]

    _env_keys: Tuple[str]
    _env_vals: Tuple[str]


@dataclass
class _ParseCtx:
    current: List[_TmpEnvEntry] = field(default_factory=list)
    result: List[List[_TmpEnvEntry]] = field(default_factory=list)


_COMMENT_REGEX = re.compile(r"\s*#.*", re.DOTALL)
_INCLUDE_REGEX = re.compile(r'^RUNALL_INCLUDE (?P<filename>.+$)')
_ENV_VAR_MULTI_ITEM_REGEX = re.compile(r'(?P<name>\w+)="(?P<value>.*?)"')
_CROSSLIST_REGEX = re.compile(r'^RUNALL_CROSSLIST$')
_EXPECTED_RESULT_REGEX = re.compile(r'^(?P<prefix>.*) (?P<result>.*?)$')


def _parse_env_line(line: str) -> Optional[_TmpEnvEntry]:
    result = _TmpEnvEntry()
    for env_match in _ENV_VAR_MULTI_ITEM_REGEX.finditer(line):
        name = env_match.group("name")
        value = env_match.group("value")
        result.env[env_match.group("name")] = env_match.group("value")
    return result


def _append_env_entries(*args) -> _TmpEnvEntry:
    result = _TmpEnvEntry()
    for entry in args:
        for k, v in entry.env.items():
            if k not in result.env:
                result.env[k] = v
            else:
                result.env[k] = ' '.join((result.env[k], v))
    return result


def _do_crosslist(ctx: _ParseCtx):
    return itertools.starmap(_append_env_entries,
                             itertools.product(*ctx.result))


def _parse_env_lst(env_lst: Path, ctx: _ParseCtx):
    for line in parse_commented_file(env_lst):
        if (m:=_INCLUDE_REGEX.match(line)) is not None:
            p = env_lst.parent / Path(m.group("filename"))
            _parse_env_lst(p, ctx)
        elif _CROSSLIST_REGEX.match(line) is not None:
            ctx.result.append(ctx.current)
            ctx.current = []
        else:
            ctx.current.append(_parse_env_line(line))


def parse_commented_file(filename: Union[str, bytes, os.PathLike]) \
        -> List[str]:
    if str(filename) in _preprocessed_file_cache:
        return _preprocessed_file_cache[str(filename)]

    filename_path = Path(filename)
    result = list()
    with filename_path.open() as f:
        for line in f.readlines():
            if (line:=_COMMENT_REGEX.sub("", line)):
                line = line.strip()
                if line:
                    result.append(line)

        _preprocessed_file_cache[str(filename)] = result
        return result


def parse_result_file(filename: Union[str, bytes, os.PathLike]) \
        -> Dict[str, lit.Test.ResultCode]:
    if str(filename) in _expected_result_entry_cache:
        return _expected_result_entry_cache[str(filename)]

    res = dict()
    for line in parse_commented_file(filename):
        m = _EXPECTED_RESULT_REGEX.match(line)
        prefix = m.group("prefix")
        result = m.group("result")
        result_code = getattr(lit.Test, result, None)
        if result_code is None:
            result_code = getattr(stl.test.tests, result)
        res[prefix] = result_code

    _expected_result_entry_cache[str(filename)] = res
    return res


def parse_env_lst_file(env_list: Union[str, bytes, os.PathLike]) \
        -> Tuple[EnvEntry, ...]:
    if str(env_list) in _envlst_cache:
        return _envlst_cache[str(env_list)]

    env_list_path = Path(env_list)
    ctx = _ParseCtx()
    _parse_env_lst(env_list_path, ctx)
    ctx.result.append(ctx.current)
    res = tuple(map(EnvEntry, _do_crosslist(ctx)))
    _envlst_cache[env_list] = res
    return res

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from collections import namedtuple
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Union
import itertools
import os
import re

import lit.Test

import stl.test.tests

_envlst_cache = dict()
_preprocessed_file_cache = dict()
_expected_result_entry_cache = dict()


@dataclass
class _TmpEnvEntry:
    env: Dict[str, str] = field(default_factory=dict)


@dataclass(frozen=True)
class EnvEntry:
    def __init__(self, tmp_env: _TmpEnvEntry):
        object.__setattr__(self, "_env_keys", tuple(tmp_env.env.keys()))
        object.__setattr__(self, "_env_vals", tuple(tmp_env.env.values()))

    def getEnvVal(self, key: str, default: Optional[str] = None) \
            -> Optional[str]:
        if key not in self._env_keys:
            return default

        # TRANSITION: All of this is to avoid having to install frozendict.
        # Reconsider this at a future date.
        return self._env_vals[self._env_keys.index(key)]

    _env_keys: Tuple[str]
    _env_vals: Tuple[str]


@dataclass
class _ParseCtx:
    current: List[_TmpEnvEntry] = field(default_factory=list)
    result: List[List[_TmpEnvEntry]] = field(default_factory=list)


_COMMENT_REGEX = re.compile(r"\s*#.*", re.DOTALL)
_INCLUDE_REGEX = re.compile(r'^RUNALL_INCLUDE (?P<filename>.+$)')
_ENV_VAR_MULTI_ITEM_REGEX = re.compile(r'(?P<name>\w+)="(?P<value>.*?)"')
_CROSSLIST_REGEX = re.compile(r'^RUNALL_CROSSLIST$')
_EXPECTED_RESULT_REGEX = re.compile(r'^(?P<prefix>.*) (?P<result>.*?)$')


def _parse_env_line(line: str) -> Optional[_TmpEnvEntry]:
    result = _TmpEnvEntry()
    for env_match in _ENV_VAR_MULTI_ITEM_REGEX.finditer(line):
        name = env_match.group("name")
        value = env_match.group("value")
        result.env[env_match.group("name")] = env_match.group("value")
    return result


def _append_env_entries(*args) -> _TmpEnvEntry:
    result = _TmpEnvEntry()
    for entry in args:
        for k, v in entry.env.items():
            if k not in result.env:
                result.env[k] = v
            else:
                result.env[k] = ' '.join((result.env[k], v))
    return result


def _do_crosslist(ctx: _ParseCtx):
    return itertools.starmap(_append_env_entries,
                             itertools.product(*ctx.result))


def _parse_env_lst(env_lst: Path, ctx: _ParseCtx):
    for line in parse_commented_file(env_lst):
        if (m:=_INCLUDE_REGEX.match(line)) is not None:
            p = env_lst.parent / Path(m.group("filename"))
            _parse_env_lst(p, ctx)
        elif _CROSSLIST_REGEX.match(line) is not None:
            ctx.result.append(ctx.current)
            ctx.current = []
        else:
            ctx.current.append(_parse_env_line(line))


def parse_commented_file(filename: Union[str, bytes, os.PathLike]) \
        -> List[str]:
    if str(filename) in _preprocessed_file_cache:
        return _preprocessed_file_cache[str(filename)]

    filename_path = Path(filename)
    result = list()
    with filename_path.open() as f:
        for line in f.readlines():
            if (line:=_COMMENT_REGEX.sub("", line)):
                line = line.strip()
                if line:
                    result.append(line)

        _preprocessed_file_cache[str(filename)] = result
        return result


def parse_result_file(filename: Union[str, bytes, os.PathLike]) \
        -> Dict[str, lit.Test.ResultCode]:
    if str(filename) in _expected_result_entry_cache:
        return _expected_result_entry_cache[str(filename)]

    res = dict()
    for line in parse_commented_file(filename):
        m = _EXPECTED_RESULT_REGEX.match(line)
        prefix = m.group("prefix")
        result = m.group("result")
        result_code = getattr(lit.Test, result, None)
        if result_code is None:
            result_code = getattr(stl.test.tests, result)
        res[prefix] = result_code

    _expected_result_entry_cache[str(filename)] = res
    return res


def parse_env_lst_file(env_list: Union[str, bytes, os.PathLike]) \
        -> Tuple[EnvEntry, ...]:
    if str(env_list) in _envlst_cache:
        return _envlst_cache[str(env_list)]

    env_list_path = Path(env_list)
    ctx = _ParseCtx()
    _parse_env_lst(env_list_path, ctx)
    ctx.result.append(ctx.current)
    res = tuple(map(EnvEntry, _do_crosslist(ctx)))
    _envlst_cache[env_list] = res
    return res

# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # TRANSITION: All of this is to avoid having to install frozendict. # Reconsider this at a future date. #.*", re.DOTALL)