"""Middleware for injecting ICARUS memory file structure into the system prompt.

Complements the SDK's MemoryMiddleware (which provides behavioral guidelines
about *when* to remember) by telling the agent *where* things live — the
MEMORY.md file, daily logs, contacts, transcripts, and voice profiles.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

from deepagents.middleware._utils import append_to_system_message
from langchain.agents.middleware.types import (
    AgentMiddleware,
    ModelRequest,
    ModelResponse,
)

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

_TEMPLATE = (Path(__file__).parent.parent / "assets" / "memory_instructions.md").read_text()


class MemoryStructureMiddleware(AgentMiddleware):
    """Inject ICARUS-specific memory file structure into the system prompt."""

    def __init__(self, memory_dir: Path) -> None:
        self._section = _TEMPLATE.replace("{memory_dir}", str(memory_dir))

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """Append memory file structure to system prompt (sync)."""
        new_system_message = append_to_system_message(request.system_message, self._section)
        return handler(request.override(system_message=new_system_message))

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelResponse:
        """Append memory file structure to system prompt (async)."""
        new_system_message = append_to_system_message(request.system_message, self._section)
        return await handler(request.override(system_message=new_system_message))


__all__ = ["MemoryStructureMiddleware"]
