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"
DAYGLO_THRESHOLD_TOPIC = "homeassistant/sensor/dayglo_threshold/state"
DISCOVERY_PREFIX = "homeassistant"

# Default threshold
dayglo_threshold = 50

def on_connect(client, userdata, flags, rc):
    print("Connected with result code", rc)
    client.subscribe(MQTT_SNAPSHOT_TOPIC)
    publish_discovery_configurations()

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"
    }
    threshold_config = {
        "name": "Dayglo Threshold",
        "state_topic": DAYGLO_THRESHOLD_TOPIC,
        "unit_of_measurement": "%",
        "value_template": "{{ value_json.threshold }}",
        "icon": "mdi:tune-vertical",
        "unique_id": "mqtt_dayglo_threshold"
    }

    client.publish(f"{DISCOVERY_PREFIX}/sensor/dayglo_rating/config", json.dumps(rating_config), retain=True)
    client.publish(f"{DISCOVERY_PREFIX}/sensor/dayglo_threshold/config", json.dumps(threshold_config), retain=True)

def on_message(client, userdata, msg):
    if msg.topic == MQTT_SNAPSHOT_TOPIC:
        print("Snapshot received")
        process_snapshot(msg.payload)
    elif msg.topic == DAYGLO_THRESHOLD_TOPIC:
        global dayglo_threshold
        dayglo_threshold = float(msg.payload.decode('utf-8'))
        print("Dayglo threshold updated to:", dayglo_threshold)

def process_snapshot(payload):
    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)
        client.publish(MQTT_TOPIC_PUBLISH, json.dumps({"rating": rating}))

def calculate_dayglo_rating(image):
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    lower_color = np.array([20, 100, 100])
    upper_color = np.array([40, 255, 255])
    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()
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()