from dataclasses import dataclass
from typing import List, Callable
from enum import Enum
import datetime

from .io_buffer import IOBuffer
from .util import format_timestamp

class MessageType(Enum):
    USER = "user"
    ASSISTANT = "assistant"

@dataclass
class ChatMessage:
    id: str  # Formatted timestamp
    content: str
    message_type: MessageType

class ChatIOBuffer(IOBuffer):
    def __init__(self):
        self._messages: List[ChatMessage] = []
        self._last_read_timestamp: str = format_timestamp(datetime.datetime.now())
        self._change_handlers: List[Callable] = []
    
    @property
    def messages(self) -> List[ChatMessage]:
        return self._messages.copy()
    
    @property
    def last_read_timestamp(self) -> str:
        return self._last_read_timestamp
    
    def read(self) -> str:
        unread_messages = self._unread_messages()
        self._last_read_timestamp = self._format_now()
        self._notify_handlers([], [])
        return unread_messages
    
    def write(self, content: str) -> None:
        message = ChatMessage(
            id=self._format_now(),
            content=content,
            message_type=MessageType.ASSISTANT
        )
        self._messages.append(message)
        self._notify_handlers(
            new_messages=[message],
            deleted_message_ids=[]
        )
    
    def buffer_length(self) -> int:
        return len(self._unread_messages())
    
    def add_user_message(self, content: str):
        message = ChatMessage(
            id=self._format_now(),
            content=content,
            message_type=MessageType.USER
        )
        self._messages.append(message)
        self._notify_handlers(
            new_messages=[message],
            deleted_message_ids=[]
        )
    
    def delete_message(self, message_id: str):
        for i, msg in enumerate(self._messages):
            if msg.id == message_id:
                self._messages.pop(i)
                self._notify_handlers(
                    new_messages=[],
                    deleted_message_ids=[message_id]
                )
                break
    
    def clear(self) -> None:
        deleted_ids = [msg.id for msg in self._messages]
        self._messages.clear()
        self._last_read_timestamp = self._format_now()
        self._notify_handlers(
            new_messages=[],
            deleted_message_ids=deleted_ids
        )
    
    def add_change_handler(self, handler: Callable[[
        List[ChatMessage], # new messages
        List[str], # deleted message ids
        str # last read timestamp
    ], None]) -> None:
        self._change_handlers.append(handler)
    
    def _notify_handlers(self, new_messages: List[ChatMessage], 
                        deleted_message_ids: List[str]) -> None:
        for handler in self._change_handlers:
            handler(new_messages, deleted_message_ids, self._last_read_timestamp)
    
    def _unread_messages(self) -> List[ChatMessage]:
        return "\n".join(msg.content for msg in self._messages 
                if msg.message_type == MessageType.USER and msg.id > self._last_read_timestamp)
    
    def _format_now(self) -> str:
        return format_timestamp(datetime.datetime.now())