#!/usr/bin/env python

import argparse
import enum
import json
import re
import shutil

from pathlib import Path
from collections import OrderedDict
from urllib.request import urlopen, URLError
from http import HTTPStatus


@enum.unique
class ProblemSpecsSource(enum.Enum):
    Local = 'local'
    Remote = 'remote'


canonical_data_url = 'https://raw.githubusercontent.com/exercism' + \
                     '/problem-specifications/master' + \
                     '/exercises/%s/canonical-data.json'

def smldir():
    """Absolute path of the sml directory."""
    return Path(__file__).absolute().parent.parent


class CanonicalDataNotFound(Exception):
    pass


def get_canonical_data_from_remote(exercise):
    try:
        with urlopen(canonical_data_url % exercise) as response:
            return json.loads(response.read().decode('utf8'))
    except URLError as e:
        raise CanonicalDataNotFound(f"{e.reason}")


def get_canonical_data_from_local(exercise):
    path = smldir().parent / 'problem-specifications' / 'exercises'

    try:
        with (path / exercise / 'canonical-data.json').open() as f:
            return json.loads(f.read())
    except OSError as e:
        raise CanonicalDataNotFound(f"{e}")


def get_canonical_data(exercise, problem_specs_source=ProblemSpecsSource.Remote):
    """Get canonical-data of the exercise as a json string."""
    if problem_specs_source == ProblemSpecsSource.Remote:
        return get_canonical_data_from_remote(exercise)
    elif problem_specs_source == ProblemSpecsSource.Local:
        return get_canonical_data_from_local(exercise)

    raise Exception(
        f"Unknown source for problem-specifications: '{problem_specs_source}'."
        f" Valid options are {list(ProblemSpecsSource)}."
    )

def indent(n, s):
    return s.rjust(len(s) + n, ' ')


def camelize(s):
    return re.sub('[-_](\w)', lambda x: x.group(1).upper(), s)


def sml_type(value, force=False):
    if force:
        return 'string'
    if value is None:
        return 'option'
    if isinstance(value, list):
        if not value:
            return "'a list"
        s = {sml_type(v)
             for v in value
             # }
             if not (isinstance(v, list) and not v)}
        s.discard('exn')
        if not s or len(s) > 2:
            raise Exception('Failed to infer a proper type')
        if len(s) == 1:
            return '%s list' % s.pop()
        if 'option' not in s:
            raise Exception('Failed to infer a proper type, got multiple types', value, s)
        s.remove('option')
        return '%s option list' % s.pop()
    if isinstance(value, dict):
        if 'error' in value:
            return 'exn'
        return '{%s}' % ', '.join(
            '%s: %s' % (key, sml_type(value[key]))
            for key in sorted(value.keys())
        )
    types = {int: 'int', str: 'string', float: 'real', bool: 'bool'}
    return types[type(value)]


def sml_value(value, smltype='', force=False):
    if force:
        return json.dumps(value)
    if value is None:
        return 'NONE'
    _value = None
    if isinstance(value, str):
        _value = json.dumps(value)
    if isinstance(value, bool):
        _value = str(value).lower()
    if isinstance(value, (int, float)):
        if value < 0:
            _value = '~%s' % -value
        else:
            _value = str(value)
    if isinstance(value, list):
        _value = '[%s]' % ', '.join(sml_value(v) for v in value)
    if isinstance(value, dict):
        if 'error' in value:
            return 'Fail "%s"' % value['error']
        _value = '{%s}' % ', '.join('%s = %s' % (k, sml_value(v))
                                  for k, v in value.items())
    return _value if 'option' not in smltype else ('SOME ' + _value)


def get_fn_args(case):
    case_input = case["input"]
    return OrderedDict(
        (k, case_input[k])
        for k in sorted(case_input.keys())
    )


