#!/usr/bin/env python
import redis, pickle
from enum import Enum
import json

class TOPICS(Enum):
    ml_vision_to_score = 'ml_vision_to_score'
    ml_vision_to_score_ext = 'ml_vision_to_score_ext'
    ml_vision_objdet_failed = 'ml_vision_objdet_failed'
    ml_vision_objdet_success = 'ml_vision_objdet_success'
    ml_vision_objdet_skipped = 'ml_vision_objdet_skipped'
    ml_vision_objdet_results_db_success = 'ml_vision_objdet_db_upload_success'
    ml_vision_objdet_results_db_failed = 'ml_vision_objdet_db_upload_failed'
    ml_vision_objdet_results_pg_success = 'ml_vision_objdet_pg_upload_success'
    ml_vision_objdet_results_pg_failed = 'ml_vision_objdet_pg_upload_failed'
    ml_vision_objdet_results_purge_success = 'ml_vision_objdet_purge_success'    
    ml_vision_objdet_results_purge_failed = 'ml_vision_objdet_purge_failed'
    ml_vision_objdet_results_purge_skipped = 'ml_vision_objdet_purge_skipped'
    ml_vision_videos_modify_success = 'ml_vision_videos_modify_success'
    ml_vision_videos_modify_failed = 'ml_vision_videos_modify_failed'
    ml_vision_embedding_success = 'ml_vision_embedding_success'
    ml_vision_embedding_fail = 'ml_vision_embedding_fail'
    ml_vision_embedding_skipped = 'ml_vision_embedding_skipped'


r = redis.StrictRedis()

def publish(topic, message):
    if isinstance(topic, TOPICS):
        topic = topic.value

    
    if isinstance(message, str):
        r.rpush(topic, message)
    else:
        r.rpush(topic, json.dumps(message))


def parse_message(msg):
    try:
        return json.loads(msg)
    except:
        return msg.decode()
        
def consume(topic):
    if isinstance(topic, TOPICS):
        topic = topic.value

    msg = r.rpop(topic)
    if msg is None:
        return None
    else:
        return parse_message(msg)
        
def peek(topic):
    if isinstance(topic, TOPICS):
        topic = topic.value

    vals = r.lrange(topic,0, 0)
    if len(vals) == 0:
        return None
    else:
        return parse_message(vals[0])

# from kafka import KafkaProducer, KafkaConsumer
# import pickle

# producer = KafkaProducer(bootstrap_servers=["192.168.1.242:9092"])
# consumers = dict()
# # %%

# def publish(topic, message):
#     producer.send(topic, value=pickle.dumps(message))
#     producer.flush()


# def get_consumer_by_topic(topic):
#     if topic in consumers:
#         return consumers[topic]

#     csumer = KafkaConsumer(
#         bootstrap_servers="localhost:9092",
#         group_id="wq",
#         auto_offset_reset="earliest",
#         enable_auto_commit=True,
#         consumer_timeout_ms=1000,
#     )
#     csumer.subscribe(topic)
#     consumers[topic] = csumer


# def get_a_message(topic):
#     csumer = get_consumer_by_topic(topic)
#     return csumer