import random

import aioredis
import orjson

# Kafka
from aiokafka import AIOKafkaProducer
from redis import StrictRedis

import util

# KAFKA_TOPIC = "msg"

log = util.get_logger("db")

# Redis (legacy)
r = StrictRedis(unix_socket_path="/var/run/redis/redis.sock", db=0)

# AIORedis
ar = aioredis.from_url("unix:///var/run/redis/redis.sock", db=0)

TYPES_MAIN = [
    "msg",
    "notice",
    "action",
    "part",
    "join",
    "kick",
    "quit",
    "nick",
    "mode",
    "topic",
    "update",
]
TYPES_META = ["who"]
TYPES_INT = ["conn", "highlight", "znc", "query", "self"]
KEYPREFIX = "queue."


async def store_kafka_batch(data):
    log.debug(f"Storing Kafka batch of {len(data)} messages")
    producer = AIOKafkaProducer(bootstrap_servers="kafka:9092")
    await producer.start()
    batch = producer.create_batch()
    for msg in data:
        if msg["type"] in TYPES_MAIN:
            index = "main"
            # schema = mc_s.schema_main
        elif msg["type"] in TYPES_META:
            index = "meta"
            # schema = mc_s.schema_meta
        elif msg["type"] in TYPES_INT:
            index = "internal"
            # schema = mc_s.schema_int

        KAFKA_TOPIC = index
        # normalise fields
        for key, value in list(msg.items()):
            if value is None:
                del msg[key]
            # if key in schema:
            #     if isinstance(value, int):
            #         if schema[key].startswith("string") or schema[key].startswith(
            #             "text"
            #         ):
            #             msg[key] = str(value)
        body = orjson.dumps(msg)
        # orjson returns bytes
        # body = str.encode(message)
        if "ts" not in msg:
            raise Exception("No TS in msg")
        metadata = batch.append(key=None, value=body, timestamp=msg["ts"])
        if metadata is None:
            partitions = await producer.partitions_for(KAFKA_TOPIC)
            partition = random.choice(tuple(partitions))
            await producer.send_batch(batch, KAFKA_TOPIC, partition=partition)
            log.debug(f"{batch.record_count()} messages sent to partition {partition}")
            batch = producer.create_batch()
            continue

    partitions = await producer.partitions_for(KAFKA_TOPIC)
    partition = random.choice(tuple(partitions))
    await producer.send_batch(batch, KAFKA_TOPIC, partition=partition)
    log.debug(f"{batch.record_count()} messages sent to partition {partition}")
    await producer.stop()


async def queue_message(msg):
    """
    Queue a message on the Redis buffer.
    """
    src = msg["src"]
    message = orjson.dumps(msg)

    key = f"{KEYPREFIX}{src}"
    # log.debug(f"Queueing single message of string length {len(message)}")
    await ar.sadd(key, message)


async def queue_message_bulk(data):
    """
    Queue multiple messages on the Redis buffer.
    """
    # log.debug(f"Queueing message batch of length {len(data)}")
    for msg in data:
        src = msg["src"]
        message = orjson.dumps(msg)

        key = f"{KEYPREFIX}{src}"
        await ar.sadd(key, message)