# -- coding: utf-8 --
# http://blog.csdn.net/luhanglei
import picamera
import time
import traceback
import ctypes
from librtmp import *
 
global meta_packet
global start_time
 
 
class Writer(): # camera可以通過(guò)一個(gè)類(lèi)文件的對(duì)象來(lái)輸出，實(shí)現(xiàn)write方法即可
  conn = None # rtmp連接
  sps = None # 記錄sps幀，發(fā)過(guò)以后就不需要再發(fā)了（抓包看到ffmpeg是這樣的）
  pps = None # 同上
  sps_len = 0 # 同上
  pps_len = 0 # 同上
 
  time_stamp = 0
 
  def __init__(self, conn):
    self.conn = conn
 
  def write(self, data):
    try:
      # 尋找h264幀間隔符
      indexs = []
      index = 0
      data_len = len(data)
      while index < data_len - 3:
        if ord(data[index]) == 0x00 and ord(data[index + 1]) == 0x00 and ord(
            data[index + 2]) == 0x00 and ord(data[index + 3]) == 0x01:
          indexs.append(index)
          index = index + 3
        index = index + 1
      # 尋找h264幀間隔符 完成
      # 通過(guò)間隔符個(gè)數(shù)確定類(lèi)型，樹(shù)莓派攝像頭的第一幀是sps+pps同時(shí)發(fā)的
      if len(indexs) == 1: # 非sps pps幀
        buf = data[4: len(data)] # 裁掉原來(lái)的頭（00 00 00 01），把幀內(nèi)容拿出來(lái)
        buf_len = len(buf)
        type = ord(buf[0]) & 0x1f
        if type == 0x05: # 關(guān)鍵幀，根據(jù)wire shark抓包結(jié)果，需要拼裝sps pps 幀內(nèi)容 三部分，長(zhǎng)度都用4個(gè)字節(jié)表示
          body0 = 0x17
          data_body_array = [bytes(bytearray(
            [body0, 0x01, 0x00, 0x00, 0x00, (self.sps_len >> 24) & 0xff, (self.sps_len >> 16) & 0xff,
             (self.sps_len >> 8) & 0xff,
             self.sps_len & 0xff])), self.sps,
            bytes(bytearray(
              [(self.pps_len >> 24) & 0xff, (self.pps_len >> 16) & 0xff, (self.pps_len >> 8) & 0xff,
               self.pps_len & 0xff])),
            self.pps,
            bytes(bytearray(
              [(buf_len >> 24) & 0xff, (buf_len >> 16) & 0xff, (buf_len >> 8) & 0xff, (buf_len) & 0xff])),
            buf
          ]
          mbody = ''.join(data_body_array)
          time_stamp = 0 # 第一次發(fā)出的時(shí)候，發(fā)時(shí)間戳0，此后發(fā)真時(shí)間戳
          if self.time_stamp != 0:
            time_stamp = int((time.time() - start_time) * 1000)
          packet_body = RTMPPacket(type=PACKET_TYPE_VIDEO, format=PACKET_SIZE_LARGE, channel=0x06,
                       timestamp=time_stamp, body=mbody)
          packet_body.packet.m_nInfoField2 = 1
        else: # 非關(guān)鍵幀
          body0 = 0x27
          data_body_array = [bytes(bytearray(
            [body0, 0x01, 0x00, 0x00, 0x00, (buf_len >> 24) & 0xff, (buf_len >> 16) & 0xff,
             (buf_len >> 8) & 0xff,
             (buf_len) & 0xff])), buf]
          mbody = ''.join(data_body_array)
          # if (self.time_stamp == 0):
          self.time_stamp = int((time.time() - start_time) * 1000)
          packet_body = RTMPPacket(type=PACKET_TYPE_VIDEO, format=PACKET_SIZE_MEDIUM, channel=0x06,
                       timestamp=self.time_stamp, body=mbody)
        self.conn.send_packet(packet_body)
      elif len(indexs) == 2: # sps pps幀
        if self.sps is not None:
          return
        data_body_array = [bytes(bytearray([0x17, 0x00, 0x00, 0x00, 0x00, 0x01]))]
        sps = data[indexs[0] + 4: indexs[1]]
        sps_len = len(sps)
        pps = data[indexs[1] + 4: len(data)]
        pps_len = len(pps)
        self.sps = sps
        self.sps_len = sps_len
        self.pps = pps
        self.pps_len = pps_len
        data_body_array.append(sps[1:4])
        data_body_array.append(bytes(bytearray([0xff, 0xe1, (sps_len >> 8) & 0xff, sps_len & 0xff])))
        data_body_array.append(sps)
        data_body_array.append(bytes(bytearray([0x01, (pps_len >> 8) & 0xff, pps_len & 0xff])))
        data_body_array.append(pps)
        data_body = ''.join(data_body_array)
        body_packet = RTMPPacket(type=PACKET_TYPE_VIDEO, format=PACKET_SIZE_LARGE, channel=0x06,
                     timestamp=0, body=data_body)
        body_packet.packet.m_nInfoField2 = 1
 
        self.conn.send_packet(meta_packet, queue=True)
        self.conn.send_packet(body_packet, queue=True)
    except Exception, e:
      traceback.print_exc()
 
  def flush(self):
    pass
 
 
