#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-

from __future__ import print_function

import socket
import struct
import sys

PROTO_HEADER_FMT = '<HHHI'
PROTO_HEADER_SIZE = struct.calcsize(PROTO_HEADER_FMT)
REQ_BODY_FMT = '<L16s'
REQ_BODY_SIZE = struct.calcsize(REQ_BODY_FMT)
ACK_BODY_FMT = '<LL'
ACK_BODY_SIZE = struct.calcsize(ACK_BODY_FMT)

SEQ_NUM = 0
TIMEOUT_SEC = 3
HOST = '203.73.22.167'
PORT = 12201
CMD_ID = 3001


def recv_all(sock, size):
    chunks = []
    received = 0
    while received < size:
        chunk = sock.recv(size - received)
        if not chunk:
            raise RuntimeError('socket closed')
        chunks.append(chunk)
        received += len(chunk)
    return b''.join(chunks)


def pack_user_pwd(user_pwd):
    raw = user_pwd[:16]
    if len(raw) < 16:
        raw += '\x00' * (16 - len(raw))
    return raw


def send_request(host, port, cmd_id, user_id, user_pwd):
    body = struct.pack(
        REQ_BODY_FMT,
        user_id,
        pack_user_pwd(user_pwd))

    packet = struct.pack(
        PROTO_HEADER_FMT,
        PROTO_HEADER_SIZE + len(body),
        cmd_id,
        0,
        SEQ_NUM) + body

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(TIMEOUT_SEC)
    try:
        sock.connect((host, port))
        sock.sendall(packet)
        header_prefix = recv_all(sock, 2)
        pkg_len = struct.unpack('<H', header_prefix)[0]
        if pkg_len < PROTO_HEADER_SIZE:
            raise RuntimeError('invalid pkg_len=%u' % pkg_len)
        payload = recv_all(sock, pkg_len - 2)
        return header_prefix + payload
    finally:
        sock.close()


def unpack_header(response):
    return struct.unpack(PROTO_HEADER_FMT, response[:PROTO_HEADER_SIZE])


def main():
    if len(sys.argv) not in (2, 3):
        print('usage: python2.7 test_query_mb_balance.py user_id [user_pwd]')
        return 1

    user_id = int(sys.argv[1])
    user_pwd = ''
    if len(sys.argv) == 3:
        user_pwd = sys.argv[2]

    try:
        response = send_request(HOST, PORT, CMD_ID, user_id, user_pwd)
    except Exception as err:
        print('cmd=0x%04x' % CMD_ID)
        print('user_id=%u' % user_id)
        print('error=%s' % err)
        return 1

    pkg_len, cmd_id, status_code, seq_num = unpack_header(response)
    print('status_code=%u' % status_code)
    print('pkg_len=%u' % pkg_len)
    print('cmd_id=0x%04x' % cmd_id)
    print('seq_num=%u' % seq_num)

    if status_code == 0 and pkg_len == PROTO_HEADER_SIZE + ACK_BODY_SIZE:
        mb_balance, user_status = struct.unpack(
            ACK_BODY_FMT,
            response[PROTO_HEADER_SIZE:PROTO_HEADER_SIZE + ACK_BODY_SIZE])
        print('mb_balance=%u' % mb_balance)
        print('user_status=%u' % user_status)

    return 0 if status_code == 0 else 1


if __name__ == '__main__':
    sys.exit(main())