#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import os
import sys
import socket
import time
import xmlrpc.client as xmlrpclib
import bz2
import http.cookies as Cookie
import unittest
from engineconfig import getConfig
from entity import Submit, Problem, TestCase, PresetCode, DataFile


class DataSourceError(Exception):
    pass


class DataSource:

    def __init__(self, driver):
        self.driver = driver
        self.config = getConfig()
        self.logger = logging.getLogger('main')

        self.judge_id = self.driver.get_judge_id()

    def _get_config_retry(self):
        ret = self.config.retry
        if ret < 1:
            ret = 1
        return ret

    def _get_config_retry_wait(self):
        ret = self.config.retry_wait
        if ret < 0.1:
            ret = 0.1
        return ret

    def _do_action(self, func, args):
        retry = self._get_config_retry()
        retry_wait = self._get_config_retry_wait()
        ret = None
        while retry > 0:
            try:
                ret = func(*args)
            except DataSourceError:
                self.logger.exception("DataSourceError")
                if retry > 0:
                    time.sleep(retry_wait)
                    retry = retry - 1
                    continue
                else:
                    self.logger.info('Quit after retried.')
                    sys.exit(1)
            break
        return ret

    def get_submits(self, limit):
        func = self.driver.get_submits
        args = (self.judge_id, limit)
        rows = self._do_action(func, args)
        ret = []
        if rows:
            for row in rows:
                ret.append(Submit(self, row))
        return ret

    def reset_submits(self):
        func = self.driver.reset_submits
        args = (self.judge_id,)
        return self._do_action(func, args)

    def get_problem(self, pid):
        func = self.driver.get_problem
        args = (pid, )
        return Problem(self, self._do_action(func, args))

    def get_tests(self, pid, full):
        func = self.driver.get_tests
        args = (pid, full)
        rows = self._do_action(func, args)
        ret = []
        if rows:
            for row in rows:
                ret.append(TestCase(self, row))
        return ret

    def get_test(self, tid, raw=True):
        func = self.driver.get_test
        args = (tid, )
        row = self._do_action(func, args)
        if raw:
            return row
        else:
            return TestCase(self, row)

    def get_presetcodes(self, pid, lang):
        func = self.driver.get_presetcodes
        args = (pid, lang)
        rows = self._do_action(func, args)
        return list(map(lambda row: PresetCode(self, row), rows))

    def get_datafiles(self, pid):
        func = self.driver.get_datafiles
        args = (pid, )
        rows = self._do_action(func, args)
        return list(map(lambda row: DataFile(self, row), rows))

    def get_datafile_data(self, datafileid):
        func = self.driver.get_datafile_data
        args = (datafileid, )
        return self._do_action(func, args)

    def update_submits_status(self, sids, status):
        func = self.driver.update_submits_status
        if not isinstance(sids, (tuple, list)):
            sids = (sids, )
        args = (sids, status)
        self.logger.debug('set submits %s status to %s', sids.__str__(), status)
        return self._do_action(func, args)

    def update_submit_compilemessage(self, sid, msg):
        func = self.driver.update_submit_compilemessage
        args = (sid, msg)
        return self._do_action(func, args)

    def get_submit_status(self, sid):
        func = self.driver.get_submit_status
        args = (sid, )
        return self._do_action(func, args)

    def get_submit_compilemessage(self, sid):
        func = self.driver.get_submit_compilemessage
        args = (sid, )
        return self._do_action(func, args)

    def update_submit_test_results(self, sid, results):
        func = self.driver.update_submit_test_results
        args = (sid, results)
        self.logger.info('update_submits_status: submit %s, %d records' % \
           (sid, len(results)))
        return self._do_action(func, args)

    def get_submit(self, sid):
        func = self.driver.get_submit
        args = (sid, )
        ret = self._do_action(func, args)
        return Submit(self, ret)

