import logging
import json
import time
from datetime import datetime
from pymodbus.client import ModbusTcpClient
from pymodbus.exceptions import ModbusException

# Add MQTT client
import paho.mqtt.client as mqtt

# Import configuration
from config import (
    MODBUS_HOST, MODBUS_PORT, UNIT_ID,
    MQTT_BROKER, MQTT_PORT, MQTT_TOPIC, MQTT_CLIENT_ID, 
    MQTT_USERNAME, MQTT_PASSWORD, LOCATION, PUBLISH_INTERVAL
)

# Setting logging basic to see output
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# MQTT callbacks
def on_connect(client, userdata, flags, rc):
    if rc == 0:
        logging.info("Connected to MQTT Broker!")
    else:
        logging.error(f"Cannot connect to MQTT Broker. Return code: {rc}")

def on_publish(client, userdata, mid):
    logging.info(f"Successfully sent message with ID: {mid}")

def read_and_publish_data(mqtt_client, modbus_client):
    """Read data from Modbus and publish to MQTT"""
    try:
        # Check connection to Modbus server
        if not modbus_client.is_socket_open():
            if not modbus_client.connect():
                logging.error("Cannot connect to Modbus server.")
                return False

        # Read temperature (register 0)
        result_temp = modbus_client.read_holding_registers(address=0, count=1, slave=UNIT_ID)

        # Read humidity (register 1)
        result_hum = modbus_client.read_holding_registers(address=1, count=1, slave=UNIT_ID)

        # Initialize data to publish
        data = {
            "time": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
            "location": LOCATION
        }

        # Process temperature
        if hasattr(result_temp, 'registers') and result_temp.registers:
            raw_temp = result_temp.registers[0]
            raw_temperature = raw_temp * 0.1
            temperature = raw_temperature - 40
            logging.info(f"Raw temperature: {raw_temperature:.1f}°C (raw: {raw_temp})")
            logging.info(f"Corrected temperature: {temperature:.1f}°C")
            data["temperature"] = round(temperature, 1)
        else:
            logging.error(f"Error reading temperature: {result_temp}")
            return False

        # Xử lý độ ẩm
        if hasattr(result_hum, 'registers') and result_hum.registers:
            raw_hum = result_hum.registers[0]
            humidity = raw_hum * 0.1
            logging.info(f"Humidity: {humidity:.1f}%RH (raw: {raw_hum})")
            data["humidity"] = round(humidity, 1)
        else:
            logging.error(f"Error reading humidity: {result_hum}")
            return False

        # Convert data to JSON with a better format
        # indent=2 creates whitespace and newlines for JSON
        payload = json.dumps(data, indent=2)
        logging.info(f"Publishing data: {payload}")
        result = mqtt_client.publish(MQTT_TOPIC, payload)
        
        # Ensure data is sent
        result.wait_for_publish()
        
        if result.is_published():
            logging.info(f"Successfully published data to topic '{MQTT_TOPIC}'")
            return True
        else:
            logging.error("Cannot publish data")
            return False

    except Exception as e:
        logging.error(f"Error in reading/writing data: {e}", exc_info=True)
        return False

def main_loop():
    """Main function to connect and publish data in cycles"""
    # Initialize MQTT client
    mqtt_client = mqtt.Client(client_id=MQTT_CLIENT_ID)
    mqtt_client.on_connect = on_connect
    mqtt_client.on_publish = on_publish
    
    # Set username and password if needed
    if MQTT_USERNAME and MQTT_PASSWORD:
        mqtt_client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
    
    # Initialize Modbus TCP client
    modbus_client = ModbusTcpClient(
        host=MODBUS_HOST,
        port=MODBUS_PORT,
        timeout=30
    )
    
    try:
        # Connect to MQTT broker
        mqtt_client.connect(MQTT_BROKER, MQTT_PORT, 60)
        mqtt_client.loop_start()
        logging.info(f"Attempting to connect to Modbus TCP server: {MODBUS_HOST}:{MODBUS_PORT}...")
        
        # Connect to Modbus server
        if not modbus_client.connect():
            logging.error("Cannot connect to Modbus server initially.")
            return

        logging.info("Successfully connected to Modbus server.")
        logging.info(f"Starting data reading cycle every {PUBLISH_INTERVAL} seconds...")
        
        try:
            # Main loop to read and publish data
            while True:
                success = read_and_publish_data(mqtt_client, modbus_client)
                
                if not success:
                    logging.warning("Error occurred in current cycle, will retry in next cycle.")
                
                # Wait for next publish cycle
                logging.info(f"Waiting {PUBLISH_INTERVAL} seconds until next publish...")
                time.sleep(PUBLISH_INTERVAL)
                
        except KeyboardInterrupt:
            logging.info("Received stop signal from user.")
    
    except Exception as e:
        logging.error(f"Unexpected error in main loop: {e}", exc_info=True)
    
    finally:
        # Close Modbus connection
        modbus_client.close()
        logging.info("Successfully closed Modbus connection.")
        
        mqtt_client.loop_stop()
        mqtt_client.disconnect()
        logging.info("Successfully closed MQTT connection.")