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

Tells the agent what PIM data is available — contact count, calendar events,
open tasks. Complements the PIM tools by giving the agent awareness of the
data landscape without requiring a tool call.
"""

from __future__ import annotations

from pathlib import Path
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,
)

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from langgraph.runtime import Runtime

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


def _count_files(directory: Path, pattern: str) -> int:
    """Count files matching a glob pattern in a directory."""
    if not directory.exists():
        return 0
    return sum(1 for _ in directory.glob(pattern))


def _count_events(calendar_dir: Path) -> int:
    """Count total calendar events (number of .ics files).

    We avoid full RRULE expansion for the middleware summary — just count
    files as a rough indicator. The actual tool does full expansion.
    """
    return _count_files(calendar_dir, "*.ics")


# Statuses that mean a todo is no longer open.
_CLOSED_STATUSES = {b"COMPLETED", b"CANCELLED"}


def _count_open_todos(todos_dir: Path) -> int:
    """Count open todos by checking STATUS in each .ics file.

    Reads each file as bytes and scans for the STATUS property.
    Excludes COMPLETED and CANCELLED todos. This is cheap — no
    icalendar parsing, just a byte-level scan per file.
    """
    if not todos_dir.exists():
        return 0
    count = 0
    for ics_path in todos_dir.glob("*.ics"):
        try:
            data = ics_path.read_bytes()
        except OSError:
            continue
        # Extract STATUS value: look for b"STATUS:" at start of a line
        closed = False
        for line in data.split(b"\n"):
            stripped = line.strip()
            if stripped.startswith(b"STATUS:"):
                status_value = stripped[7:].strip()
                if status_value in _CLOSED_STATUSES:
                    closed = True
                break
        if not closed:
            count += 1
    return count


class PIMContextState(AgentState):
    """State fields used by :class:`PIMContextMiddleware`."""

    _pim_contact_count: NotRequired[Annotated[int, PrivateStateAttr]]
    _pim_event_count: NotRequired[Annotated[int, PrivateStateAttr]]
    _pim_open_todo_count: NotRequired[Annotated[int, PrivateStateAttr]]


class PIMContextMiddleware(AgentMiddleware):
    """Inject PIM context summary into the system prompt.

    Counts are computed once per turn (in :meth:`before_agent`) and stored
    in graph state so that multiple model calls within a single turn
    (e.g. after tool results) reuse the same values — keeping the system
    prompt stable for prompt caching.
    """

    state_schema = PIMContextState

    def __init__(
        self,
        contacts_dir: Path,
        calendar_dir: Path,
        todos_dir: Path,
    ) -> None:
        self._contacts_dir = contacts_dir
        self._calendar_dir = calendar_dir
        self._todos_dir = todos_dir

    def before_agent(  # type: ignore[override]
        self,
        state: PIMContextState,
        runtime: Runtime,  # noqa: ARG002
    ) -> dict[str, Any]:
        """Refresh PIM counts from disk (once per turn)."""
        return {
            "_pim_contact_count": _count_files(self._contacts_dir, "*.vcf"),
            "_pim_event_count": _count_events(self._calendar_dir),
            "_pim_open_todo_count": _count_open_todos(self._todos_dir),
        }

    def _build_section(self, state: PIMContextState) -> str:
        """Build the PIM context section from state-cached counts."""
        return (
            _TEMPLATE
            .replace("{contacts_dir}", str(self._contacts_dir))
            .replace("{calendar_dir}", str(self._calendar_dir))
            .replace("{todos_dir}", str(self._todos_dir))
            .replace("{contact_count}", str(state.get("_pim_contact_count", 0)))
            .replace("{event_count}", str(state.get("_pim_event_count", 0)))
            .replace("{open_todo_count}", str(state.get("_pim_open_todo_count", 0)))
        )

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """Append PIM context to system prompt (sync)."""
        state = cast("PIMContextState", request.state)
        section = self._build_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:
        """Append PIM context to system prompt (async)."""
        state = cast("PIMContextState", request.state)
        section = self._build_section(state)
        new_system_message = append_to_system_message(request.system_message, section)
        return await handler(request.override(system_message=new_system_message))


__all__ = ["PIMContextMiddleware"]
