package com.yupi.project.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.yupi.project.config.properties.MqttProperties;
import com.yupi.project.model.dto.mqtt.RobotStatusMessage;
import com.yupi.project.model.enums.RobotStatusEnum;       // For robot's own status
import com.yupi.project.model.enums.RobotTaskStatusEnum;  // For task status from message
import com.yupi.project.service.RobotTaskService;
import com.yupi.project.service.ChargingRobotService; // Added import
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.paho.client.mqttv3.IMqttMessageListener;
import org.eclipse.paho.client.mqttv3.MqttMessage;
import org.springframework.stereotype.Component;
import org.apache.commons.lang3.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.Date;

@Slf4j
@Component
@RequiredArgsConstructor
public class MqttMessageHandler implements IMqttMessageListener {

    private final RobotTaskService robotTaskService;
    private final ChargingRobotService chargingRobotService; // Added final for RequiredArgsConstructor
    private final MqttProperties mqttProperties;
    private final ObjectMapper objectMapper; // For JSON parsing

    /**
     * This method is called when a message arrives from the server.
     *
     * @param topic   Name of the topic on the message was published to
     * @param message The actual message.
     * @throws Exception if a terminal error occurs, this call will be client driven
     *                   re-delivery processes rely on the client storing messages which is not implemented.
     */
    @Override
    public void messageArrived(String topic, MqttMessage message) throws Exception {
        String payload = new String(message.getPayload(), StandardCharsets.UTF_8);
        log.info("MQTT Message Arrived. Topic: {}, QoS: {}, Payload: {}", topic, message.getQos(), payload);

        String[] topicParts = topic.split("/");
        if (topicParts.length == 0) {
            log.error("Received message on invalid topic format: {}", topic);
            return;
        }
        String robotUID = topicParts[topicParts.length - 1]; 

        if (topic.startsWith(mqttProperties.getStatusTopicBase())) {
            handleRobotStatusUpdate(robotUID, payload);
        } else {
            log.warn("Received message on unhandled topic structure: {}. Ignoring.", topic);
        }
    }

    private void handleRobotStatusUpdate(String robotUIDFromTopicSource, String payloadJson) {
        try {
            RobotStatusMessage statusMessage = objectMapper.readValue(payloadJson, RobotStatusMessage.class);
            log.info("Parsed RobotStatusMessage for robot {}: {}", robotUIDFromTopicSource, statusMessage);

            // Scenario 1: Message is related to a specific task (has taskId)
            if (statusMessage.getTaskId() != null) {
                log.debug("Handling task-specific status update for task ID: {} from robot {}", statusMessage.getTaskId(), robotUIDFromTopicSource);
                RobotTaskStatusEnum taskNewStatus = RobotTaskStatusEnum.fromValue(statusMessage.getStatus());
                if (taskNewStatus == null) {
                    log.error("Invalid task status '{}' received for task {} from robot {}. Payload: {}", 
                            statusMessage.getStatus(), statusMessage.getTaskId(), robotUIDFromTopicSource, payloadJson);
                    // Do not return, general status might still need processing.
                } else {
                boolean taskUpdated; // Defined outside switch
                switch (taskNewStatus) {
                    case PROCESSING: 
                        taskUpdated = robotTaskService.markTaskAsProcessing(statusMessage.getTaskId(), new Date(), statusMessage.getMessage());
                        break;
                    case COMPLETED:
                        taskUpdated = robotTaskService.markTaskAsCompleted(statusMessage.getTaskId(), new Date(), statusMessage.getMessage());
                        break;
                    case FAILED:
                        String combinedTaskErrorMessage = statusMessage.getErrorCode() != null ?
                                statusMessage.getErrorCode() + ": " + statusMessage.getMessage() :
                                statusMessage.getMessage();
                        taskUpdated = robotTaskService.markTaskAsFailed(statusMessage.getTaskId(), combinedTaskErrorMessage, new Date());
                        break;
                    default:
                        log.warn("Received unhandled RobotTaskStatusEnum: {} for task {} from robot {}. Ignoring.",
                                taskNewStatus, statusMessage.getTaskId(), robotUIDFromTopicSource);
                            taskUpdated = false; // To prevent logging success
                            break; 
                }
                 if (taskUpdated) {
                    log.info("Successfully updated task {} to status {} for robot {} based on MQTT message.",
                            statusMessage.getTaskId(), taskNewStatus, robotUIDFromTopicSource);
                } else {
                    log.warn("Failed to update task {} to status {} for robot {} (or task not found/invalid state transition). Message: {}", 
                            statusMessage.getTaskId(), taskNewStatus, robotUIDFromTopicSource, statusMessage.getMessage());
                }
            } 
            } 
            
            // Scenario 2: Message contains a general robot status update.
            // This is now a separate "if", not an "else if", so it can execute even if a taskId was present.
            // This is crucial for ACK messages that also convey a change in the robot's main status (e.g., from MOVING to CHARGING).
            if (statusMessage.getActualRobotStatus() != null) { 
                log.debug("Handling general status update from robot {}", robotUIDFromTopicSource);
                
                String actualRobotUIDToUse = StringUtils.isNotBlank(statusMessage.getRobotUid()) ? statusMessage.getRobotUid() : robotUIDFromTopicSource;
                if (StringUtils.isBlank(actualRobotUIDToUse)) {
                     log.warn("Cannot determine a valid robot UID for general status update (from topic: {}, from message body: {}). Ignoring update.", 
                              robotUIDFromTopicSource, statusMessage.getRobotUid());
                     return; // Critical to have a UID, so we can return here.
                }

                RobotStatusEnum robotStatus = RobotStatusEnum.getEnumByValue(statusMessage.getActualRobotStatus());
                if (robotStatus == null) {
                    log.warn("Received unknown status value '{}' from robot {}. Message: {}", 
                             statusMessage.getActualRobotStatus(), actualRobotUIDToUse, statusMessage);
                }

                String location = statusMessage.getLocation(); 
                Integer batteryLevel = statusMessage.getBatteryLevel(); 
                Long currentRobotTask = statusMessage.getActiveTaskId(); 

                log.info("Processing general status update for robot {}: Status={}, Location={}, Battery={}, ActiveTaskID={}. (Original topic UID: {})",
                        actualRobotUIDToUse, robotStatus, location, batteryLevel, currentRobotTask, robotUIDFromTopicSource);

                chargingRobotService.updateRobotStatus(actualRobotUIDToUse, robotStatus, location, batteryLevel, currentRobotTask, new Date());
            } 
            // A final check for completely empty/un-actionable messages.
            else if (statusMessage.getTaskId() == null) {
                log.warn("RobotStatusMessage for robot {} has no taskId and no general status fields. Payload: {}. Ignoring.", 
                         robotUIDFromTopicSource, payloadJson);
            }

        } catch (Exception e) {
            log.error("Error processing robot status update for robot {}. Payload: {}. Error: {}",
                    robotUIDFromTopicSource, payloadJson, e.getMessage(), e);
        }
    }
} 