import asyncio
import threading
from datetime import datetime, timezone
from typing import List, Optional

from letta.agents.base_agent import BaseAgent
from letta.agents.letta_agent import LettaAgent
from letta.groups.helpers import stringify_message
from letta.interface import AgentInterface
from letta.orm import User
from letta.schemas.agent import AgentState
from letta.schemas.enums import JobStatus
from letta.schemas.job import JobUpdate
from letta.schemas.letta_message_content import TextContent
from letta.schemas.message import Message, MessageCreate
from letta.schemas.run import Run
from letta.schemas.usage import LettaUsageStatistics
from letta.server.rest_api.interface import StreamingServerInterface
from letta.services.group_manager import GroupManager
from letta.services.job_manager import JobManager
from letta.services.message_manager import MessageManager


class SleeptimeMultiAgent(BaseAgent):
    def __init__(
        self,
        interface: AgentInterface,
        agent_state: AgentState,
        user: User,
        # mcp_clients: Optional[Dict[str, BaseMCPClient]] = None,
        # custom
        group_id: str = "",
        agent_ids: List[str] = [],
        description: str = "",
        sleeptime_agent_frequency: Optional[int] = None,
    ):
        super().__init__(interface, agent_state, user)
        self.group_id = group_id
        self.agent_ids = agent_ids
        self.description = description
        self.sleeptime_agent_frequency = sleeptime_agent_frequency
        self.group_manager = GroupManager()
        self.message_manager = MessageManager()
        self.job_manager = JobManager()
        # TODO: add back MCP support with new agent loop
        self.mcp_clients = {}

    def _run_async_in_new_thread(self, coro):
        """Run an async coroutine in a new thread with its own event loop"""

        def run_async():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                loop.run_until_complete(coro)
            finally:
                loop.close()
                asyncio.set_event_loop(None)

        thread = threading.Thread(target=run_async)
        thread.daemon = True
        thread.start()

    def _issue_background_task(
        self,
        participant_agent_id: str,
        messages: List[Message],
        chaining: bool,
        max_chaining_steps: Optional[int],
        token_streaming: bool,
        metadata: Optional[dict],
        put_inner_thoughts_first: bool,
        last_processed_message_id: str,
    ) -> str:
        run = Run(
            user_id=self.user.id,
            status=JobStatus.created,
            metadata={
                "job_type": "background_agent_send_message_async",
                "agent_id": participant_agent_id,
            },
        )
        run = self.job_manager.create_job(pydantic_job=run, actor=self.user)

        self._run_async_in_new_thread(
            self._perform_background_agent_step(
                participant_agent_id=participant_agent_id,
                messages=messages,
                chaining=chaining,
                max_chaining_steps=max_chaining_steps,
                token_streaming=token_streaming,
                metadata=metadata,
                put_inner_thoughts_first=put_inner_thoughts_first,
                last_processed_message_id=last_processed_message_id,
                run_id=run.id,
            )
        )

        return run.id

    async def _perform_background_agent_step(
        self,
        participant_agent_id: str,
        messages: List[Message],
        chaining: bool,
        max_chaining_steps: Optional[int],
        token_streaming: bool,
        metadata: Optional[dict],
        put_inner_thoughts_first: bool,
        last_processed_message_id: str,
        run_id: str,
    ) -> LettaUsageStatistics:
        try:
            job_update = JobUpdate(status=JobStatus.running)
            self.job_manager.update_job_by_id(job_id=run_id, job_update=job_update, actor=self.user)

            participant_agent_state = self.agent_manager.get_agent_by_id(participant_agent_id, actor=self.user)
            participant_agent = LettaAgent(
                agent_state=participant_agent_state,
                interface=StreamingServerInterface(),
                user=self.user,
                mcp_clients=self.mcp_clients,
            )

            prior_messages = []
            if self.sleeptime_agent_frequency:
                try:
                    prior_messages = self.message_manager.list_messages_for_agent(
                        agent_id=self.agent_state.id,
                        actor=self.user,
                        after=last_processed_message_id,
                        before=messages[0].id,
                    )
                except Exception as e:
                    print(f"Error fetching prior messages: {str(e)}")
                    # continue with just latest messages

            transcript_summary = [stringify_message(message) for message in prior_messages + messages]
            transcript_summary = [summary for summary in transcript_summary if summary is not None]
            message_text = "\n".join(transcript_summary)

            participant_agent_messages = [
                Message(
                    id=Message.generate_id(),
                    agent_id=participant_agent.agent_state.id,
                    role="user",
                    content=[TextContent(text=message_text)],
                    group_id=self.group_id,
                )
            ]

            # Convert Message objects to MessageCreate objects
            message_creates = [
                MessageCreate(
                    role=m.role,
                    content=m.content[0].text if m.content and len(m.content) == 1 else m.content,
                    name=m.name,
                    otid=m.otid,
                    sender_id=m.sender_id,
                )
                for m in participant_agent_messages
            ]

            result = participant_agent.step(
                input_messages=message_creates,
                chaining=chaining,
                max_chaining_steps=max_chaining_steps,
                stream=token_streaming,
                skip_verify=True,
                metadata=metadata,
                put_inner_thoughts_first=put_inner_thoughts_first,
            )
            job_update = JobUpdate(
                status=JobStatus.completed,
                completed_at=datetime.now(timezone.utc),
                metadata={
                    "result": result.model_dump(mode="json"),
                    "agent_id": participant_agent.agent_state.id,
                },
            )
            self.job_manager.update_job_by_id(job_id=run_id, job_update=job_update, actor=self.user)
            return result
        except Exception as e:
            job_update = JobUpdate(
                status=JobStatus.failed,
                completed_at=datetime.now(timezone.utc),
                metadata={"error": str(e)},
            )
            self.job_manager.update_job_by_id(job_id=run_id, job_update=job_update, actor=self.user)
            raise

    def step(
        self,
        input_messages: List[MessageCreate],
        chaining: bool = True,
        max_chaining_steps: Optional[int] = None,
        put_inner_thoughts_first: bool = True,
        **kwargs,
    ) -> LettaUsageStatistics:
        run_ids = []

        # Load settings
        token_streaming = self.interface.streaming_mode if hasattr(self.interface, "streaming_mode") else False
        metadata = self.interface.metadata if hasattr(self.interface, "metadata") else None

        # Prepare new messages
        new_messages = []
        for message in input_messages:
            if isinstance(message.content, str):
                message.content = [TextContent(text=message.content)]
            message.group_id = self.group_id
            new_messages.append(message)

        try:
            # Load main agent
            main_agent = LettaAgent(
                agent_state=self.agent_state,
                interface=self.interface,
                user=self.user,
                mcp_clients=self.mcp_clients,
            )
            # Perform main agent step
            usage_stats = main_agent.step(
                input_messages=new_messages,
                chaining=chaining,
                max_chaining_steps=max_chaining_steps,
                stream=token_streaming,
                skip_verify=True,
                metadata=metadata,
                put_inner_thoughts_first=put_inner_thoughts_first,
            )

            # Update turns counter
            turns_counter = None
            if self.sleeptime_agent_frequency is not None and self.sleeptime_agent_frequency > 0:
                turns_counter = self.group_manager.bump_turns_counter(group_id=self.group_id, actor=self.user)

            # Perform participant steps
            if self.sleeptime_agent_frequency is None or (
                turns_counter is not None and turns_counter % self.sleeptime_agent_frequency == 0
            ):
                last_response_messages = [message for sublist in usage_stats.steps_messages for message in sublist]
                last_processed_message_id = self.group_manager.get_last_processed_message_id_and_update(
                    group_id=self.group_id, last_processed_message_id=last_response_messages[-1].id, actor=self.user
                )
                for participant_agent_id in self.agent_ids:
                    try:
                        run_id = self._issue_background_task(
                            participant_agent_id,
                            last_response_messages,
                            chaining,
                            max_chaining_steps,
                            token_streaming,
                            metadata,
                            put_inner_thoughts_first,
                            last_processed_message_id,
                        )
                        run_ids.append(run_id)

                    except Exception as e:
                        # Handle individual task failures
                        print(f"Agent processing failed: {str(e)}")
                        raise e

        except Exception as e:
            raise e
        finally:
            self.interface.step_yield()

        self.interface.step_complete()

        usage_stats.run_ids = run_ids
        return LettaUsageStatistics(**usage_stats.model_dump())

    async def step_stream(
        self,
        input_messages: List[MessageCreate],
        max_steps: int = 10,
    ):
        """
        Streaming is not currently supported for SleeptimeMultiAgent.
        This method is required by the BaseAgent abstract class but not yet implemented.
        """
        from letta.log import get_logger

        logger = get_logger(__name__)
        logger.error(
            f"step_stream called on SleeptimeMultiAgent (agent_id={self.agent_state.id}, "
            f"group_id={self.group_id}). Streaming is not supported for multi-agent groups."
        )
        raise NotImplementedError("Streaming is not supported for SleeptimeMultiAgent. Use the non-streaming step() method instead.")
