from dataclasses import dataclass, field
from datetime import timedelta
from pathlib import Path
from typing import TypeVar
from zenlog import log

from numpy import nan

from mccode_antlr.reader import Registry

TestInstrExampleType = TypeVar('TestInstrExampleType', bound='TestInstrExample')


@dataclass
class TestInstrExample:
    """Holds parameters for a single instr file %Example: test case"""
    sourcefile: Path
    test_number: int
    parameter_values: str = ""
    detector: str = ""
    target_value: float = nan
    test_value: float = nan
    compiled: bool = False
    compile_time: timedelta = field(default_factory=timedelta)
    ran: bool = False
    run_time: timedelta = field(default_factory=timedelta)
    error_message: str = ""
    stdout: str = ""
    test_complete: bool = False

    def __post_init__(self):
        import re
        if '.instr' in self.parameter_values:
            rinstr = re.compile(r'[a-zA-Z][a-zA-Z0-9_\-]*\.instr')
            new_param = rinstr.sub('', self.parameter_values)
            log.info(f'Removing .instr from {self.parameter_values} now {new_param}')
            self.parameter_values = new_param
        if self.sourcefile.stem in self.parameter_values:
            new_param = self.parameter_values
            self.parameter_values = new_param.replace(self.sourcefile.stem, '')
            log.info(f'Removing {self.sourcefile.stem} from {new_param} now {self.parameter_values}')
        if '-n' in self.parameter_values:
            rn = re.compile(r'-n[0-9eE\.]*')
            new_param = rn.sub('', self.parameter_values)
            log.info(f'Removing -n from {self.parameter_values}, now {new_param}')
            self.parameter_values = new_param

    @classmethod
    def list_from_file(cls, filename: Path) -> list[TestInstrExampleType]:
        import re
        rspec = re.compile(f'%Example:([^\n]*)Detector:([^\n]*)_I=([0-9.+-e]+)')
        with open(filename, 'r') as file:
            contents = file.read()
            matches = rspec.findall(contents)
        if len(matches):
            groups = [(m[0].strip(), m[1].strip(), float(m[2].strip())) for m in matches]
            return [cls(filename, i+1, p, d, t) for i, (p, d, t) in enumerate(groups)]
        return [cls(filename, -1)]

    def get_display_name(self):
        return f'{self.sourcefile.stem}_{self.test_number}' if self.test_number > 1 else self.sourcefile.stem

    def get_json_repr(self):
        jr = dict(displayname=self.get_display_name(),
                  sourcefile=self.sourcefile.as_posix(),
                  testnb=self.test_number,
                  parvals=self.parameter_values,
                  detector=self.detector,
                  targetval=self.target_value,
                  testval=self.test_value,
                  compiled=self.compiled,
                  compiletime=str(self.compile_time),
                  didrun=self.ran,
                  runtime=str(self.run_time),
                  errmsg=self.error_message,
                  stdout=self.stdout,
                  complete=self.test_complete
                  )
        return jr

    def save(self, path: Path):
        import json
        with open(path.joinpath(self.get_display_name() + '.json'), 'w') as file:
            json.dump(self.get_json_repr(), file)


@dataclass
class LineLogger:
    lines: list = field(default_factory=list)

    def logline(self, line: str):
        self.lines.append(line)

    def save(self, filename: Path):
        with open(filename, 'w') as file:
            file.write('\n'.join(self.lines))
            file.write('\n')

    def find(self, search: str):
        import re
        for line in self.lines:
            if re.search(search, line):
                return True
        return False


def get_cpu_name():
    from platform import system
    from subprocess import check_output
    if system() == "Windows":
        from platform import processor
        return processor()
    elif system() == "Darwin":
        from os import environ, pathsep
        environ['PATH'] = environ['PATH'] + pathsep + '/usr/sbin'
        command = "sysctl -n machdep.cpu.brand_string"
        return check_output(command).strip()
    elif system() == "Linux":
        command = "cat /proc/cpuinfo"
        all_info = check_output(command, shell=True).decode().strip()
        for line in all_info.split("\n"):
            if "model name" in line:
                from re import sub
                return sub(".*model name.*:", "", line, 1).strip()
    return ""


def get_nvidia_gpu_name():
    """Identify GPU (Uses GPUtil which only knows about Nvidia GPUs)"""
    import GPUtil
    gpus = GPUtil.getGPUs()
    names = ','.join(gpu.name for gpu in gpus)
    return names if len(names) else "none"


