from __future__ import print_function

import logging
import os
import sys
import traceback

import six
import yaml
from six import iteritems

from dmoj import judgeenv, executors
from dmoj.judge import Judge
from dmoj.judgeenv import get_supported_problems, get_problem_root
from dmoj.utils.ansi import ansi_style

all_executors = executors.executors


class TestManager(object):
    def output(self, message):
        print(message)

    def fail(self, message):
        self.output('\t\t' + message.replace('\r\n', '\n').replace('\n', '\r\n\t\t'))
        self.failed = True

    def set_expected(self, codes_all, codes_cases, feedback_all, feedback_cases):
        self.failed = False
        self.codes_all = codes_all
        self.codes_cases = codes_cases
        self.feedback_all = feedback_all
        self.feedback_cases = feedback_cases

    def _receive_packet(self, packet):
        pass

    def supported_problems_packet(self, problems):
        pass

    def test_case_status_packet(self, position, result):
        code = result.readable_codes()[0]
        if position in self.codes_cases:
            if code not in self.codes_cases[position]:
                self.fail('Unexpected code for case %d: %s, expecting %s' %
                          (position, code, ', '.join(self.codes_cases[position])))
        elif code not in self.codes_all:
            self.fail('Unexpected global code: %s, expecting %s' %
                      (code, ', '.join(self.codes_all)))

        feedback = self.feedback_all
        if position in self.feedback_cases:
            feedback = self.feedback_cases[position]
        if feedback is not None and result.feedback not in feedback:
            self.fail('Unexpected feedback: "%s", expected: "%s"' %
                      (result.feedback, '", "'.join(feedback)))

    def compile_error_packet(self, log):
        if 'CE' not in self.codes_all:
            self.fail('Unexpected compile error')

    def compile_message_packet(self, log):
        pass

    def internal_error_packet(self, message):
        self.fail('Unexpected internal error:\n' + message)

    def begin_grading_packet(self, is_pretested):
        pass

    def grading_end_packet(self):
        pass

    def batch_begin_packet(self):
        pass

    def batch_end_packet(self):
        pass

    def current_submission_packet(self):
        pass

    def submission_terminated_packet(self):
        pass

    def submission_acknowledged_packet(self, sub_id):
        pass


class TestJudge(Judge):
    def __init__(self, manager):
        super(TestJudge, self).__init__()
        self.packet_manager = manager


class Tester(object):
    all_codes = {'AC', 'IE', 'TLE', 'MLE', 'OLE', 'RTE', 'IR', 'WA', 'CE', 'SC'}

    def __init__(self, problem_regex=None, case_regex=None):
        self.manager = TestManager()
        self.manager.output = self.error_output
        self.judge = TestJudge(self.manager)
        self.sub_id = 0
        self.problem_regex = problem_regex
        self.case_regex = case_regex

        self.case_files = ['test.yml']
        if os.name == 'nt':
            self.case_files += ['test.windows.yml']
        elif os.name == 'posix':
            self.case_files += ['test.posix.yml']
            if 'freebsd' in sys.platform:
                self.case_files += ['test.freebsd.yml']
                if not sys.platform.startswith('freebsd'):
                    self.case_files += ['test.kfreebsd.yml']
            elif sys.platform.startswith('linux'):
                self.case_files += ['test.linux.yml']

    def output(self, message=''):
        print(message)

    def error_output(self, message):
        print(ansi_style('#ansi[%s](red)') % message)

    def test_all(self):
        total_fails = 0

        for problem, _ in get_supported_problems():
            if self.problem_regex is not None and not self.problem_regex.match(problem):
                continue
            root = get_problem_root(problem)
            test_dir = os.path.join(root, 'tests')
            if os.path.isdir(test_dir):
                fails = self.test_problem(problem, test_dir)
                if fails:
                    self.output(ansi_style('Problem #ansi[%s](cyan|bold) #ansi[failed %d case(s)](red|bold).') %
                                (problem, fails))
                else:
                    self.output(ansi_style('Problem #ansi[%s](cyan|bold) passed with flying colours.') % problem)
                self.output()
                total_fails += fails

        return total_fails

    def test_problem(self, problem, test_dir):
        self.output(ansi_style('Testing problem #ansi[%s](cyan|bold)...') % problem)
        fails = 0

        dirs = [case for case in os.listdir(test_dir) if self.case_regex is None or self.case_regex.match(case)]
        for i in range(len(dirs)):
            case = dirs[i]
            case_dir = os.path.join(test_dir, case)
            if os.path.isdir(case_dir):
                self.output(ansi_style('\tRunning test case #ansi[%s](yellow|bold) for #ansi[%s](cyan|bold)...')
                            % (case, problem))
                try:
                    case_fails = self.run_test_case(problem, case, case_dir)
                except Exception:
                    fails += 1
                    self.output(ansi_style('\t#ansi[Test case failed with exception:](red|bold)'))
                    self.output(traceback.format_exc())
                else:
                    self.output(ansi_style('\tResult of case #ansi[%s](yellow|bold) for #ansi[%s](cyan|bold): ')
                                % (case, problem) +
                                ansi_style(['#ansi[Failed](red|bold)', '#ansi[Success](green|bold)'][not case_fails]))
                    fails += case_fails

                if i != len(dirs) - 1:
                    self.output()

        return fails

    def run_test_case(self, problem, case, case_dir):
        config = {}
        for file in self.case_files:
            try:
                with open(os.path.join(case_dir, file)) as f:
                    config.update(yaml.safe_load(f.read()))
            except IOError:
                pass

        if not config:
            self.output(ansi_style('\t\t#ansi[Skipped](magenta|bold) - No usable test.yml'))
            return 0

        if 'skip' in config and config['skip']:
            self.output(ansi_style('\t\t#ansi[Skipped](magenta|bold) - Unsupported on current platform'))
            return 0

        language = config['language']
        if language not in all_executors:
            self.output(ansi_style('\t\t#ansi[Skipped](magenta|bold) - Language not supported'))
            return 0
        time = config['time']
        memory = config['memory']
        if isinstance(config['source'], six.string_types):
            with open(os.path.join(case_dir, config['source'])) as f:
                sources = [f.read()]
        else:
            sources = []
            for file in config['source']:
                with open(os.path.join(case_dir, file)) as f:
                    sources += [f.read()]
        codes_all, codes_cases = self.parse_expect(config.get('expect', 'AC'),
                                                   config.get('cases', {}),
                                                   self.parse_expected_codes)
        feedback_all, feedback_cases = self.parse_expect(config.get('feedback'),
                                                         config.get('feedback_cases', {}),
                                                         self.parse_feedback)

        def output_case(data):
            self.output('\t\t' + data.strip())

        fails = 0
        for source in sources:
            self.sub_id += 1
            self.manager.set_expected(codes_all, codes_cases, feedback_all, feedback_cases)
            self.judge.begin_grading(self.sub_id, problem, language, source, time, memory, False, False, blocking=True, report=output_case)
            fails += self.manager.failed
        return fails

    def parse_expect(self, all, cases, func):
        expect = func(all)
        if isinstance(cases, list):
            cases = enumerate(cases, 1)
        else:
            cases = iteritems(cases)
        case_expect = {id: func(codes) for id, codes in cases}
        return expect, case_expect

    def parse_expected_codes(self, codes):
        if codes == '*':
            return self.all_codes
        elif isinstance(codes, six.string_types):
            assert codes in self.all_codes
            return {codes}
        else:
            result = set(codes)
            assert not (result - self.all_codes)
            return result

    def parse_feedback(self, feedback):
        if feedback is None or feedback == '*':
            return None
        elif isinstance(feedback, six.string_types):
            return {feedback}
        else:
            return set(feedback)


def main():
    judgeenv.load_env(cli=True, testsuite=True)

    # Emulate ANSI colors with colorama
    if os.name == 'nt' and not judgeenv.no_ansi_emu:
        try:
            from colorama import init
            init()
        except ImportError:
            pass

    logging.basicConfig(filename=judgeenv.log_file, level=logging.INFO,
                        format='%(levelname)s %(asctime)s %(module)s %(message)s')

    executors.load_executors()

    tester = Tester(judgeenv.problem_regex, judgeenv.case_regex)
    fails = tester.test_all()
    print()
    print('Test complete')
    if fails:
        print(ansi_style('#ansi[A total of %d case(s) failed](red|bold).') % fails)
    else:
        print(ansi_style('#ansi[All cases passed.](green|bold)'))
    raise SystemExit(int(fails != 0))

if __name__ == '__main__':
    main()
