from awscrt import mqtt, http
from awsiot import mqtt_connection_builder
import sys
import threading
import time
import json
import csv
import os
import boto3
import uuid
from utils.command_line_utils import CommandLineUtils

# Parse command-line arguments
cmdData = CommandLineUtils.parse_sample_input_pubsub()

received_count = 0
received_all_event = threading.Event()

# 📌 CSV File Configuration
CSV_FOLDER = "/home/jacob/CSV Files/"  # Change this to your actual CSV folder

# 📌 DynamoDB Configuration
DYNAMODB_TABLE = "Raw_Data_SeniorDesign"
AWS_REGION = "us-east-1"  # Change this to your AWS region

# 📌 Initialize DynamoDB Client
dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION)
table = dynamodb.Table(DYNAMODB_TABLE)

# 📌 MQTT connection setup
mqtt_connection = mqtt_connection_builder.mtls_from_path(
    endpoint=cmdData.input_endpoint,
    port=cmdData.input_port,
    cert_filepath=cmdData.input_cert,
    pri_key_filepath=cmdData.input_key,
    ca_filepath=cmdData.input_ca,
    client_id=cmdData.input_clientId,
    clean_session=False,
    keep_alive_secs=30)

# 📌 Function to insert received message into DynamoDB
def insert_into_dynamodb(topic, message):
    try:
        # Convert DynamoDB message format to a normal dictionary
        item = {
            "key_1": str(uuid.uuid4()),  # Unique identifier
            "key_2": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),  # ISO timestamp
            "topic": topic,
        }

        # Extract key-value pairs from the message and store them as separate columns
        for key, value in message.items():
            # Handle cases where DynamoDB might store data as {'S': 'value'}
            if isinstance(value, dict) and 'S' in value:
                item[key] = value['S']  # Extract actual value
            else:
                item[key] = str(value)  # Store as string to avoid type mismatches

        # Insert the formatted item into DynamoDB
        table.put_item(Item=item)
        print("✅ Message stored in DynamoDB with structured columns")

    except Exception as e:
        print(f"❌ DynamoDB insert error: {e}")


# 📌 Callback for receiving MQTT messages
def on_message_received(topic, payload, dup, qos, retain, **kwargs):
    global received_count
    received_count += 1
    print(f"📩 Received message on topic '{topic}': {payload}")

    # Convert payload to JSON & store in DynamoDB
    try:
        message_data = json.loads(payload.decode("utf-8"))  # Decode from bytes
        insert_into_dynamodb(topic, message_data)
    except Exception as e:
        print(f"❌ Error processing message: {e}")

    if received_count == cmdData.input_count:
        received_all_event.set()

# 📌 Subscribe to MQTT topic
message_topic = cmdData.input_topic
print(f"🔌 Connecting to {cmdData.input_endpoint} with client ID '{cmdData.input_clientId}'...")
connect_future = mqtt_connection.connect()
connect_future.result()
print("✅ Connected to AWS IoT!")

print(f"📡 Subscribing to topic '{message_topic}'...")
subscribe_future, packet_id = mqtt_connection.subscribe(
    topic=message_topic,
    qos=mqtt.QoS.AT_LEAST_ONCE,
    callback=on_message_received)

subscribe_result = subscribe_future.result()
print("✅ Subscribed successfully!")

# 📌 Function to read CSV and publish messages
def process_csv_file(file_path):
    try:
        with open(file_path, "r") as file:
            reader = csv.DictReader(file)
            for row in reader:
                try:
                    payload = json.dumps(row)
                    print(f"📤 Publishing message to topic '{message_topic}': {payload}")
                    mqtt_connection.publish(
                        topic=message_topic,
                        payload=payload,
                        qos=mqtt.QoS.AT_LEAST_ONCE)
                    
                    time.sleep(1)  # 2-second delay between messages
                except Exception as e:
                    print(f"❌ Error sending row from {file_path}: {e}")
    except Exception as e:
        print(f"❌ Error reading file {file_path}: {e}")

# 📌 Process all CSV files in the folder
def process_all_csv_files():
    try:
        csv_files = [f for f in os.listdir(CSV_FOLDER) if f.endswith(".csv")]
        
        if not csv_files:
            print("⚠️ No CSV files found in the folder.")
            return
        
        for filename in csv_files:
            file_path = os.path.join(CSV_FOLDER, filename)
            print(f"📄 Processing file: {file_path}")
            process_csv_file(file_path)
    
    except Exception as e:
        print(f"❌ Error accessing folder {CSV_FOLDER}: {e}")

# 📌 Start Publishing CSV Data
print("🚀 Starting CSV data publishing...")
process_all_csv_files()

# 📌 Wait for incoming messages
print("📡 Waiting for incoming messages...")
received_all_event.wait()

# 📌 Disconnect from AWS IoT
print("🔌 Disconnecting from AWS IoT...")
disconnect_future = mqtt_connection.disconnect()
disconnect_future.result()
print("✅ Disconnected!")