import logging, os, sys, socket, time, xmlrpclib, bz2, 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, e:
                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 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 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.SmartCookie()

    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 self.__cookies.has_key('MoodleSession'):
            return self.__cookies['MoodleSession'].value
        return None


    def __sendJsessionCookie(self, connection):
        if self.__cookies.has_key('MoodleSession'):
            connection.putheader(
                'Cookie',
                'MoodleSession=%s' % self.__cookies['MoodleSession'].value)
        if self.__cookies.has_key('MoodleSessionTest'):
            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, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                raise DataSourceError(e)
            except socket.error, 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, unicode)):
                        submit['code'] = submit['code'].__str__()
                return submits
            except xmlrpclib.Error, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                raise DataSourceError(e)
            except socket.error, 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, unicode)):
                        test['input'] = test['input'].__str__()
                    if not isinstance(test['output'], (str, unicode)):
                        test['output'] = test['output'].__str__()
                self.logger.debug('Got %d test case(s)', len(tests))
                return tests
            except xmlrpclib.Error, e:
                raise DataSourceError(e)
            except socket.error, 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, unicode)):
                    test['input'] = test['input'].__str__()
                if not isinstance(test['output'], (str, unicode)):
                    test['output'] = test['output'].__str__()
                test['input'] = bz2.decompress(test['input'])
                test['output'] = bz2.decompress(test['output'])
                return test
            except xmlrpclib.Error, e:
                raise DataSourceError(e)
            except socket.error, 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, unicode)):
                        code['code'] = code['code'].__str__()
                self.logger.debug('Got %d presetcodes', len(codes))
                return codes
            except xmlrpclib.Error, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                raise DataSourceError(e)
            except socket.error, 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, unicode)): r['stdout'] = ''
            if not isinstance(r['stderr'], (str, unicode)): 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, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                raise DataSourceError(e)
            except socket.error, 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, e:
                return DataSourceError(e)
            except socket.error, 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, unicode)):
                    return msg
                else:
                    return msg.__str__()
            except xmlrpclib.Error, e:
                return DataSourceError(e)
            except socket.error, 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, e:
                return DataSourceError(e)
            except socket.error, e:
                self.logger.exception('Failed to get_submit')
                time.sleep(self.config.retry_wait)

class DataSourceTest(unittest.TestCase):

    def setUp(self):
        execfile(os.path.join('..', 'testdata', 'test_config.py'))
        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: