"""Middleware for injecting temporal context into the system prompt.

Provides the agent with awareness of the current date and session start
time.  Only includes values that are **stable within a session** so the
system prompt remains identical across model calls — a requirement for
prompt caching on Anthropic, OpenAI, Gemini, and compatible providers.

For exact current time the agent can call the ``get_current_time`` tool.
"""

from __future__ import annotations

import logging
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Annotated,
    Any,
    NotRequired,
    cast,
)

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

_log = logging.getLogger(__name__)

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from langgraph.runtime import Runtime


# ---------------------------------------------------------------------------
# State schema
# ---------------------------------------------------------------------------


class TimeContextState(AgentState):
    """State fields used by :class:`TimeContextMiddleware`."""

    _time_session_start: NotRequired[Annotated[str, PrivateStateAttr]]
    """ISO 8601 timestamp of session start (first ``before_agent`` call)."""


# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------


class TimeContextMiddleware(AgentMiddleware):
    """Inject stable time context (date + session start) into the system
    prompt.  No backend or shell access required.

    The section is fully determined by ``_time_session_start`` which is
    set once in :meth:`before_agent`, so the system prompt never changes
    between turns — enabling prompt-prefix caching on all providers.
    """

    state_schema = TimeContextState

    def before_agent(  # type: ignore[override]
        self,
        state: TimeContextState,
        runtime: Runtime,  # noqa: ARG002
    ) -> dict[str, Any] | None:
        """Record session start time (once, on the first turn)."""
        if not state.get("_time_session_start"):
            return {"_time_session_start": datetime.now().astimezone().isoformat()}
        return None

    @staticmethod
    def _build_time_section(state: TimeContextState) -> str:
        """Build the stable time-context markdown section."""
        session_iso = state.get("_time_session_start")
        if not session_iso:
            _log.warning("_time_session_start missing — before_agent may not have run")
            now = datetime.now().astimezone()
            return f"## Current Time\n\n- **Date**: {now.strftime('%A, %Y-%m-%d')}\n"

        # Date is pinned to session start — may go stale across midnight.
        # Agent can call get_current_time tool for the real date.
        session_start = datetime.fromisoformat(session_iso)
        return (
            "## Current Time\n"
            "\n"
            f"- **Date**: {session_start.strftime('%A, %Y-%m-%d')}\n"
            f"- **Session started**: {session_start.strftime('%H:%M %Z')}\n"
        )

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """Inject time context into system prompt (sync)."""
        state = cast("TimeContextState", request.state)
        section = self._build_time_section(state)
        new_system_message = append_to_system_message(request.system_message, section)
        return handler(request.override(system_message=new_system_message))

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelResponse:
        """Inject time context into system prompt (async)."""
        state = cast("TimeContextState", request.state)
        section = self._build_time_section(state)
        new_system_message = append_to_system_message(request.system_message, section)
        return await handler(request.override(system_message=new_system_message))


__all__ = ["TimeContextMiddleware"]