def extract_signatures(data, force=False):
    funcs = OrderedDict()
    exercise = data['exercise']

    def type_aux(values):
        s = {sml_type(val, force=force) for val in values}
        if len(s) == 1:
            return s.pop()
        if len(s) == 2 and 'option' in s:
            s.remove('option')
            return '%s option' % s.pop()
        raise Exception('Failed to infer a proper type, can not unify different types: %s' % s)

    def signature(xs):
        output = []
        args = OrderedDict()
        for x in xs:
            output.append(x['output'])
            for arg, val in x['args'].items():
                args.setdefault(arg, []).append(val)
        return {
            'args': OrderedDict((arg, type_aux(vals)) for arg, vals in args.items()),
            'output': type_aux(output),
        }

    def collect(cases):
        for case in cases:
            if 'cases' in case:
                collect(case['cases'])
            else:
                # assumes there's at least one test with non "nullish" vars
                fn = camelize(case.get('property', exercise))
                if sml_type(case['expected'], force=force) == 'exn':
                    continue
                args = get_fn_args(case)
                if not all(args.values()) or sml_type(case['expected'], force=force) == "'a list":# or not case['expected']:
                    continue
                funcs.setdefault(fn, []).append({
                    'args': args,
                    'output': case['expected']
                })

    collect(data['cases'])
    return {fn: signature(funcs[fn]) for fn in funcs}


def expectation(signature, fn, args, expected, force=False):
    tmpl = '(fn _ => %s |> Expect.%s)'
    output = sml_value(expected, signature['output'], force=force)
    invocation_tmpl = '%s %s' if len(args) == 1 else '%s (%s)'
    invocation = invocation_tmpl % (
        fn,
        ', '.join((
            sml_value(val, signature['args'][arg], force=force)
            for arg, val in args.items())
        )
    )
    if sml_type(expected, force=force) == 'exn':
        return tmpl % (
            '(fn _ => %s)' % invocation,
            'error (%s)' % output
        )
    if signature['output'] == 'bool':
        return tmpl % (
            invocation,
            'truthy' if expected else 'falsy'
        )
    if signature['output'] == 'real':
        return tmpl % (
            invocation,
            'nearTo 0.001 %s' % output
        )
    return tmpl % (
        invocation,
        ('equalTo (%s)' if 'SOME ' in output else 'equalTo %s') % output
    )


def generate_test_content(data, signature, force=False):
    exercise = data['exercise']

    def fmt(case, depth=0):
        if 'property' not in  case:
            print('WARNING: property not found using exercise name as function name')
        fn = camelize(case.get('property', exercise))
        expected = case['expected']
        args = get_fn_args(case)
        return '\n'.join((
            indent(depth, 'test "%s"' % case.get('description', fn)),
            indent(depth + 2, expectation(signature[fn], fn, args, expected, force=force))
        ))

    def traverse(cases, depth=0):
        acc = []
        for case in cases:
            if 'cases' in case:
                acc.append('\n'.join((
                    indent(depth * 2, 'describe "%s" [' % case['description']),
                    ',\n\n'.join(traverse(case['cases'], depth + 1)),
                    indent(depth * 2, ']')
                )))
            else:
                acc.append(fmt(case, depth * 2))
        return acc

    return '\n'.join([
        'use "testlib.sml";',
        'use "%s.sml";' % data['exercise'],
        '',
        'infixr |>',
        'fun x |> f = f x',
        '',
        'val testsuite =',
        '  describe "%s" [' % data['exercise'],
        ',\n\n'.join(traverse(data['cases'], depth=2)),
        '  ]',
        '',
        'val _ = Test.run testsuite'
    ])


def generate_exercise_content(signatures):
    def fmt(args):
        return ', '.join('%s: %s' % (name, value)
                        for name, value in args.items())

    return '\n\n'.join(
        'fun {0} ({1}): {2} =\n'
        '  raise Fail "\'{0}\' is not implemented"'.format(
            fn,
            fmt(signature['args']),
            signature['output']
        )
        for fn, signature in signatures.items()
    )


def write(path, content):
    with path.open('w') as f:
        f.write(content)


# flags
TEST = 1
STUB = 2
EXAMPLE = 4
FLAGS = TEST | STUB | EXAMPLE


