"""
Nexus MQTT Enigma Client

Python client for publishing encrypted MQTT messages to Nexus.
Uses AES-256-GCM with daily key rotation via HKDF.
"""

import base64
import json
import os
import time
from datetime import datetime, timezone
from typing import Any, Dict, Optional

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

try:
    import paho.mqtt.client as mqtt
except ImportError:
    raise ImportError("Please install paho-mqtt: pip install paho-mqtt")


class MQTTEnigmaClient:
    """
    MQTT client with Nexus Enigma encryption support.
    
    Encrypts payloads before publishing to MQTT broker.
    The Nexus MQTT subscriber will decrypt messages automatically.
    
    Example:
        client = MQTTEnigmaClient(
            host="localhost",
            port=1883,
            master_secret="BASE64_MASTER_SECRET",
            source_id=1
        )
        
        client.connect()
        client.publish("sensors/temperature", {"temp": 25.5})
        client.disconnect()
    """
    
    NONCE_LENGTH = 12  # 96 bits for GCM
    KEY_LENGTH = 32    # 256 bits for AES-256
    
    def __init__(
        self, 
        host: str, 
        port: int, 
        master_secret: str, 
        source_id: int,
        username: Optional[str] = None, 
        password: Optional[str] = None, 
        client_id: Optional[str] = None
    ):
        """
        Initialize the MQTT Enigma client.
        
        Args:
            host: MQTT broker hostname
            port: MQTT broker port
            master_secret: Base64-encoded master secret from Nexus
            source_id: MQTT source ID from Nexus
            username: Optional MQTT username
            password: Optional MQTT password
            client_id: Optional client ID (auto-generated if not provided)
        """
        self.host = host
        self.port = port
        self.master_secret = base64.b64decode(master_secret)
        self.source_id = source_id
        self.app_key = f"mqtt_source_{source_id}"
        
        # Generate client ID if not provided
        if client_id is None:
            client_id = f"nexus-mqtt-enigma-{source_id}-{os.getpid()}"
        
        # Initialize MQTT client
        self.client = mqtt.Client(client_id=client_id)
        
        if username:
            self.client.username_pw_set(username, password)
        
        # Connect callbacks
        self.client.on_connect = self._on_connect
        self.client.on_disconnect = self._on_disconnect
        
        self._connected = False
    
    def _on_connect(self, client, userdata, flags, rc):
        """Called when connected to MQTT broker."""
        if rc == 0:
            self._connected = True
        
    def _on_disconnect(self, client, userdata, rc):
        """Called when disconnected from MQTT broker."""
        self._connected = False
    
    def connect(self, timeout: int = 10) -> bool:
        """
        Connect to the MQTT broker.
        
        Args:
            timeout: Connection timeout in seconds
            
        Returns:
            True if connected successfully
        """
        try:
            self.client.connect(self.host, self.port)
            self.client.loop_start()
            
            # Wait for connection
            start = time.time()
            while not self._connected and (time.time() - start) < timeout:
                time.sleep(0.1)
            
            return self._connected
        except Exception as e:
            return False
    
    def disconnect(self):
        """Disconnect from the MQTT broker."""
        self.client.loop_stop()
        self.client.disconnect()
    
    @property
    def is_connected(self) -> bool:
        """Check if client is connected."""
        return self._connected
    
    def _derive_key(self, key_date: str) -> bytes:
        """
        Derive a daily encryption key using HKDF.
        
        Args:
            key_date: Date string in YYYY-MM-DD format
            
        Returns:
            32-byte derived key
        """
        info = f"nexus-enigma-{self.app_key}-{key_date}".encode('utf-8')
        
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=self.KEY_LENGTH,
            salt=None,
            info=info
        )
        
        return hkdf.derive(self.master_secret)
    
    def _encrypt(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Encrypt data using AES-256-GCM with daily key rotation.
        
        Args:
            data: Dictionary to encrypt
            
        Returns:
            Encrypted message structure
        """
        # Get today's date for key derivation
        key_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        
        # Derive daily key
        derived_key = self._derive_key(key_date)
        
        # Generate random nonce
        nonce = os.urandom(self.NONCE_LENGTH)
        
        # Encrypt the payload
        aesgcm = AESGCM(derived_key)
        plaintext = json.dumps(data).encode('utf-8')
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)
        
        # Return encrypted message structure
        return {
            "encrypted": True,
            "keyDate": key_date,
            "nonce": base64.b64encode(nonce).decode('utf-8'),
            "data": base64.b64encode(ciphertext).decode('utf-8')
        }
    
    def publish(
        self, 
        topic: str, 
        data: Dict[str, Any], 
        qos: int = 0, 
        retain: bool = False
    ) -> bool:
        """
        Publish an encrypted message to a topic.
        
        Args:
            topic: MQTT topic to publish to
            data: Dictionary data to encrypt and publish
            qos: Quality of Service level (0, 1, or 2)
            retain: Retain message on broker
            
        Returns:
            True if published successfully
        """
        if not self._connected:
            if not self.connect():
                return False
        
        # Encrypt the payload
        encrypted = self._encrypt(data)
        
        # Publish
        payload = json.dumps(encrypted)
        result = self.client.publish(topic, payload, qos=qos, retain=retain)
        
        return result.rc == mqtt.MQTT_ERR_SUCCESS


def create_mqtt_client(
    host: str,
    port: int,
    master_secret: str,
    source_id: int,
    username: Optional[str] = None,
    password: Optional[str] = None
) -> MQTTEnigmaClient:
    """Create a new Nexus MQTT Enigma client."""
    return MQTTEnigmaClient(
        host=host,
        port=port,
        master_secret=master_secret,
        source_id=source_id,
        username=username,
        password=password
    )
