#! /usr/bin/python3
import argparse
import base64
import requests

RR_IP="192.168.250.1"
RR_IP_PORT="8080"
RAW_IMT_BASE_URL = "http://" + RR_IP +":" + RR_IP_PORT + "/api/public/imt"

def send_message(topic, message, skip_queue):
    if topic:
        payload = {
            "topicId": f"{topic}",
            "message": base64.b64encode(message.encode()).decode(),
            "skipQueue": f"{skip_queue}" if skip_queue else "0"
        }
    else:
        payload = {
            "message": base64.b64encode(message.encode()).decode(),
            "skipQueue": f"{skip_queue}" if skip_queue else "0"
        }
    response = requests.post(f"{RAW_IMT_BASE_URL}/mo", data=payload)
    if response.status_code == 200:
        data = response.json()
        if data['status']:
            print(f"Message ID: {data['messageId']}")
        else:
            for error in data['errors']:
                print(f"{error}")
    else:
        print(f"{response.text}")

def get_message_status(message_id):
    payload = {"messageId": message_id}
    response = requests.get(f"{RAW_IMT_BASE_URL}/mo", params=payload)
    if response.status_code == 200:
        data = response.json()
        if data['status']:
            status = data['imtStatus']
            date = status['at']['date']
            time = status['at']['time']
            print("Status:")
            print(f"  Topic ID: {status['topic']}")
            print(f"  State: {status['state']}")
            print(f"  Timestamp: {date['Year']}-{date['Month']}-{date['Day']} {time['Hour']}:{time['Minute']}:{time['Second']}")
        else:
            for error in data['errors']:
                print(f"{error}")
    else:
        print(f"{response.text}")

def get_message(topic=None):
    if topic:
        payload = {"topicId": topic} if topic else {}
        response = requests.get(f"{RAW_IMT_BASE_URL}/mt", params=payload)
    else:
        response = requests.get(f"{RAW_IMT_BASE_URL}/mt")

    if response.status_code == 200:
        data = response.json()
        if data['status']:
            date = data['at']['date']
            time = data['at']['time']
            payload = base64.b64decode(data['payload'])
            print(f"  Message ID: {data['messageId']}")
            print(f"  Timestamp: {date['Year']}-{date['Month']}-{date['Day']} {time['Hour']}:{time['Minute']}:{time['Second']}")
            try:
                ascii_data = payload.decode('ascii')
            except UnicodeDecodeError:
                # Not ASCII so just print base64
                print(f"  Payload: {data['payload']}")
            else:
                print(f"  Payload: {ascii_data}")
        else:
            for error in data['errors']:
                print(f"{error}")
    else:
        print(f"{response.text}")

def acknowledge_message(message_id):
    payload = {"messageId": message_id}
    response = requests.post(f"{RAW_IMT_BASE_URL}/mt/acknowledge", data=payload)
    if response.status_code == 200:
        data = response.json()
        if data['status']:
            print(f"Message ID: {message_id} has been acknowledged")
        else:
            for error in data['errors']:
                print(f"{error}")
    else:
        print(f"{response.text}")

def get_topics():
    response = requests.get(f"{RAW_IMT_BASE_URL}/topics")
    if response.status_code == 200:
        data = response.json()
        if data['status']:
            print("Topics:")
            for topic in data['topics']:
                print(f"  ID: {topic['id']}\n  Name: {topic['name']}")
        else:
            for error in data['errors']:
                print(f"{error}")
    else:
        print(f"{response.text}")

def send_command(command, topic_id=None, message_id=None):
    command_strings = ["UNKNOWN", "CLEAR_MT_QUEUE", "CLEAR_MO_QUEUE", "CANCEL_MO"]
    payload = {"command": command_strings[command]}
    if command in [1, 2]:
        payload["topicId"] = topic_id
    elif command == 3:
        payload["messageId"] = message_id

    response = requests.post(f"{RAW_IMT_BASE_URL}/command", data=payload)
    if response.status_code == 200:
        data = response.json()
        if data['status']:
            print("Command was successful")
        else:
            for error in data['errors']:
                print(f"{error}")
    else:
        print(f"{response.text}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='raw-imt', description="REST API client for RAW IMT.")
    parser.add_argument("action",
                        choices=["send", "status", "get", "ack", "topics", "cmd"],
                        help="Action to perform: send, status, get, ack, topics, or cmd")
    parser.add_argument("--topic",
                        type=int,
                        help="Topic ID")
    parser.add_argument("--message",
                        help="Message to send")
    parser.add_argument("--message-id",
                        help="Message ID")
    parser.add_argument("--skip-queue",
                        type=bool,
                        help="Is set when sending a message it is sent with high priority to skip the queue")
    parser.add_argument("--command",
                        type=int,
                        help="Command: 1 for CLEAR_MT_QUEUE, 2 for CLEAR_MO_QUEUE, 3 for CANCEL_MO")

    args = parser.parse_args()

    if args.action == "send":
        send_message(args.topic, args.message, args.skip_queue)
    elif args.action == "status":
        get_message_status(args.message_id)
    elif args.action == "get":
        get_message(args.topic)
    elif args.action == "ack":
        acknowledge_message(args.message_id)
    elif args.action == "topics":
        get_topics()
    elif args.action == "cmd":
        send_command(args.command, args.topic, args.message_id)