class JspAuthTransport(xmlrpclib.Transport):
    def __init__(self):
        xmlrpclib.Transport.__init__(self)
        self.__cookies = Cookie.SimpleCookie()

    def request(self, host, handler, request_body, verbose=0):
        # issue XML-RPC request

        h = self.make_connection(host)
        if verbose:
            h.set_debuglevel(1)

        self.send_request(h, handler, request_body)
        self.__sendJsessionCookie(h)
        self.send_host(h, host)
        self.send_user_agent(h)
        self.send_content(h, request_body)

        errcode, errmsg, headers = h.getreply()

        if errcode != 200:
            raise xmlrpclib.ProtocolError(
                host + handler,
                errcode, errmsg,
                headers
                )
        self.verbose = verbose
        self.__processCookies(headers)
        return self.parse_response(h.getfile())


    def get_jsession_id(self):
        if 'MoodleSession' in self.__cookies:
            return self.__cookies['MoodleSession'].value
        return None


    def __sendJsessionCookie(self, connection):
        if 'MoodleSession' in self.__cookies:
            connection.putheader(
                'Cookie',
                'MoodleSession=%s' % self.__cookies['MoodleSession'].value)
        if 'MoodleSessionTest' in self.__cookies:
            connection.putheader(
                'Cookie',
                'MoodleSessionTest=%s' % self.__cookies['MoodleSessionTest'].value)


    def __processCookies(self, headers):
        if headers.getheader('Set-Cookie'):
            self.__cookies.load(headers.getheader('Set-Cookie'))


    def send_content(self, connection, request_body):
        connection.putheader("Content-Type", "text/xml")
        connection.putheader("Content-Length", str(len(request_body)))

        connection.endheaders()
        if request_body:
            connection.send(request_body)

