#!/usr/bin/env python

import argparse
import subprocess
import glob
import shutil
import os
import re
import requests
import textwrap
from bs4 import BeautifulSoup

problems_dir = '/usr/src/app/problems'

def main():
    choices = ['create', 'test', 'clean']
    parser = argparse.ArgumentParser(description='create a task / exec test / clean directory')
    parser.add_argument('command', choices=choices, help='command to execute')
    parser.add_argument('-t', '--task_name', help='target task name')
    parser.add_argument('-g', '--get_samples', help='to get task\'s sample (optional)', action='store_true')

    args = parser.parse_args()
    if args.command == 'create':
        create_project(args.task_name, args.get_samples)
    elif args.command == 'test':
        exec_test(args.task_name)
    elif args.command == 'clean':
        clean_dir()
    else:
        parser.print_help()


def create_project(task_name, is_get_samples):
    """
    プロジェクトのテンプレート作成
    """
    # バリデーション
    if task_name is None:
        print('ERROR: Please input task name.')
        exit(1)
    task_url = None
    if is_get_samples:
        contest_name = '-'.join(task_name.split('_')[:-1])
        task_url = 'https://atcoder.jp/contests/{}/tasks/{}'.format(contest_name, task_name)
    # 入出力例
    samples = get_samples(task_url)
    try:
        # ディレクトリ作成
        problem_dir_full_path = '{}/{}'.format(problems_dir, task_name)
        os.makedirs(problem_dir_full_path)
        # 問題とテストケースのファイルを作成
        with open('{}/main.py'.format(problem_dir_full_path), 'w') as file:
            code = textwrap.dedent('''\
                import sys
                input = sys.stdin.readline
                sys.setrecursionlimit(10 ** 9)

                def main():
                    pass

                if __name__ == '__main__':
                    main()
            ''')
            file.write(code)
        for i, sample in enumerate(samples):
            with open('{}/{}.in'.format(problem_dir_full_path, str(i + 1).zfill(2)), 'w') as file:
                file.write(sample['input'])
            with open('{}/{}.out'.format(problem_dir_full_path, str(i + 1).zfill(2)), 'w') as file:
                file.write(sample['output'])
        print('Task directory is created successfully.')
    except FileExistsError:
        print('ERROR: Target directory already exists.')


def get_samples(task_url):
    if task_url is None:
        return [{ 'input': '', 'output': '' }]
    html = None
    try:
        login_res = login()
        session = requests.session()
        res = session.get(task_url, cookies = login_res.cookies)
        res.raise_for_status()
        html = res.text
    except requests.exceptions.HTTPError:
        print('WARN: Error occured when scraping AtCoder task\'s page.')
        return [{ 'input': '', 'output': '' }]
    soup = BeautifulSoup(html, 'html.parser')
    parts = soup.find_all('div', class_='part')
    samples = []
    in_pattern = r'^入力例\s\d+$'
    out_pattern = r'^出力例\s\d+$'
    for part in parts:
        h3 = part.find('h3').string
        if re.match(in_pattern, h3):
            if len(samples) == 0 or samples[-1]['input'] is not None:
                samples.append({})
            pre = part.find('pre')
            content = pre.string.replace('\r', '')
            samples[-1]['input'] = content
        elif re.match(out_pattern, h3):
            pre = part.find('pre')
            content = ''
            if pre.string is not None:
                content = pre.string.replace('\r', '')
            samples[-1]['output'] = content
    return samples


def login():
    login_url = 'https://atcoder.jp/login'
    session = requests.session()
    response = session.get(login_url)
    bs = BeautifulSoup(response.text, 'html.parser')
    login_data = {
        'username': os.getenv('AC_USERNAME', 'username'),
        'password': os.getenv('AC_PASSWORD', 'password')
    }
    csrf_token = bs.find(attrs={'name': 'csrf_token'}).get('value')
    login_data['csrf_token'] = csrf_token
    return session.post(login_url, data = login_data, cookies = response.cookies)


def exec_test(task_name):
    """
    テストケースの実行
    """
    problem_dir_full_path = '{}/{}'.format(problems_dir, task_name)

    input_file_paths = sorted(glob.glob('{}/*.in'.format(problem_dir_full_path)))
    problem_file_path = glob.glob('{}/*.py'.format(problem_dir_full_path))[0]
    wrong_cases = []

    if len(input_file_paths) == 0:
        print('ERROR: Test cases are missing.')
        exit(1)

    for input_file_path in input_file_paths:
        command = 'cat {} | python {}'.format(input_file_path, problem_file_path)

        result = subprocess.check_output(command, shell=True)
        result_lines = result.decode('utf-8').strip().split("\n")

        output_file_path = input_file_path.replace('.in', '.out')

        answer = subprocess.check_output('cat {}'.format(output_file_path), shell=True)
        answer_lines = answer.decode('utf-8').strip().split("\n")

        is_correct = check(input_file_path, output_file_path, result_lines, answer_lines)
        if is_correct:
            print_color('.', pycolor.GREEN, newline=False)
        else:
            wrong_cases.append({
                'input_file_path': input_file_path,
                'results': result_lines,
                'answers': answer_lines
            })
            print_color('F', pycolor.RED, newline=False)

    print('')
    print('')
    if len(wrong_cases) == 0:
        print_color('All test cases are passed.', pycolor.GREEN)
    else:
        print_color('Failures:', pycolor.RED)
        for wc in wrong_cases:
            print_wrong_answer(wc['input_file_path'], wc['results'], wc['answers'])


def check(input_file_path, output_file_path, results, answers):
    """
    正解か判定
    """
    if len(results) != len(answers):
        return False

    for i in range(len(results)):
        if results[i] != answers[i]:
            return False

    return True


def print_wrong_answer(input_file_path, results, answers):
    """
    失敗したテストケースの表示
    """
    print('')
    print_color('  Test case: ' + input_file_path, pycolor.RED)
    print_color('    expected:', pycolor.RED)
    for answer in answers:
        print_color('      ' + answer, pycolor.RED)
    print_color('    got:', pycolor.RED)
    for result in results:
        print_color('      ' + result, pycolor.RED)


def print_color(msg, color, newline = True):
    """
    色付きのprint
    """
    end = '\n' if newline else ''
    print(color + msg + pycolor.END, end=end)


# ref. https://qiita.com/ironguy/items/8fb3ddadb3c4c986496d
class pycolor:
    BLACK = '\033[30m'
    RED = '\033[31m'
    GREEN = '\033[32m'
    YELLOW = '\033[33m'
    BLUE = '\033[34m'
    PURPLE = '\033[35m'
    CYAN = '\033[36m'
    WHITE = '\033[37m'
    END = '\033[0m'
    BOLD = '\038[1m'
    UNDERLINE = '\033[4m'
    INVISIBLE = '\033[08m'
    REVERCE = '\033[07m'


def clean_dir():
    """
    問いた問題をsolvedディレクトリに移行する
    """
    unsolved = '.unsolved'
    unsolved_dirs = []
    with open(problems_dir + '/' + unsolved) as f:
        for line in f:
            dir_name = line.strip()
            unsolved_dirs.append(dir_name)
    files = os.listdir(problems_dir)
    for f in files:
        if f.startswith('.'):
            continue
        if f == 'sample':
            continue
        if f == 'solved':
            continue
        if f in unsolved_dirs:
            continue
        shutil.move(problems_dir + '/' + f, problems_dir + '/solved/')
    print('Done.')


main()
