test

#!/usr/bin/env python3

import getpass
import argparse
import os
import sys
import subprocess
import re
import tempfile
import traceback
from modules import atcoder, aoj, judge_result

def input_user_id():
  return input('User Id: ')

def input_password():
  return getpass.getpass()

class Atcoder:
  def __init__(self, config):
    assert 'contest_name' in config, 'missing config parameter: `contest_name`\nexample: `contest_name: arc055`'
    assert 'task_name' in config, 'missing config parameter: `task_name`\nexample: `task_name: arc055_c`'
    self.contest_name = config['contest_name']
    self.task_name = config['task_name']

  def submit(self, source):
    atcoder.submit(
      source_code=source,
      contest_name=self.contest_name,
      task_name=self.task_name,
    )
    return judge_result.Browser('https://beta.atcoder.jp/contests/{0}/submissions/me'.format(contest_name))

class Aoj:
  def __init__(self, config):
    assert 'problem_id' in config, 'missing config parameter: `problem_id`\nexample: `problem_id: 0005`'
    self.problem_id = config['problem_id']

  def submit(self, source):
    try:
      if not aoj.is_login():
        print('AOJ Session was expired. So you need to sign in again.')
        user_id = input_user_id()
        password = input_password()
        aoj.login(user_id=user_id, password=password)
        print('Logged in successfully!!')
      return aoj.submit(
        source_code=source,
        problem_id=self.problem_id,
      )
    except Exception as e:
      sys.stderr.write(str(e) + '\n')
      exit(1)

def parser():
  parser = argparse.ArgumentParser(description="Compile 'verify/<test-name>.<suffix>.cpp' and test it")
  parser.add_argument('-o', '--output-source', dest='output', type=str, help='only merge source and output it to this file')
  parser.add_argument('-l', '--local', dest='local', help='compile and run locally', action='store_true')
  parser.add_argument('-i', '--stdin', dest='stdin', help='force change input to stdin', action='store_true')
  parser.add_argument('test_name', metavar='test-name', type=str)
  parser.add_argument('suffix', type=str, nargs='?', default='test')
  return parser

