#!/usr/bin/env python

import bs4
import os
import re
import sys
from requests import get

PYTHON_PROJECT_HOME = '/mnt/Storage/Lab/C++/competitive-programming/python'
CPP_PROJECT_HOME = '/mnt/Storage/Lab/C++/competitive-programming'
JAVA_PROJECT_HOME = '/mnt/Storage/Lab/Java/javaProject'


def extract_title(url, html):
    num = re.search(r'/(\d+)/', url).groups()[0]
    raw_title = html.find(attrs={'class': 'problem-statement'}).find(
        attrs={'class': 'title'}).text
    title = re.sub(r'\W+', '_', raw_title)
    return f'{num}{title}'


def copy(frompath, topath):
    with open(frompath, 'r') as fromfile, open(topath, 'w') as tofile:
        tofile.writelines(fromfile.readlines())
    print('Created code file at', topath)


def write_testcases(html, infile_path, outfile_path):
    def extract_text(tag):
        lines = []
        for line in tag.children:
            if isinstance(line, bs4.element.NavigableString):
                lines.append(line)
        return '\n'.join(lines)

    input_output_sections = [*html.find(attrs={'class': 'sample-test'})]
    input_tags = [*filter(lambda tag: tag.attrs and 'input' in tag.attrs['class'], input_output_sections)]
    output_tags = [*filter(lambda tag: tag.attrs and 'output' in tag.attrs['class'], input_output_sections)]
    assert len(input_tags) == len(output_tags)
    testcases = len(input_tags)

    with open(infile_path, 'w') as infile, open(outfile_path, 'w') as outfile:
        print(testcases, file=infile)
        for i, (input_tag, output_tag) in enumerate(zip(input_tags, output_tags)):
            input_text = extract_text([*input_tag][1])
            output_text = extract_text([*output_tag][1])

            print(file=infile)
            print(input_text, file=infile)

            print(f'Case #{i + 1}:', file=outfile)
            print(output_text, file=outfile)
            print(file=outfile)

    print(f'Written {testcases} test cases to \n'
          f'  {infile_path}, \n'
          f'  {outfile_path}')


def generate(url, html, language, project_home, template_relpath,
             code_path=None):
    complete_name = extract_title(url, html)

    if code_path is None:
        code_path = os.path.join(project_home, f'{complete_name}.{language}')

    template_path = os.path.join(project_home, template_relpath)
    infile_path = os.path.join(project_home, f'{complete_name}.in')
    outfile_path = os.path.join(project_home, f'{complete_name}.correct.out')

    copy(template_path, code_path)
    write_testcases(html, infile_path, outfile_path)

    return code_path, infile_path, outfile_path


def generate_cpp(url, html):
    code_path, infile_path, _ = generate(url=url, html=html, language='cpp',
                                         project_home=CPP_PROJECT_HOME,
                                         template_relpath='templates/codeforces.cpp')

    cmake_file = os.path.join(CPP_PROJECT_HOME, 'CMakeLists.txt')
    cpp_file_name = os.path.basename(code_path)
    replace_pattern(cmake_file,
                    r'set\(SOURCE_FILES \w+\.cpp\)',
                    f'set(SOURCE_FILES {cpp_file_name})')
    problem_name = re.findall(r'\/(\w+)\.cpp$', code_path)[0]
    replace_pattern(cmake_file,
                    r'add_definitions\(-DINPUT_PATH=".*?"\)',
                    f'add_definitions(-DINPUT_PATH="../{problem_name}.in")')


def replace_pattern(file_path, pattern, substitute):
    with open(file_path, 'r') as f:
        lines = f.readlines()
    with open(file_path, 'w') as f:
        for line in lines:
            if re.match(pattern, line):
                line = re.sub(pattern, substitute, line)
            f.write(line)


def generate_py(url, html):
    generate(url=url, html=html, language='py',
             project_home=PYTHON_PROJECT_HOME,
             template_relpath='templates/codeforces.py')


def generate_java(url, html):
    project_home = JAVA_PROJECT_HOME
    complete_name = extract_title(url, html)
    class_name = f'_{complete_name}.java'
    code_path = os.path.join(project_home, class_name)
    generate(url=url, html=html, language='java',
             project_home=project_home,
             template_relpath='Contests/Template.java',
             code_path=code_path)
    replace_pattern(code_path,
                    r'public class Template',
                    f'public class _{complete_name}')


def main(*argv):
    if len(argv) < 3:
        print('Usage: codeforces <url> (cpp|py|java)')
        sys.exit(1)

    url = argv[1]
    language = argv[2]

    print('Loading from', url, '...')
    response_text = get(url).content

    html = bs4.BeautifulSoup(response_text, 'html5lib')

    if language == 'cpp':
        generate_cpp(url, html)
    elif language == 'py':
        generate_py(url, html)
    elif language == 'java':
        generate_java(url, html)
    else:
        raise ValueError('Invalid language')


if __name__ == '__main__':
    main(*sys.argv)