def get_hostname():
    import socket
    return socket.gethostname()


def get_username():
    import os
    import pwd
    info = pwd.getpwuid(os.getuid())
    # Return the full name 'pw_gecos' if present, otherwise the username 'pw_name'
    return info[4] if len(info[4]) else info[0]


def _monitor_name_file_name_match(folder, monitor_name):
    import re
    look_for_filename = False
    with open(folder.joinpath('mccode_antlr.sim'), 'r') as file:
        lines = file.readlines()
    for line in lines:
        if re.match(rf"\s*component:\s*{monitor_name}$", line):
            look_for_filename = True
        if look_for_filename:
            m = re.match(r'\s*filename:\s+(.+)', line)
            if m:
                filename = folder.joinpath(m.group(1))
                if filename.is_file():
                    return filename
            if re.match('end data', line):
                # try the monitor named dat file
                zero_d_filename = folder.joinpath(monitor_name + '.dat')
                if zero_d_filename.is_file():
                    return zero_d_filename
                return None
    return None


def mccode_test_compiler(work_dir, file_path, target, registry, generator, dump, **kwargs):
    from pathlib import Path
    from mccode_antlr.reader import Reader
    from mccode_antlr.compiler.c import compile_instrument
    # only the provided (remote) registry should be necessary for test instruments
    reader = Reader(registries=[registry])
    inst = reader.get_instrument(file_path)
    output = Path(work_dir)
    config = dict(default_main=True, enable_trace=False, portable=False, include_runtime=True,
                  embed_instrument_file=False, verbose=False)
    try:
        binary_path = compile_instrument(inst, target, output, generator=generator, config=config, dump_source=dump, **kwargs)
    except RuntimeError as err:
        log.critical(f'Failed to produce a binary for {file_path}')
        log.info(err)
        return False, Path()
    return True, binary_path


def mcstas_test_compiler(target, work_dir, file_path, dump, **kwargs):
    from mccode_antlr.reader import MCSTAS_REGISTRY
    from mccode_antlr.translators.target import MCSTAS_GENERATOR
    return mccode_test_compiler(work_dir, file_path, target, MCSTAS_REGISTRY, MCSTAS_GENERATOR, dump, **kwargs)


def mcxtrace_test_compiler(target, work_dir, file_path, dump, **kwargs):
    from mccode_antlr.reader import MCXTRACE_REGISTRY
    from mccode_antlr.translators.target import MCXTRACE_GENERATOR
    return mccode_test_compiler(work_dir, file_path, target, MCXTRACE_REGISTRY, MCXTRACE_GENERATOR, dump, **kwargs)


def mccode_test_runner(target, binary_path, test_parameters: str, n_particles: int):
    import shutil
    from tempfile import mkdtemp
    from mccode_antlr.compiler.c import run_compiled_instrument

    output_path = Path(mkdtemp(prefix=binary_path.stem, dir=binary_path.parent.resolve()))
    # The McCode C runtime won't use an existing directory ...
    shutil.rmtree(output_path)
    # common default parameters: output directory, seed value, number of particles
    parameters = f'--dir {output_path} -s 1000 -n {n_particles} ' + test_parameters
    try:
        output = run_compiled_instrument(binary_path, target, parameters, capture=True)
    except RuntimeError as error:
        return False, str(error), output_path
    return True, output, output_path


def mcstas_test_runner(target, binary_path, test_parameters: str, n_particles: int):
    # If NeXus output is requested and the InstrumentDescriptionFile is needed, run a different script entirely...
    #   TODO think about actually doing this?
    # if target.flags & CBinaryTarget.Type.nexus and idf_required:
    #     run([config['idfgen'].get(str), str(binary), *options])
    return mccode_test_runner(target, binary_path, test_parameters, n_particles)


def mcxtrace_test_runner(target, binary_path, test_parameters: str, n_particles: int):
    return mccode_test_runner(target, binary_path, test_parameters, n_particles)


def mccode_test(compiler, runner, mpi: int = None, acc: bool = False, nexus: bool = False, **kwargs):
    """Specialize the compiler and runner based on requested number of processors and use of OpenACC, then test"""
    from functools import partial
    from mccode_antlr.compiler.c import CBinaryTarget
    target = CBinaryTarget(mpi=mpi is not None, acc=acc, count=1 if mpi is None else mpi, nexus=nexus)
    return _mccode_test(partial(compiler, target), partial(runner, target), **kwargs)


