from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator, Optional

import openai

from letta.schemas.letta_message import UserMessage
from letta.schemas.letta_response import LettaResponse
from letta.schemas.user import User
from letta.services.agent_manager import AgentManager
from letta.services.message_manager import MessageManager


class BaseAgent(ABC):
    """
    Abstract base class for AI agents, handling message management, tool execution,
    and context tracking.
    """

    def __init__(
        self,
        agent_id: str,
        # TODO: Make required once client refactor hits
        openai_client: Optional[openai.AsyncClient],
        message_manager: MessageManager,
        agent_manager: AgentManager,
        actor: User,
    ):
        self.agent_id = agent_id
        self.openai_client = openai_client
        self.message_manager = message_manager
        self.agent_manager = agent_manager
        self.actor = actor

    @abstractmethod
    async def step(self, input_message: UserMessage, max_steps: int = 10) -> LettaResponse:
        """
        Main execution loop for the agent.
        """
        raise NotImplementedError

    @abstractmethod
    async def step_stream(self, input_message: UserMessage, max_steps: int = 10) -> AsyncGenerator[str, None]:
        """
        Main async execution loop for the agent. Implementations must yield messages as SSE events.
        """
        raise NotImplementedError

    def pre_process_input_message(self, input_message: UserMessage) -> Any:
        """
        Pre-process function to run on the input_message.
        """
        return input_message.model_dump()