class XmlRpcDataSource:

    def __init__(self, url):
        self.transport = JspAuthTransport()
        self.server = xmlrpclib.Server(url, transport = self.transport)
        self.config = getConfig()
        self.logger = logging.getLogger('main')

    def get_judge_id(self):
        while True:
            try:
                return self.server.oj.get_judge_id()
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_judge_id')
                time.sleep(self.config.retry_wait)

    def reset_submits(self, judgeid):
        while True:
            try:
                return self.server.oj.reset_submits(judgeid)
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to reset_submits')
                time.sleep(self.config.retry_wait)

    def get_submits(self, judgeid, limit):
        while True:
            try:
                submits = self.server.oj.get_submits(judgeid, limit)
                for submit in submits:
                    if not isinstance(submit['code'], str):
                        submit['code'] = submit['code'].__str__()
                return submits
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                time.sleep(self.config.retry_wait)
                self.logger.exception('Failed to get_submits')

    def get_problem(self, problemid):
        while True:
            try:
                return self.server.oj.get_problem(problemid)
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_problem')
                time.sleep(self.config.retry_wait)

    def get_tests(self, problemid, full):
        while True:
            try:
                tests = self.server.oj.get_tests(problemid, full)
                for test in tests:
                    if not isinstance(test['input'], str):
                        test['input'] = test['input'].__str__()
                    if not isinstance(test['output'], str):
                        test['output'] = test['output'].__str__()
                self.logger.debug('Got %d test case(s)', len(tests))
                return tests
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_tests')
                time.sleep(self.config.retry_wait)

    def get_test(self, testid):
        while True:
            try:
                test = self.server.oj.get_gztest(testid)
                if not isinstance(test['input'], str):
                    test['input'] = test['input'].__str__()
                if not isinstance(test['output'], str):
                    test['output'] = test['output'].__str__()
                test['input'] = bz2.decompress(test['input'])
                test['output'] = bz2.decompress(test['output'])
                return test
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_tests')
                time.sleep(self.config.retry_wait)

    def get_presetcodes(self, problemid, lang):
        while True:
            try:
                codes = self.server.oj.get_presetcodes(problemid, lang)
                for code in codes:
                    if not isinstance(code['code'], str):
                        code['code'] = code['code'].__str__()
                self.logger.debug('Got %d presetcodes', len(codes))
                return codes
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_presetcodes')
                time.sleep(self.config.retry_wait)
              
    def get_datafiles(self, problemid):
        while True:
            try:
                files = self.server.oj.get_datafiles(problemid)
                self.logger.debug('Got %d datafiles', len(files))
                return files
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_datafiles')
                time.sleep(self.config.retry_wait)

    def get_datafile_data(self, datafileid):
        while True:
            try:
                data = self.server.oj.get_datafile_data(datafileid)
                return str(data)
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_datafiles')
                time.sleep(self.config.retry_wait)

    def update_submit_compilemessage(self, id, compilemsg):
        compilemsg = xmlrpclib.Binary(compilemsg)
        while True:
            try:
                return self.server.oj.update_submit_compilemessage(
                    id, compilemsg)
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to update_submit_compilemessage')
                time.sleep(self.config.retry_wait)

    def update_submit_test_results(self, id, results):
        for r in results:
            if not isinstance(r['stdout'], str): r['stdout'] = ''
            if not isinstance(r['stderr'], str): r['stderr'] = ''
            r['stdout'] = xmlrpclib.Binary(r['stdout'])
            r['stderr'] = xmlrpclib.Binary(r['stderr'])
        while True:
            try:
                return self.server.oj.update_submit_test_results(id, results)
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to update_submit_compilemessage')
                time.sleep(self.config.retry_wait)

    def update_submits_status(self, submitids, newstatus):
        while True:
            mc = xmlrpclib.MultiCall(self.server)
            for id in submitids:
                mc.oj.update_submit_status(id, newstatus)
            try:
                return mc()
            except xmlrpclib.Error as e:
                raise DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to update_submits_status')
                time.sleep(self.config.retry_wait)

    def get_submit_status(self, sid):
        while True:
            try:
                return self.server.oj.get_submit_status(sid)
            except xmlrpclib.Error as e:
                return DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_submit_status')
                time.sleep(self.config.retry_wait)

    def get_submit_compilemessage(self, sid):
        while True:
            try:
                msg = self.server.oj.get_submit_compilemessage(sid)
                if isinstance(msg, str):
                    return msg
                else:
                    return msg.__str__()
            except xmlrpclib.Error as e:
                return DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_submit_status')
                time.sleep(self.config.retry_wait)

    def get_submit(self, sid):
        while True:
            try:
                return self.server.oj.get_submit(sid)
            except xmlrpclib.Error as e:
                return DataSourceError(e)
            except socket.error as e:
                self.logger.exception('Failed to get_submit')
                time.sleep(self.config.retry_wait)

