import json
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.layers import get_channel_layer
from asgiref.sync import sync_to_async


class CitizenChatConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        self.session_id = self.scope["url_route"]["kwargs"]["session_id"]
        self.group_name = f"session_{self.session_id}"

        await self.channel_layer.group_add(self.group_name, self.channel_name)
        await self.accept()

    async def disconnect(self, close_code):
        await self.channel_layer.group_discard(self.group_name, self.channel_name)

    async def receive_json(self, content, **kwargs):
        message_type = content.get("type", None)

        if message_type == "chat_message":
            message = content["message"]
            await self.channel_layer.group_send(
                self.group_name,
                {
                    "type": "group_message",
                    "message": message,
                    "user": "citizen",  # Or fetch from scope if authenticated
                },
            )

    async def group_message(self, event):
        message = event["message"]
        user = event["user"]

        await self.send_json({"type": "chat_message", "message": message, "user": user})


class AgentConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        self.queue_name = "agents"
        self.session_id = self.scope["url_route"]["kwargs"].get("session_id")

        if self.session_id:
            self.group_name = f"session_{self.session_id}"
            await self.channel_layer.group_add(self.group_name, self.channel_name)
        else:
            await self.channel_layer.group_add(self.queue_name, self.channel_name)

        await self.accept()

    async def disconnect(self, close_code):
        if self.session_id:
            await self.channel_layer.group_discard(self.group_name, self.channel_name)
        else:
            await self.channel_layer.group_discard(self.queue_name, self.channel_name)

    async def receive_json(self, content, **kwargs):
        message = content["message"]
        await self.channel_layer.group_send(
            self.group_name,
            {"type": "agent_message", "message": message, "user": "agent"},
        )

    async def agent_message(self, event):
        message = event["message"]
        user = event["user"]
        await self.send_json({"type": "chat_message", "message": message, "user": user})