class Executable:
  def __init__(self, file_path, delete=True):
    self.file_path = file_path
    self.delete = delete
    self.input_data = ''
  def input(self, input_data):
    self.input_data = input_data
    return self
  def run(self):
    command_str = self.file_path
    proc = subprocess.Popen(command_str, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    proc.stdin.write(self.input_data)
    proc.stdin.close()
    status = proc.wait()
    getattr(sys.stdout, 'buffer').write(proc.stdout.read())
    getattr(sys.stderr, 'buffer').write(proc.stderr.read())
    return status
  def __enter__(self):
    return self
  def __exit__(self, exception_type, exception_value, traceback):
    os.remove(self.file_path)

class Compiler:
  # command_str: '... [src] ... [dest] ...'
  def __init__(self, command_str):
    self.command_str = command_str
  def compile(self, src_file_path, dest_file_path):
    command_str = self.command_str \
      .replace('[src]', src_file_path) \
      .replace('[dest]', dest_file_path)
    proc = subprocess.Popen(command_str.strip().split(" "), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    status = proc.wait()
    getattr(sys.stdout, 'buffer').write(proc.stdout.read())
    getattr(sys.stderr, 'buffer').write(proc.stderr.read())
    if status != 0:
      exit(status)

  @staticmethod
  def get_gcc_compiler():
    return Compiler("g++ -o [dest] -O2 -fconstexpr-loop-limit=5000000 --std=c++14 [src]")
  @staticmethod
  def selectFromFname(ext):
    if ext == '.cpp':
      return Compiler.get_gcc_compiler()
    else:
      raise RuntimeError("no compiler for {0}".format(ext))

class Program:
  def __init__(self, compiler, source, config={}):
    self.compiler = compiler
    self.source = source
    self.tempfile = tempfile.NamedTemporaryFile(mode='w', suffix='.cpp', delete=False)
    self.config = config
  @staticmethod
  def load(file_path):
    src = ""
    include_reg = re.compile(r'^#include "([^"]+)"')
    config_start_reg = re.compile(r'.*==start==')
    config_parse_reg = re.compile(r'[ \t]*([^:]+): ?(.+)')
    config_end_reg = re.compile(r'.*==end==')
    is_config = False
    config = {}
    with open(file_path, encoding='utf-8') as f:
      for line in f:
        if config_start_reg.match(line):
          is_config = True
        elif config_end_reg.match(line):
          is_config = False
        elif is_config:
          mc = config_parse_reg.match(line)
          if mc:
            config[mc.group(1)] = mc.group(2)
        else:
          mc = include_reg.match(line)
          if mc:
            with open(os.path.join(os.path.dirname(file_path), mc.group(1))) as flib:
              lib_str = flib.read()
              src += include_reg.sub(lib_str, line)
          else:
            src += line
    return Program(Compiler.selectFromFname(os.path.splitext(file_path)[1]), src, config=config)

  def _outname(self):
    return self.tempfile.name + '.out'

  def write_to_temp(self):
    self.tempfile.write(self.source)
    self.tempfile.close()

  def compile(self):
    self.write_to_temp()
    self.compiler.compile(self.tempfile.name, self._outname())
    return Executable(self._outname())

  def save(self, target_file):
    with open(target_file, 'w') as fs:
      fs.write(self.source)

  def submit(self):
    assert 'judge' in self.config, 'write `judge: [atcoder/aoj]` between `==start==` and `==end==` line in the program'
    if self.config['judge'] == 'atcoder':
      judge = Atcoder
    elif self.config['judge'] == 'aoj':
      judge = Aoj
    return judge(self.config).submit(self.source)

  def __enter__(self):
    self.tempfile.__enter__()
    return self

  def __exit__(self, exception_type, exception_value, traceback):
    self.tempfile.__exit__(exception_type, exception_value, traceback)
    os.remove(self.tempfile.name)

def get_test_and_program_name(args):
  test_name = os.path.splitext(args.test_name)[0]
  program_name = 'verify/{0}.{1}.cpp'.format(test_name, args.suffix)
  return test_name, program_name

def write_log(title, sub_title, message):
  print("[{0}.{1}] {2}".format(title, sub_title, message))

def test_local(args):
  test_name, program_name = get_test_and_program_name(args)
  input_file_path = 'verify/{0}.{1}.in'.format(test_name, args.suffix)
  write_log(test_name, args.suffix, 'Loading...')
  status = 1
  with Program.load(program_name) as program:
    write_log(test_name, args.suffix, 'Compiling...')
    with program.compile() as executable:
      write_log(test_name, args.suffix, 'ReadingInput...')
      if not args.stdin and os.path.exists(input_file_path):
        with open(input_file_path, mode='r+b') as input_fs:
          input_data = input_fs.read()
      else:
        input_data = getattr(sys.stdin, 'buffer').read()
      write_log(test_name, args.suffix, 'Running')
      status = executable.input(input_data).run()
  exit(status)

def test_online(args):
  test_name, program_name = get_test_and_program_name(args)
  # write_log(test_name, args.suffix, 'Loading...')
  with Program.load(program_name) as program:
    # write_log(test_name, args.suffix, 'Judging...')
    result = program.submit()
    result.output()

def merge_only(args):
  _, program_name = get_test_and_program_name(args)
  with Program.load(program_name) as program:
    program.save(args.output)

def main():
  try:
    args = parser().parse_args()
    if args.output:
      merge_only(args)
    elif args.local:
      test_local(args)
    else:
      test_online(args)
  except Exception as e:
    sys.stderr.write(traceback.format_exc()+'\n')
    sys.stderr.write(str(e) + '\n')
    exit(1)

main()