class DataSourceTest(unittest.TestCase):

    def setUp(self):
        exec(open(os.path.join('..', 'testdata', 'test_config.py')).read())
        self.config = getConfig()
        self.datasource = self.config.datasources[0]
        self.dbname = os.path.join('..', 'testdata', self.config.testdb)


    def testGetSubmits(self):
        submits = self.datasource.get_submits(4)
        self.assertTrue(len(submits), 4)
        s = submits[0]
        self.assertEqual(s.id, '0000000001')
        self.assertEqual(s.problem_id, '0000000001')
        self.assertEqual(s.language, 'gcc-3.3')
        self.assertEqual(s.code, '')
        s = submits[1]
        self.assertEqual(s.code, '#include <stdio.h>')

    def testGetSubmit(self):
        s = self.datasource.get_submit(1)
        self.assertEqual(s.id, '0000000001')
        self.assertEqual(s.problem_id, '0000000001')
        self.assertEqual(s.language, 'gcc-3.3')
        self.assertEqual(s.code, '')

    def testResetSubmits(self):
        self.datasource.reset_submits()

    def testGetProblem(self):
        p = self.datasource.get_problem('0000000001')
        self.assertEqual(p.id, '0000000001')

    def testGetTests(self):
        tests = self.datasource.get_tests('0000000001', True)
        self.assertTrue(len(tests), 2)
        t = tests[0]
        self.assertEqual(t.id, '0000000001')
        self.assertEqual(t.problem_id, '0000000001')
        self.assertEqual(t.timemodified, 0)
        self.assertEqual(t.input, '1 2')
        self.assertEqual(t.output, '3')
        self.assertEqual(t.timelimit, 10)
        self.assertEqual(t.memlimit, 1024 * 1024)

    def testGetTest(self):
        t = self.datasource.get_test('0000000001', False)
        self.assertEqual(t.id, '0000000001')
        self.assertEqual(t.problem_id, '0000000001')
        self.assertEqual(t.timemodified, 0)
        self.assertEqual(t.input, '1 2')
        self.assertEqual(t.output, '3')
        self.assertEqual(t.timelimit, 10)
        self.assertEqual(t.memlimit, 1024 * 1024)

    def testUpdateSubmitCompileMessage(self):
        self.datasource.update_submit_compilemessage('0000000001', 'hello')
        msg = self.datasource.get_submit_compilemessage('0000000001')
        self.assertEqual(msg, 'hello')
        self.datasource.update_submit_compilemessage('0000000001', '')

    def testUpdateSubmitStatus(self):
        self.datasource.update_submits_status(('0000000004', '0000000001'),
                                              'compiling')
        s = self.datasource.get_submit_status('0000000001')
        self.assertEqual(s, 'compiling')
        s = self.datasource.get_submit_status('0000000004')
        self.assertEqual(s, 'compiling')
        s = self.datasource.get_submit_status('0000000002')
        self.assertEqual(s, 'new')
        self.datasource.update_submits_status(('0000000004', '0000000001'), 'new')
        s = self.datasource.get_submit_status('0000000001')
        self.assertEqual(s, 'new')

    def testUpdateSubmitTestResults(self):
        from pysqlite2 import dbapi2 as sqlite
        conn = sqlite.connect(self.dbname)
        cur = conn.cursor()
        cur.execute('DELETE FROM submit_test_result')
        conn.commit()
        cur.close()
        conn.close()

        id = '0000000001'
        r1 = { 'test_id' : '0000000001',
               'stdout' : '3',
               'stderr' : '',
               'exitcode' : 0,
               'signal' : 0,
               'timeused' : 0,
               'memused' : 0}
        r2 = { 'test_id' : '0000000002',
               'stdout' : '4',
               'stderr' : 'abc',
               'exitcode' : 0,
               'signal' : 11,
               'timeused' : 1,
               'memused' : 2}
        self.datasource.update_submit_test_results(id, (r1, r2))

        conn = sqlite.connect(self.dbname)
        cur = conn.cursor()
        cur.execute('SELECT * FROM submit_test_result')
        rows = cur.fetchall()
        self.assertEqual(len(rows), 2)
        cur.close()
        conn.close()
        row = rows[0]
        self.assertEqual(row[1], 1)   #sid
        self.assertEqual(row[2], 1)   #tid
        self.assertEqual(row[3], '3') #stdout
        self.assertEqual(row[4], '')  #stderr
        self.assertEqual(row[5], 0)   #exitcode
        self.assertEqual(row[6], 0)   #signal
        self.assertEqual(row[7], 0)   #timeused
        self.assertEqual(row[8], 0)   #memused
        row = rows[1]
        self.assertEqual(row[1], 1)   #sid
        self.assertEqual(row[2], 2)   #tid
        self.assertEqual(row[3], '4') #stdout
        self.assertEqual(row[4], 'abc')  #stderr
        self.assertEqual(row[5], 0)   #exitcode
        self.assertEqual(row[6], 11)   #signal
        self.assertEqual(row[7], 1)   #timeused
        self.assertEqual(row[8], 2)   #memused

if __name__ == '__main__':
    unittest.main()

# vim: set expandtab tabstop=4 shiftwidth=4: