#!/usr/bin/env python

import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.abspath(os.path.split(sys.argv[0])[0]), '..', 'gen-py'))
from thrift.transport.TTransport import *
from cStringIO import StringIO

import httplib
from boto.sqs.connection import SQSConnection
from boto.sqs.message import Message

class TSqsClient(TTransportBase):

  """Http implementation of TTransport base."""

  def __init__(self, accessKey, secretAccessKey, queueName):
    self.accessKey=accessKey
    self.secretAccessKey=secretAccessKey
    self.queueName = queueName
    self.inBuff = None
    self.outBuff = StringIO()
    self.conn=None
    self.q=None

  def open(self):
    print 'Opening'
    self.conn = SQSConnection(self.accessKey, self.secretAccessKey)
    self.q=self.conn.create_queue(self.queueName)

  def close(self):
    print 'Closing'
    if self.conn:
      self.conn.close()
      self.conn=None
      self.q=None
    
  def isOpen(self):
    return self.conn != None

  def read(self, sz):
    if not self.q:
      self.open()
    if self.inBuff==None:
      ms=self.q.get_messages(1)
      while len(ms)==0:
        time.sleep(1)
        ms=self.q.get_messages(1)

      m=ms[0]
      s=m.get_body()
      if len(s)<=sz:
        return s
      else:
        self.inBuff=StringIO()
        self.inBuff.write(s)
        self.inBuff.reset()
        return self.read(sz)
    else:
      s=self.inBuff.read(sz)
      if s=='':
        self.inBuff=None
        return self.read(sz)
      else:
        return s

  def write(self, buf):
    print 'Writing', len(buf)
    self.outBuff.write(buf)

  def flush(self):
    # Pull data out of buffer
    data = self.outBuff.getvalue()
    self.outBuff = StringIO()
    print 'Flushing', len(data)

    m=Message()
    m.set_body(data)
    self.q.write(m)

class TSqsServer(TServerTransportBase):
  def __init__(self, accessKey, secretAccessKey, queueName):
    self.accessKey=accessKey
    self.secretAccessKey=secretAccessKey
    self.queueName = queueName
    self.client=None
 
  def listen(self):
    pass

  def accept(self):
    self.client = TSqsClient(self.accessKey, self.secretAccessKey, self.queueName)
    return self.client

  def close(self):
    if self.client:
      self.client.close()
      self.client = None