"""Process memory profiling utilities.

Three instruments:

1. ``rss_mb()`` — current RSS in MiB (cheap, single call).

2. ``profile_task(name)`` — async context manager that logs RSS delta before
   and after a block.  Use it to wrap every background task so we can see
   exactly which one is the memory hog.

3. ``top_allocations(limit)`` — returns the top tracemalloc frames by
   cumulative memory.  Requires tracemalloc to be started at boot
   (``tracemalloc.start()`` in ``main.py``).

4. ``MemoryLogMiddleware`` — ASGI middleware that logs RSS after every
   request above a configurable threshold.

Usage in background tasks::

    from musehub.debug.memory import profile_task

    async with profile_task("symbol-index repo=abc123"):
        await build_symbol_index(session, repo_id, head)

Usage in route handlers for a one-off snapshot::

    from musehub.debug.memory import rss_mb
    logger.info("rss after push: %.1f MiB", rss_mb())
"""

import logging
import os
import time
import tracemalloc
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from starlette.types import ASGIApp, Receive, Scope, Send
from musehub.types.json_types import MemoryFrame

logger = logging.getLogger(__name__)

try:
    import psutil as _psutil
    _PROC = _psutil.Process(os.getpid())

    def rss_mb() -> float:
        """Return current RSS in MiB."""
        return float(_PROC.memory_info().rss) / 1024 / 1024

except ImportError:  # pragma: no cover — psutil optional in test envs
    def rss_mb() -> float:  # noqa: F811
        return -1.0

@asynccontextmanager
async def profile_task(name: str) -> AsyncGenerator[None, None]:
    """Log RSS before and after an async block.

    Logs at INFO so it always appears in production logs.  If the delta
    exceeds 50 MiB it is also logged at WARNING so it's easy to grep.

    Example log output::

        [memory] START symbol-index repo=abc  rss=210.3 MiB
        [memory] END   symbol-index repo=abc  rss=310.7 MiB  delta=+100.4 MiB  elapsed=2.14s  ← WARNING
    """
    rss_before = rss_mb()
    t0 = time.monotonic()
    try:
        yield
    finally:
        elapsed = time.monotonic() - t0
        rss_after = rss_mb()
        delta = rss_after - rss_before
        sign = "+" if delta >= 0 else ""
        msg = (
            f"[memory] END   {name}  "
            f"rss={rss_after:.1f} MiB  delta={sign}{delta:.1f} MiB  elapsed={elapsed:.2f}s"
        )
        if abs(delta) >= 50:
            logger.warning(msg)
        else:
            logger.info(msg)
        logger.info(
            "[memory] START %s  rss=%.1f MiB",
            name, rss_before,
        )

def top_allocations(limit: int = 20) -> list[MemoryFrame]:
    """Return the top *limit* tracemalloc frames by cumulative size.

    Returns an empty list if tracemalloc was not started.
    """
    if not tracemalloc.is_tracing():
        return []
    snapshot = tracemalloc.take_snapshot()
    stats = snapshot.statistics("lineno")
    result = []
    for stat in stats[:limit]:
        frame = stat.traceback[0]
        result.append({
            "file": frame.filename,
            "line": frame.lineno,
            "size_kb": round(stat.size / 1024, 1),
            "count": stat.count,
        })
    return result

class MemoryLogMiddleware:
    """ASGI middleware that logs RSS after every response.

    Only logs when RSS exceeds ``warn_above_mb`` (default 400 MiB) so that
    normal traffic doesn't flood the logs.  Set ``warn_above_mb=0`` to log
    every request (noisy but useful during a profiling session).
    """

    def __init__(self, app: ASGIApp, warn_above_mb: float = 400.0) -> None:
        self.app = app
        self.warn_above_mb = warn_above_mb

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] not in ("http", "websocket"):
            await self.app(scope, receive, send)
            return

        await self.app(scope, receive, send)

        rss = rss_mb()
        if rss >= self.warn_above_mb:
            path = scope.get("path", "?")
            method = scope.get("method", "?")
            logger.warning(
                "[memory] HIGH RSS %.1f MiB after %s %s",
                rss, method, path,
            )