def get_property_string(string): # 返回兩字節(jié)string長(zhǎng)度及string
  length = len(string)
  return ''.join([chr((length >> 8) & 0xff), chr(length & 0xff), string])
 
 
def get_meta_string(string): # 按照meta packet要求格式返回bytes,帶02前綴
  return ''.join([chr(0x02), get_property_string(string)])
 
 
def get_meta_double(db):
  nums = [0x00]
  fp = ctypes.pointer(ctypes.c_double(db))
  cp = ctypes.cast(fp, ctypes.POINTER(ctypes.c_longlong))
  for i in range(7, -1, -1):
    nums.append((cp.contents.value >> (i * 8)) & 0xff)
  return ''.join(bytes(bytearray(nums)))
 
 
def get_meta_boolean(isTrue):
  nums = [0x01]
  if (isTrue):
    nums.append(0x01)
  else:
    nums.append(0x00)
  return ''.join(bytes(bytearray(nums)))
 
 
conn = RTMP(
  'rtmp://192.168.199.154/oflaDemo/test', # 推流地址
  live=True)
librtmp.RTMP_EnableWrite(conn.rtmp)
conn.connect()
start_time = time.time()
# 拼裝視頻格式的數(shù)據(jù)包
meta_body_array = [get_meta_string('@setDataFrame'), get_meta_string('onMetaData'),
          bytes(bytearray([0x08, 0x00, 0x00, 0x00, 0x06])), # 兩個(gè)字符串和ECMA array頭，共計(jì)6個(gè)元素,注釋掉了音頻相關(guān)數(shù)據(jù)
          get_property_string('width'), get_meta_double(640.0),
          get_property_string('height'), get_meta_double(480.0),
          get_property_string('videodatarate'), get_meta_double(0.0),
          get_property_string('framerate'), get_meta_double(25.0),
          get_property_string('videocodecid'), get_meta_double(7.0),
          # get_property_string('audiodatarate'), get_meta_double(125.0),
          # get_property_string('audiosamplerate'), get_meta_double(44100.0),
          # get_property_string('audiosamplesize'), get_meta_double(16.0),
          # get_property_string('stereo'), get_meta_boolean(True),
          # get_property_string('audiocodecid'), get_meta_double(10.0),
          get_property_string('encoder'), get_meta_string('Lavf57.56.101'),
          bytes(bytearray([0x00, 0x00, 0x09]))
          ]
meta_body = ''.join(meta_body_array)
print meta_body.encode('hex')
meta_packet = RTMPPacket(type=PACKET_TYPE_INFO, format=PACKET_SIZE_LARGE, channel=0x04,
             timestamp=0, body=meta_body)
meta_packet.packet.m_nInfoField2 = 1 # 修改stream id
stream = conn.create_stream(writeable=True)
with picamera.PiCamera() as camera:
  camera.start_preview()
  time.sleep(2)
  camera.start_recording(Writer(conn), format='h264', resize=(640, 480), intra_period=25,
              quality=25) # 開(kāi)始錄制，數(shù)據(jù)輸出到Writer的對(duì)象里
  while True:#永遠(yuǎn)不停止
    time.sleep(60)
  camera.stop_recording()
  camera.stop_preview()