#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import glob
import os.path
import sys
import time

sys.path.insert(0, './gen-py')
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
lib_path = glob.glob('../../lib/py/build/lib.*')
if lib_path:
    sys.path.insert(0, lib_path[0])
from optparse import OptionParser

from ThriftTest import ThriftTest, SecondService
from ThriftTest.ttypes import *
from thrift import TMultiplexedProcessor
from thrift.transport import TTransport
from thrift.transport import TSocket, TSSLSocket
from thrift.transport.THeaderTransport import THeaderTransport
from thrift.protocol import TBinaryProtocol
from thrift.protocol import THeaderProtocol
from thrift.server import TServer, TNonblockingServer, \
    TProcessPoolServer, THttpServer, TCppServer

class SecondHandler(SecondService.Iface):
    def blahBlah(self):
        print('blahBlah()')

class SecondContextHandler(SecondService.ContextIface):
    def __init__(self,):
        self.th = SecondHandler()

    def blahBlah(self, handler_ctx):
        self.th.blahBlah()

class TestHandler(ThriftTest.Iface):

    def testVoid(self):
        print('testVoid()')

    def testString(self, str):
        print('testString(%s)' % str)
        return str

    def testByte(self, byte):
        print('testByte(%d)' % byte)
        return byte

    def testI16(self, i16):
        print('testI16(%d)' % i16)
        return i16

    def testI32(self, i32):
        print('testI32(%d)' % i32)
        return i32

    def testI64(self, i64):
        print('testI64(%d)' % i64)
        return i64

    def testDouble(self, dub):
        print('testDouble(%f)' % dub)
        return dub

    def testFloat(self, flt):
        print('testFloat(%f)' % flt)
        return flt

    def testStruct(self, thing):
        print('testStruct({%s, %d, %d, %d})' % (thing.string_thing,
            thing.byte_thing, thing.i32_thing, thing.i64_thing))
        return thing

    def testException(self, str):
        print('testException(%s)' % str)
        if str == 'Xception':
            x = Xception()
            x.errorCode = 1001
            x.message = str
            raise x
        elif str == "throw_undeclared":
            raise ValueError("foo")

    def testOneway(self, seconds):
        print('testOneway(%d) => sleeping...' % seconds)
        time.sleep(seconds)
        print('done sleeping')

    def testNest(self, thing):
        return thing

    def testMap(self, thing):
        return thing

    def testSet(self, thing):
        return thing

    def testList(self, thing):
        return thing

    def testEnum(self, thing):
        return thing

    def testTypedef(self, thing):
        return thing

    def testRequestCount(self):
        return event_handler.request_count

    def testNewConnection(self):
        return event_handler.num_new_conns

    def testPreServe(self):
        return event_handler.num_pre_serve

    def testConnectionDestroyed(self):
        return event_handler.num_conns_destroyed


class TestContextHandler(ThriftTest.ContextIface):

    def __init__(self, server_port):
        self.th = TestHandler()
        self._server_port = server_port

    def testVoid(self, handler_ctx):
        self.th.testVoid()
        # This is here so we can check that handler_ctx is getting set,
        # without modifying the service definition which would require
        # modifying all the languages.
        if not (handler_ctx[0].endswith("127.0.0.1") or \
            handler_ctx[0].endswith("::1")) or \
                handler_ctx[1] == self._server_port:
            raise ValueError("handler_ctx not set properly " +
                    str(handler_ctx))

    def testString(self, handler_ctx, str):
        return self.th.testString(str)

    def testByte(self, handler_ctx, byte):
        return self.th.testByte(byte)

    def testI16(self, handler_ctx, i16):
        return self.th.testI16(i16)

    def testI32(self, handler_ctx, i32):
        return self.th.testI32(i32)

    def testI64(self, handler_ctx, i64):
        return self.th.testI64(i64)

    def testDouble(self, handler_ctx, dub):
        return self.th.testDouble(dub)

    def testFloat(self, handler_ctx, flt):
        return self.th.testFloat(flt)

    def testStruct(self, handler_ctx, thing):
        return self.th.testStruct(thing)

    def testException(self, handler_ctx, str):
        return self.th.testException(str)

    def testOneway(self, handler_ctx, seconds):
        return self.th.testOneway(seconds)

    def testNest(self, handler_ctx, thing):
        return self.th.testNest(thing)

    def testMap(self, handler_ctx, thing):
        return self.th.testMap(thing)

    def testSet(self, handler_ctx, thing):
        return self.th.testSet(thing)

    def testList(self, handler_ctx, thing):
        return self.th.testList(thing)

    def testEnum(self, handler_ctx, thing):
        return self.th.testEnum(thing)

    def testTypedef(self, handler_ctx, thing):
        return self.th.testTypedef(thing)

    def testRequestCount(self, handler_ctx):
        return self.th.testRequestCount()

    def testNewConnection(self, handler_ctx):
        return self.th.testNewConnection()

    def testPreServe(self, handler_ctx):
        return self.th.testPreServe()

    def testConnectionDestroyed(self, handler_ctx):
        return self.th.testConnectionDestroyed()


class ContextEventHandler(TProcessorEventHandler):
    def getHandlerContext(self, fn_name, server_context):
        # this is a tuple ("hostname", port)
        return server_context.getPeerName()

class HeaderEventHandler(ContextEventHandler):
    def getHandlerContext(self, fn_name, server_context):
        self.htrans = server_context.iprot.trans
        return ContextEventHandler.getHandlerContext(self,
                fn_name, server_context)

    def preWrite(self, handler_context, fn_name, result):
        for str_key, str_value in self.htrans.get_headers().items():
            # Just spit them back for testing.
            self.htrans.set_header(str_key, str_value)

class TestServerEventHandler(TServer.TServerEventHandler):

    def __init__(self):
        self.num_pre_serve = 0
        self.request_count = 0
        self.num_new_conns = 0
        self.num_conns_destroyed = 0

    def newConnection(self, context):
        self.num_new_conns += 1

    def preServe(self, address):
        self.num_pre_serve += 1

    def clientBegin(self, iprot, oprot):
        self.request_count += 1

    def connectionDestroyed(self, context):
        self.num_conns_destroyed += 1


if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option(
        "--ssl",
        action="store_true",
        dest="ssl",
        default=False,
        help="use SSL for encrypted transport")
    parser.add_option(
        "--multiple",
        action="store_true",
        dest="multiple",
        default=False,
        help="use multiple service")
    parser.add_option(
        "--header",
        action="store_true",
        dest="header",
        default=False,
        help="use the Header protocol")
    parser.add_option(
        "--context",
        action="store_true",
        dest="context",
        default=False,
        help="Use the context-passing Handler")
    parser.add_option(
        "--port",
        action="store",
        type="int",
        dest="port",
        default=9090)
    parser.add_option(
        "--timeout",
        action="store",
        type="int",
        dest="timeout",
        default=60)
    options, args = parser.parse_args()

    event_handler = TestServerEventHandler()

    if options.header:
        pfactory = THeaderProtocol.THeaderProtocolFactory(True,
                [THeaderTransport.HEADERS_CLIENT_TYPE,
                 THeaderTransport.FRAMED_DEPRECATED,
                 THeaderTransport.UNFRAMED_DEPRECATED,
                 THeaderTransport.HTTP_CLIENT_TYPE])
    else:
        pfactory = TBinaryProtocol.TBinaryProtocolFactory()

    if options.context:
        processor = ThriftTest.ContextProcessor(TestContextHandler(
                options.port))
    else:
        processor = ThriftTest.Processor(TestHandler())

    if options.multiple:
        processor = TMultiplexedProcessor.TMultiplexedProcessor()
        if options.context:
            processor.registerProcessor("ThriftTest",
                    ThriftTest.ContextProcessor(TestContextHandler(
                            options.port)))
            processor.registerProcessor("SecondService",
                    SecondService.ContextProcessor(SecondContextHandler()))
        else:
            processor.registerProcessor("ThriftTest",
                    ThriftTest.Processor(TestHandler()))
            processor.registerProcessor("SecondService",
                    SecondService.Processor(SecondHandler()))

    if args[0] == "THttpServer":
        try:
            server = THttpServer.THttpServer(
                    processor, ('', options.port), pfactory)
        except:
            print("Could not load THttpServer")
            raise
    elif args[0] == "TCppServer":
        server = TCppServer.TCppServer(processor)
        server.setPort(options.port)
    else:
        if options.ssl:
            cert_path = 'thrift/test/py/test_cert.pem'
            transport = TSSLSocket.TSSLServerSocket(options.port,
                    certfile=cert_path)
        else:
            transport = TSocket.TServerSocket(options.port)
        tfactory = TTransport.TBufferedTransportFactory()

        server_class_name = args[0]

        if args[0] == "TNonblockingServer":
            server = TNonblockingServer.TNonblockingServer(
                    processor, transport, pfactory,
                    readTimeout=options.timeout)
        elif args[0] == "TProcessPoolServer":
            server = TProcessPoolServer.TProcessPoolServer(
                    processor, transport, tfactory, pfactory)
        else:
            ServerClass = getattr(TServer, server_class_name)
            server = ServerClass(processor, transport, tfactory, pfactory)

    if options.header:
        server.processor.setEventHandler(HeaderEventHandler())
    elif options.context:
        server.processor.setEventHandler(ContextEventHandler())
    server.setServerEventHandler(event_handler)

    server.serve()
