from collections import deque
import psycopg2 as db


# The table and the function are created by the setup script `postgres_setup.py`
query_select = "SELECT a, b, c, d, e FROM RandomTable LIMIT 10"
query_sleep = "SELECT * FROM sleep(1)"


# These constants are defined in the WSGI environment but their value
# is know
WSGI_POLLIN = 0x01
WSGI_POLLOUT = 0x04


# Size of the connection pool
POOL_SIZE = 20

# Free connections available
free_connections = deque()

# Connections waiting for a free slot
waiting_requests = deque()

# Number of concurrent connections
connections = 0

# State to be kept between requests
request_state = {}



def get_connection(environ):
    global connections

    print 'open', connections, len(free_connections), len(waiting_requests)

    if free_connections:
        print 'reuse'
        # reuse existing connection
        dbconn, c = free_connections.pop()
    elif connections < POOL_SIZE:
        print 'new'
        # create a new connection
        dbconn = db.connect(database='test')

        curs = dbconn.cursor()
        # XXX bad API, fileno should be a property of the connection object
        fd = curs.fileno()
        c = environ['ngx.connection_wrapper'](fd)

        connections = connections + 1
    else:
        print 'wait'
        # no free slots, this request will have to wait
        ctx = environ['ngx.request_context']()
        waiting_requests.append(ctx)

        return None, None

    # XXX check me
    environ['ngx.poll_register'](c, WSGI_POLLIN)

    return dbconn, c

def close_connection(environ, dbconn, c):
    print 'close', connections, len(free_connections), len(waiting_requests)

    environ['ngx.poll_unregister'](c)

    # resume a waiting request, if any
    if waiting_requests:
        print 'waiting'
        ctx = waiting_requests.pop()
        try:
            ctx.resume()
        except RuntimeError, e:
            # XXX check me, the request has been finalized
            # (in this case we need to increment the timeout)
            print 'Error:', e
            # push the connection in the free connections
            free_connections.append((dbconn, c))

            return

        # make the connection directly available to the waiting
        # request
        request_state[ctx] = (dbconn, c)
    else:
        print 'no waiting'
        # push the connection in the free connections
        free_connections.append((dbconn, c))


def application(environ, start_response):
    global connections

    headers = [
        ('Server', 'Test'),
        ('Content-Type', 'text/plain'),
        ('X-Powered-By', 'Python'),
        ]

    qs = environ.get('QUERY_STRING', 'sleep')
    poll = environ['ngx.poll']

    dbconn, c = get_connection(environ)
    if dbconn is None:
        # no free connection, wait until one became available or a
        # timeout expires
        poll(50000)
        yield ''

        ctx = environ['ngx.request_context']()

        dbconn, c = request_state.get(ctx, (None, None))
        if dbconn is None:
            start_response('500 ERROR', headers)
            print 'query wait timedout'
            yield 'query wait timedout'
            return

        # XXX check me
        environ['ngx.poll_register'](c, WSGI_POLLIN)

        ##del request_state[ctx]

    curs = dbconn.cursor()

    if qs == 'sleep':
        query = query_sleep
    else:
        query = query_select

    # XXX error handling
    curs.execute(query, async=True)
    while not curs.isready():
        state = poll(5000)
        yield ''

        c_, flags = state()
        if c_ is None:
            start_response('500 ERROR', headers)
            print 'query timedout'
            yield 'query timedout'

            # XXX we just throw away the connection
            dbconn.close()
            c.close()
            connections = connections - 1

            return

    r = curs.fetchall()
    close_connection(environ, dbconn, c)

    if qs == 'sleep':
        data = 'empty return set'
    else:
        data = '\n'.join([' | '.join(['%5d' % col for col in row]) for row in r])

    start_response('200 OK', headers)
    yield data