#!/usr/bin/python3
import requests
import urllib.request
import time
from bs4 import BeautifulSoup
import sys
import os
import argparse
from config import config

parser = argparse.ArgumentParser()

parser.add_argument(
    "problemid",
    help="The problem id on kattis.com, can be a portion of the problem id, just enough to be unique")
parser.add_argument("-l", "--language",
                    help="The language you wish to use, default=cpp",
                    default=config['default_language'],
                    choices=["cpp", "py", "py2"])


args = parser.parse_args()

language_to_extension = {
    'py' :'.py',
    'py2':'.py',
    'py3':'.py',
    'cpp':'.cpp'
}

language_to_hattis = {
    'py' :'python 3',
    'py2':'python 2',
    'py3':'python 3',
    'cpp':'C++',
}

language_to_template = {
    'cpp' : """#include <bits/stdc++.h>
using namespace std;

int main(){

  return 0;
}
""",
    'py' : """print('IMPLEMENT ME')
    """
}

def setup_testcases(response, problemid, problemfile, problemDir):
    soup = BeautifulSoup(response.text, "html.parser")
    pres = [s.find_all('pre') for s in soup.find_all('table', class_='sample')]
    tcs = []
    for x in pres:
        tcs.extend(x)
    tcs = [s.text.replace('<pre>', '') for s in tcs]
    tcs = [s.replace('</pre>', '') for s in tcs]

    tccount = 1

    for k in range(0, len(tcs), 2):
        with open(problemDir+'/in'+str(tccount)+".txt", 'w') as f:
            f.write(tcs[k])
        with open(problemDir+'/out'+str(tccount)+".txt", 'w') as f:
            f.write(tcs[k+1])
        tccount += 1
    with open(problemDir+'/user_in.txt', 'w') as f:
        f.write(tcs[0])
    with open(problemDir+'/user_out.txt', 'w') as f:
        f.write(tcs[1])

    tccount = 1
    hattis_lang = ""

    if args.language == "cpp":
        hattis_lang = language_to_hattis["cpp"]
        makeInit = "compile:\n\tg++ -std=c++17 -Wall -static -O2 {0} -o run.out\ntest:\n\tg++ -std=c++17 -Wall -static -O2 {0} -o run.out".format(
            problemId+".cpp")
        for k in range(0, len(tcs), 2):
            makeInit += "\n\t./run.out < in{}.txt > temp".format(tccount)
            makeInit += "\n\tdiff -y out{}.txt temp".format(tccount)
            tccount += 1

    elif args.language.startswith('py'):
        hattis_lang = language_to_hattis["py"]
        makeInit = "test:\n"
        for k in range(0, len(tcs), 2):
            makeInit += "\n\tpython {}.py  < in{}.txt > temp".format(problemId, tccount)
            makeInit += "\n\tdiff -y out{}.txt temp".format(tccount)
            tccount += 1

    makeInit += "\n\t./run.out < user_in.txt > temp".format(tccount)
    makeInit += "\n\tdiff -y user_out.txt temp".format(tccount)

    if(config["enable-hattis"]):
        makeInit += "\nsubmit:"
        for rc in config["kattisrcs"]:
            makeInit += "\n\thattis {} {} --conf {}".format(problemid, problemfile, rc)
            if hattis_lang:
                makeInit += " -l \"{}\"".format(hattis_lang)

    with open(problemDir+'/makefile', 'w') as f:
        f.write(makeInit)

def setup_description(response, problemId, problemDir):
    """
        This is a place holder function to be fixed
        when more nenn is found.
    """
    soup = BeautifulSoup(response.text, "html.parser")

    desc_text = """
    {}

    {}
    """.format(
        soup.select('.headline-wrapper')[0].get_text(),
        soup.select('.problembody')[0].get_text()
    )
    with open(problemDir+"/"+problemId+".txt", 'w') as f:
        f.write(desc_text)

if __name__ == '__main__':
    url = 'https://open.kattis.com/search?q='+args.problemid
    response = requests.get(url)
    if "search" in response.url:
        print("please write the exact problem id")
        sys.exit()

    problemId = response.url.split('/')[-1]
    problemDir = './'+problemId
    problemFile = problemId + language_to_extension[args.language]

    problemPath = problemDir+"/"+ problemFile

    if not os.path.exists(problemDir):
        os.mkdir(problemDir)
    if not os.path.exists(problemPath):
        with open(problemPath, 'w') as f:
            f.write(language_to_template[args.language])

    setup_testcases(response, problemId, problemFile, problemDir)
    setup_description(response, problemId, problemDir)