class EnumAction(argparse.Action):
    """Better enum support for argparse.

    Usage: If MyEnum is an enum.Enum, e.g.:

    @enum.unique
    class MyEnum(enum.Enum):
        Foo = 'foo'
        Bar = 'bar'

    you can use it with argparse like this:

    parser.add_argument(
        '--my-option',
        action=EnumAction,
        type=MyEnum,
        choices=list(MyEnum), # Can be omitted
        default=MyEnum.Foo,
    )

    On the command-line the *values* of MyEnum should be provided (i.e. 'foo' and 'bar' for MyEnum).
    Therefore it makes sense to use '@enum.unique'.
    """

    def __init__(self, **kwargs):
        enum_type = kwargs.pop("type", None)

        if enum_type is None or not issubclass(enum_type, enum.Enum):
            raise TypeError("The 'type' parameter must be an Enum when using EnumAction")

        # For convinience we deduce 'choices' if none are provided.
        kwargs.setdefault("choices", list(enum_type))

        for choice in kwargs["choices"]:
            if not isinstance(choice, enum_type):
                raise ValueError(f"Choices for '{kwargs['dest']}' must be of type {enum_type},"
                    + f"this choice is not: '{choice}'.")

        # Translate choices into strings:
        kwargs["choices"] = [e.value for e in kwargs["choices"]]

        super(EnumAction, self).__init__(**kwargs)

        self._enum = enum_type

    def __call__(self, parser: argparse.ArgumentParser, namespace: argparse.Namespace, value, option_string: str = None):
        # Convert value back into an Enum
        if isinstance(value, str):
            value = self._enum(value)
            setattr(namespace, self.dest, value)
        else:
            raise argparse.ArgumentTypeError()


def generate(exercise, flags, problem_specs_source, force=False):
    root = Path(__file__).parent.parent.absolute()
    path = root / 'exercises' / 'practice' / exercise

    if not path.exists():
        path.mkdir()

    data = get_canonical_data(exercise, problem_specs_source)
    signatures = extract_signatures(data, force=force)
    content = generate_exercise_content(signatures)

    if flags & TEST:
        write(path / 'test.sml', generate_test_content(data, signatures, force=force))

    if flags & STUB:
        write(path / ('%s.sml' % exercise), content)

    if flags & EXAMPLE:
        write(path / 'example.sml', content)

    shutil.copyfile(
        (root / 'lib/testlib.sml').as_posix(),
        (path / 'testlib.sml').as_posix()
    )


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('exercises', nargs='+')
    parser.add_argument(
        '--force',
        action='store_true',
        help=(
            'Type inference will be disabled and "string" will be assumed. ' +
            'Test cases will need to be modified to match the right data type.'
        )
    )
    parser.add_argument(
        '--test-only',
        action='store_true',
        help='Generate only "test.sml"'
    )
    parser.add_argument(
        '--stub-only',
        action='store_true',
        help='Generate only "<exercise>.sml"'
    )
    parser.add_argument(
        '--example-only',
        action='store_true',
        help='Generate only "example.sml"'
    )
    parser.add_argument(
        '--problem-specs-source',
        action=EnumAction,
        type=ProblemSpecsSource,
        default=ProblemSpecsSource.Remote,
        help='Choose whether to use remote (default) or another checkout of problem-specifications.'
    )
    args = parser.parse_args()

    for exercise in args.exercises:
        try:
            flags = 0
            if args.test_only:
                flags |= ~(STUB | EXAMPLE)
            if args.example_only:
                flags |= ~(STUB | TEST)
            if args.stub_only:
                flags |= ~(TEST | EXAMPLE)
            generate(exercise, force=args.force, problem_specs_source=args.problem_specs_source, flags=flags or FLAGS)
        except CanonicalDataNotFound as e:
            other_sources = [f"'{s.value}'" for s in ProblemSpecsSource if s != args.problem_specs_source]

            msg = f"Could not load canonical data from '{args.problem_specs_source.value}'-repo "\
                + f"of problem-specifications. Try to specify another source in '--problem-specs-source' "\
                + f"at command-line (e.g.: {', '.join(other_sources)}). Original Error: {e}"

            print(f"[{exercise}]: {msg}")
        except Exception as e:
            print(f"[{exercise}]: {e}")
