from __future__ import annotations
import asyncio
import dataclasses
import threading

LINE_ENDING = b"\r\n"

class ThreadSafeDict(object):
    def __init__(self):
        self.lock = threading.Lock()
        self.dict = {}
    def set(self, key, value):
        with self.lock:
            self.dict[key] = value
    def get(self, key):
        with self.lock:
            return self.dict.get(key)
    def __getitem__(self, key):
        return self.get(key)
    def __setitem__(self, key, value):
        self.set(key, value)

def get_senddata(data: bytes) -> bytes:
    """Return the data to send back to the client."""
    return b"+" + data + LINE_ENDING
def get_nildata() -> bytes:
    return b"$-1" + LINE_ENDING


@dataclasses.dataclass
class RedisCommand:
    """A Redis command."""
    command: bytes
    args: list[bytes]
    @staticmethod
    def parse(data: bytes) -> RedisCommand:
        """Parse the data into a RedisCommand."""
        array = data.split(LINE_ENDING)
        command = array[2].upper()
        args = []
        for i in range(3, len(array)):
            if i % 2 == 0:
                args.append(array[i])
        return RedisCommand(command, args)
async def redis(reader, writer):
    itemMap = ThreadSafeDict()
    while True:
        data = await reader.read(1024)
        if not data:
            break
        cmd = RedisCommand.parse(data)
        resp = get_nildata()
        if cmd.command == b"ECHO":
            resp = get_senddata(cmd.args[0])
        elif cmd.command == b"SET":
            itemMap[cmd.args[0]] = cmd.args[1]
            resp = get_senddata(b"OK")
        elif cmd.command == b"GET":
            resp = itemMap.get(cmd.args[0])
            if resp:
                resp = get_senddata(resp)
            else:
                resp = get_nildata()
        elif cmd.command == b"PING":
            resp = get_senddata(b"PONG")
        writer.write(resp)
        await writer.drain()

async def main(host, port):
    srv = await asyncio.start_server(redis, host, port)
    async with srv:
        await srv.serve_forever()

if __name__ == "__main__":
    asyncio.run(main("localhost", 6379))

