#!/usr/bin/env python3

import socket
import json
from datetime import datetime, timedelta
import os
import threading
import time
import logging
from logging.handlers import RotatingFileHandler

# Define device type mappings
DEVICE_TYPES = {
    "10000002": "DoubleRoller",
    "02000001": "Host",
    "10000000": "SubSet",
    "10000001": "TDBU",
    "22000000": "WifiCurtain",
    "22000005": "WifiModule",
    "22000001": "WifiReceiver",
    "22000007": "WifiSwitch",
    "22000002": "WifiTubularMotor"
}

# UDP configuration
UDP_PORT_LISTEN = 32101
UDP_PORT_SEND = 32100
MULTICAST_GROUP = "238.0.0.18"
BUFFER_SIZE = 1024
CONFIG_DIR = os.path.expanduser("~/.config/KN_Connect")
CONFIG_FILE = os.path.join(CONFIG_DIR, "known_devices")
LOG_DIR = os.path.expanduser("~/.cache/KN_Connect")
LOG_FILE = os.path.join(LOG_DIR, "logfile.log")

# Logging configuration
os.makedirs(LOG_DIR, exist_ok=True)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        RotatingFileHandler(LOG_FILE, maxBytes=1024 * 1024, backupCount=7),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger()

# Device mapping
device_mapping = {}

# Ensure configuration directory exists
def ensure_config_dir():
    try:
        os.makedirs(CONFIG_DIR, exist_ok=True)
        logger.info(f"Ensured configuration directory exists: {CONFIG_DIR}")
    except Exception as e:
        logger.critical(f"Error ensuring configuration directory: {e}")

# Function to load known devices from config file
def load_known_devices():
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r") as file:
                known_devices = json.load(file)
                for mac, details in known_devices.items():
                    device_mapping[mac] = details
            logger.info("Loaded known devices from config file.")
        except Exception as e:
            logger.error(f"Error loading config file: {e}")

# Function to save known devices to config file
def save_known_devices():
    try:
        # Set all devices to offline before saving and include optional details
        for mac, details in device_mapping.items():
            details["online"] = False

        with open(CONFIG_FILE, "w") as file:
            json.dump(device_mapping, file, indent=4)
        logger.info("Saved known devices to config file.")
    except Exception as e:
        logger.error(f"Error saving config file: {e}")

# Start UDP listener
def start_udp_listener():
    logger.info(f"Starting UDP listener on port {UDP_PORT_LISTEN}")
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(("", UDP_PORT_LISTEN))

        # Join multicast group
        group = socket.inet_aton(MULTICAST_GROUP)
        mreq = group + socket.inet_aton("0.0.0.0")
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)

        while True:
            try:
                data, addr = sock.recvfrom(BUFFER_SIZE)
                message = json.loads(data.decode("utf-8"))

                # Check for unexpected data
                unexpected_keys = set(message.keys()) - {"msgType", "mac", "deviceType", "msgID", "token", "data", "AccessToken"}
                if unexpected_keys:
                    logger.warning(f"Unexpected keys in message: {unexpected_keys} from {addr}")

                if message.get("msgType") == "GetDeviceListAck":
                    process_get_device_list_ack(message)
                elif message.get("msgType") == "Heartbeat":
                    process_heartbeat(message)
                elif message.get("msgType") == "Report":
                    process_report(message)
                elif message.get("msgType") == "WriteDevice":
                    process_write_device(message)
                else:
                    logger.warning(f"Unknown message type received: {message} from {addr}")

            except Exception as e:
                logger.error(f"Error receiving or processing message: {e}")

# Function to send GetDeviceList message
def send_get_device_list():
    message = {
        "msgType": "GetDeviceList",
        "msgID": datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
    }
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(("", UDP_PORT_LISTEN))
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 2)
        sock.sendto(json.dumps(message, separators=(",", ":")).encode("utf-8"), (MULTICAST_GROUP, UDP_PORT_SEND))
    logger.info("Sent GetDeviceList message.")