def mcstas_test(search_pattern=None, instr_count=None, skip_non_test: bool = False,
                mpi: int = None, acc: bool = False, nexus: bool = False, n_particles: int = 1000,
                dump: bool = False, workdir: Path = None):
    from mccode_antlr.reader import MCSTAS_REGISTRY as REGISTRY
    return mccode_test(mcstas_test_compiler, mcstas_test_runner, mpi, acc, nexus, registry=REGISTRY,
                       search_pattern=search_pattern, instr_count=instr_count, skip_non_test=skip_non_test,
                       n_particles=n_particles, dump=dump, workdir=workdir)


def mcxtrace_test(search_pattern=None, instr_count=None, skip_non_test: bool = False,
                  mpi: int = None, acc: bool = False, nexus: bool = False, n_particles: int = 1000,
                  dump: bool = False, workdir: Path = None):
    from mccode_antlr.reader import MCXTRACE_REGISTRY as REGISTRY
    return mccode_test(mcxtrace_test_compiler, mcxtrace_test_runner, mpi, acc, nexus, registry=REGISTRY,
                       search_pattern=search_pattern, instr_count=instr_count, skip_non_test=skip_non_test,
                       n_particles=n_particles, dump=dump, workdir=workdir)


def _mccode_test(compiler, runner, registry: Registry, search_pattern=None, instr_count=None,
                skip_non_test: bool = False, n_particles: int = 1000, workdir: Path = None, dump: bool=False):
    """Run McCode instrument test(s) using the provided `compiler` and `runner` functions to handle those tasks.

    Each takes the expected working directory name (a string) and the path to the instr file as input.
    The compiler should return only a boolean success indication.
    The runner should return a boolean success indicator, the contents of the run standard output, and the directory
    (if any) used to store run output files -- these files are intended to reside below the working directory and will
    be automatically removed by this function.

    Parameters:
        compiler: a function taking a string and Path and producing the necessary runtime executable
        runner: a function taking a string and Path and running the executable, storing its output
        registry: a Registry object which knows where to find all 'examples/*' instr files
        search_pattern: an optional positive filter for the registered 'examples/*' instr files
        instr_count: an optional positive count of filtered instr files to include
        skip_non_test: a flag to control whether instr files which do not specify test cases should be compiled and run
    """
    import re
    # import logging
    from zenlog import log as logging
    import tempfile
    from datetime import datetime
    logging.info(f"Finding test instruments in: {registry}")
    # we _require_ that all test instr files are in a folder called "examples" ...
    filenames = registry.match(re.compile('examples'))
    # allow the user to limit which test cases can be found via a regular expression search
    if search_pattern is not None:
        if not isinstance(search_pattern, re.Pattern):
            search_pattern = re.compile(search_pattern)
        filenames = [x for x in filenames if search_pattern.search(x) is not None]
    # the examples folder has at least one script in it:
    filenames = [x for x in filenames if Path(x).suffix == '.instr']
    # optionally test case count limiter:
    if instr_count is not None and isinstance(instr_count, int) and instr_count > 0:
        filenames = filenames[:instr_count]
    # grab full paths to each file (in case this is a remote registry and they need to be fetched)
    filepaths = [registry.path(filename) for filename in filenames]

    longest_name = max(len(x.stem) for x in filepaths)

    # Build test objects from the '%Example:' line(s) in each file:
    tests = [x for filepath in filepaths for x in TestInstrExample.list_from_file(filepath)]

    binaries = dict()

    # By default, use a temporary directory -- but allow the user to specify someplace else:
    created = False
    if workdir is None:
        created = True
        workdir = tempfile.mkdtemp()
    if not isinstance(workdir, Path):
        workdir = Path(workdir)
    if not workdir.exists():
        created = True
        workdir.mkdir(parents=True)
    if not workdir.is_dir():
        raise RuntimeError(f'Expected {workdir} to be a directory, but it is not.')

    # Move into a temporary directory to ensure compilation/run times are from scratch

    logging.info("Compiling instruments")
    for test in tests:
        print(test)
        if test.sourcefile not in binaries and (test.test_number != 0 or not skip_non_test):
            t1 = datetime.now()
            binaries[test.sourcefile] = compiler(workdir, test.sourcefile, dump=dump, recompile=True, replace=False)
            test.compile_time = datetime.now() - t1
            test.compiled = True

            if binaries[test.sourcefile][0]:
                logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: {test.compile_time}')
            else:
                logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: COMPILE ERROR')
        else:
            logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: COMPILE skipped')

    logging.info("Running tests")
    for test in tests:
        if not binaries.get(test.sourcefile, (False, None))[0]:
            logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: No compile')
            continue
        if test.test_number < 1:
            logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: No test')
            continue

        t1 = datetime.now()
        log.info(f'run {binaries[test.sourcefile][1]} -n {n_particles} {test.parameter_values}')
        test.ran, stdout, output_dir = runner(binaries[test.sourcefile][1], test.parameter_values, n_particles)
        test.run_time = datetime.now() - t1
        test.stdout = stdout.decode() if isinstance(stdout, bytes) else stdout

        if test.ran:
            logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: {test.run_time}')
        else:
            logging.info(f'%-{test.sourcefile.stem:>{longest_name}s}: RUNTIME ERROR')
            logging.debug(test.stdout)
            test.error_message = test.stdout
            continue

        detector_output = _monitor_name_file_name_match(output_dir, test.detector)
        logging.debug(f'Detector output stored in {detector_output}')
        test.test_value = - 1
        if detector_output is None or not detector_output.is_file():
            rdet = re.compile(f'{test.detector}_I= ([0-9+-eE.]+)')
            matches = rdet.findall(test.stdout)
            if len(matches):
                test.test_value = float(matches[0][0])
        else:
            with open(detector_output, 'r') as detector_file:
                detector_lines = detector_file.readlines()
            detector_pattern = re.compile('# values: ([0-9+-e.]+) ([0-9+-e.]+) ([0-9]+)')
            detector_matches = [detector_pattern.match(line) for line in detector_lines if detector_pattern.match(line)]
            if len(detector_matches):
                test.test_value = float(detector_matches[0].group(1))

        test.test_complete = True

    if created:
        import shutil
        shutil.rmtree(workdir)
    else:
        logging.info(f'Test compiled binaries and test case output directories are located under {workdir}.\n'
                     'You may wish to clean-up this directory to recover disk space.')

    # Since the output directories may no longer exist, save the test results ... here?
    for test in tests:
        test.save(Path())

    results = {t.get_display_name(): t.get_json_repr() for t in tests}
    # the 'ncount' and 'mpi' keys should be defined elsewhere
    results['_meta'] = {'date': datetime.now().isoformat(), 'hostname': get_hostname(),
                        'user': get_username(), 'cpu_type': get_cpu_name(), 'gpu_type': get_nvidia_gpu_name()}
    return results



