import os
import paho.mqtt.client as mqtt
import numpy as np
import cv2
import base64
import json

# Configuration
MQTT_BROKER = os.environ.get('MQTT_BROKER', '10.59.221.172')
MQTT_PORT = int(os.environ.get('MQTT_PORT', '1883'))
MQTT_USERNAME = os.getenv('MQTT_USERNAME', 'your_username')
MQTT_PASSWORD = os.getenv('MQTT_PASSWORD', 'your_password')
MQTT_SNAPSHOT_TOPIC = "/frigate/patiocam/person/snapshot"
MQTT_TOPIC_PUBLISH = "homeassistant/sensor/dayglo_rating/state"
DISCOVERY_PREFIX = "homeassistant"

# Default rating
last_rating = 0

# Default color thresholds for dayglo detection
LOWER_COLOR = np.array([20, 100, 100])
UPPER_COLOR = np.array([40, 255, 255])

def on_connect(client, userdata, flags, rc):
    print("Connected with result code", rc)
    client.subscribe(MQTT_SNAPSHOT_TOPIC)
    publish_discovery_configurations()
    # Publish initial rating of 0
    publish_rating(0)
    # Attempt to process the latest snapshot upon connection
    process_latest_snapshot()

def publish_discovery_configurations():
    rating_config = {
        "name": "Dayglo Rating",
        "state_topic": MQTT_TOPIC_PUBLISH,
        "unit_of_measurement": "%",
        "value_template": "{{ value_json.rating }}",
        "icon": "mdi:brush",
        "unique_id": "mqtt_dayglo_rating"
    }
    client.publish(f"{DISCOVERY_PREFIX}/sensor/dayglo_rating/config", json.dumps(rating_config), retain=True)

def publish_rating(rating):
    global last_rating
    last_rating = rating
    client.publish(MQTT_TOPIC_PUBLISH, json.dumps({"rating": rating}))

def on_message(client, userdata, msg):
    if msg.topic == MQTT_SNAPSHOT_TOPIC:
        print("Snapshot received")
        process_snapshot(msg.payload)

def process_latest_snapshot():
    # Placeholder for processing the latest snapshot (if available)
    print("Attempting to process the latest snapshot...")
    # Assuming we have a way to get the latest snapshot payload
    # For now, this is just a debug statement

def process_snapshot(payload):
    print("Processing snapshot...")
    image_data = base64.b64decode(payload)
    nparr = np.frombuffer(image_data, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if image is not None:
        rating = calculate_dayglo_rating(image)
        print("Dayglo Rating calculated:", rating)
        publish_rating(rating)
    else:
        print("Invalid image received")

def calculate_dayglo_rating(image):
    print("Calculating dayglo rating...")
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv_image, LOWER_COLOR, UPPER_COLOR)
    dayglo_pixels = cv2.countNonZero(mask)
    total_pixels = image.shape[0] * image.shape[1]
    rating = (dayglo_pixels / total_pixels) * 100
    return rating

client = mqtt.Client(protocol=mqtt.MQTTv311)
client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
client.on_connect = on_connect
client.on_message = on_message
client.connect(MQTT_BROKER, MQTT_PORT, 60)
client.loop_forever()