# Function to handle received Heartbeat messages
def process_heartbeat(message):
    try:
        if message.get("msgType") != "Heartbeat":
            return

        mac = message.get("mac")
        device_type_code = message.get("deviceType")
        device_type = DEVICE_TYPES.get(device_type_code, "Unknown")
        token = message.get("token")
        data = message.get("data", {})

        # Extract optional parameters
        current_angle = data.get("currentAngle")
        current_position = data.get("currentPosition")
        rssi = data.get("RSSI")

        # Generate user-friendly name
        user_friendly_name = f"{mac}_{device_type}"

        # Add or update device mapping
        if mac not in device_mapping:
            device_mapping[mac] = {
                "name": user_friendly_name,
                "deviceType": device_type,
                "deviceTypeCode": device_type_code,
                "token": token,
                "fwVersion": None,
                "ProtocolVersion": None,
                "currentAngle": current_angle,
                "currentPosition": current_position,
                "RSSI": rssi,
                "online": True
            }
        else:
            existing_entry = device_mapping[mac]
            existing_entry.update({
                "currentAngle": current_angle if current_angle is not None else existing_entry.get("currentAngle"),
                "currentPosition": current_position if current_position is not None else existing_entry.get("currentPosition"),
                "RSSI": rssi if rssi is not None else existing_entry.get("RSSI"),
                "online": True
            })

        logger.info(f"Processed Heartbeat for {user_friendly_name}")

    except Exception as e:
        logger.error(f"Error processing Heartbeat: {e}")

# Function to handle GetDeviceListAck messages
def process_get_device_list_ack(message):
    try:
        if message.get("msgType") != "GetDeviceListAck":
            return

        mac = message.get("mac")
        device_type_code = message.get("deviceType")
        device_type = DEVICE_TYPES.get(device_type_code, "Unknown")
        token = message.get("token")
        fw_version = message.get("fwVersion")
        protocol_version = message.get("ProtocolVersion")

        # Add or update device mapping
        if mac not in device_mapping:
            device_mapping[mac] = {
                "name": f"{mac}_{device_type}",
                "deviceType": device_type,
                "deviceTypeCode": device_type_code,
                "token": token,
                "fwVersion": fw_version,
                "ProtocolVersion": protocol_version,
                "currentAngle": None,
                "currentPosition": None,
                "RSSI": None,
                "online": True
            }
        else:
            existing_entry = device_mapping[mac]
            existing_entry.update({
                "fwVersion": fw_version if fw_version is not None else existing_entry.get("fwVersion"),
                "ProtocolVersion": protocol_version if protocol_version is not None else existing_entry.get("ProtocolVersion"),
                "online": True
            })

        logger.info(f"Processed GetDeviceListAck for {mac}")

    except Exception as e:
        logger.error(f"Error processing GetDeviceListAck: {e}")

# Function to handle Report messages
def process_report(message):
    try:
        if message.get("msgType") != "Report":
            return

        mac = message.get("mac")
        data = message.get("data", {})

        # Extract optional parameters
        current_angle = data.get("currentAngle")
        current_position = data.get("currentPosition")
        rssi = data.get("RSSI")

        # Update device mapping
        if mac in device_mapping:
            existing_entry = device_mapping[mac]
            existing_entry.update({
                "currentAngle": current_angle if current_angle is not None else existing_entry.get("currentAngle"),
                "currentPosition": current_position if current_position is not None else existing_entry.get("currentPosition"),
                "RSSI": rssi if rssi is not None else existing_entry.get("RSSI"),
                "online": True
            })
            logger.info(f"Processed Report for {mac}")
        else:
            logger.warning(f"Report received for unknown device: {mac}")

    except Exception as e:
        logger.error(f"Error processing Report: {e}")

# Function to handle WriteDevice messages
def process_write_device(message):
    try:
        if message.get("msgType") != "WriteDevice":
            return

        mac = message.get("mac")
        device_type = message.get("deviceType")
        access_token = message.get("AccessToken")
        data = message.get("data", {})

        if "operation" not in data:
            logger.error(f"WriteDevice message missing 'operation': {message}")
            return

        operation = data.get("operation")
        if operation not in [0, 1, 2]:
            logger.error(f"Invalid operation value in WriteDevice message: {operation}")
            return

        if mac in device_mapping:
            logger.info(f"Triggering device {mac} of type {device_type} with operation: {operation}")
        else:
            logger.warning(f"WriteDevice command received for unknown device: {mac}")

    except Exception as e:
        logger.error(f"Error processing WriteDevice: {e}")

# Function to output mapping as JSON
def get_device_mapping_json():
    return json.dumps(device_mapping, indent=4)

if __name__ == "__main__":
    try:
        # Ensure config directory exists
        ensure_config_dir()

        # Load known devices
        load_known_devices()

        # Start the UDP listener first
        logger.info("Initializing listener...")
        listener_thread = threading.Thread(target=start_udp_listener, daemon=True)
        listener_thread.start()

        # Delay to ensure listener starts before sending GetDeviceList
        time.sleep(1)

        # Send GetDeviceList message
        send_get_device_list()

        # Keep the main thread alive
        listener_thread.join()
    except KeyboardInterrupt:
        logger.info("Shutting down UDP listener.")
        logger.info("Final device mapping:")
        logger.info(get_device_mapping_json())

        # Save known devices
        save_known_devices()