def ispath(arg: str):
    from pathlib import Path
    return Path(arg)


def main(name, program):
    from argparse import ArgumentParser
    from zenlog import log
    parser = ArgumentParser(name, description=f'Test instrument compilation and runtime for {name}')
    parser.add_argument('-s', '--search', help='Regular expression positive filter for instrument names', default=None)
    parser.add_argument('-c', '--count', type=int, help='Maximum number of instruments to test', default=None)
    parser.add_argument('-k', '--skip', action='store_true', help='Skip instruments without test cases')
    parser.add_argument('--mpi', type=int, help='Number of mpi processes to use', default=None)
    parser.add_argument('--nexus', action='store_true', help='Whether to use NeXus output')
    parser.add_argument('-n', '--ncount', type=int, help='Number of particles to simulate per test', default=1000)
    parser.add_argument('-d','--dump', action='store_true', help='Output C source to file', default=False)
    parser.add_argument('-w','--workdir', type=ispath, help='Work directory, temporary if None', default=None)

    # logging.basicConfig(level=logging.DEBUG, format='{asctime} {levelname} {message}', style='{')
    # logging.basicConfig(level=logging.DEBUG)

    log.level('debug')

    args = parser.parse_args()
    results = program(search_pattern=args.search, instr_count=args.count, skip_non_test=args.skip, mpi=args.mpi,
                      nexus=args.nexus, n_particles=args.ncount, dump=args.dump, workdir=args.workdir)
    results['_meta']['mpi'] = args.mpi
    results['_meta']['ncount'] = args.ncount

def mcstas_test_main():
    return main('mcstas_test', mcstas_test)

def mcxtract_test_main():
    return main('mcxtrace_test', mcxtrace_test)


if __name__ == '__main__':
    mcstas_test_main()