#!/usr/bin/env python

"""scribe_pipe: A simple script for sending messages to scribe."""

import sys
from scribe import scribe
from thrift.transport import TTransport, TSocket
from thrift.protocol import TBinaryProtocol

result = 0

if len(sys.argv) == 2:
  category = sys.argv[1]
  host = '127.0.0.1'
  port = 1463
elif len(sys.argv) == 4 and sys.argv[1] == '-h':
  category = sys.argv[3]
  host_port = sys.argv[2].split(':')
  host = host_port[0]
  if len(host_port) > 1:
    port = int(host_port[1])
  else:
    port = 1463
else:
  sys.exit('usage (message is stdin): scribe_pipe [-h host[:port]] category')

socket = TSocket.TSocket(host=host, port=port)
transport = TTransport.TFramedTransport(socket)
protocol = TBinaryProtocol.TBinaryProtocol(trans=transport, strictRead=False, strictWrite=False)
client = scribe.Client(iprot=protocol, oprot=protocol)

transport.open()

try:
  while True:
    message = sys.stdin.readline()
    if not message:
      break
    log_entry = scribe.LogEntry(dict(category=category, message=message))
    result = client.Log(messages=[log_entry])
except KeyboardInterrupt:
  pass

transport.close()

if result == scribe.ResultCode.OK:
  sys.exit()
elif result == scribe.ResultCode.TRY_LATER:
  print >> sys.stderr, 'TRY_LATER'
  sys.exit(84)
else:
  sys.exit('Unknown error code (%s).' % result)