"""
Memory profiling — track memory allocations per function.

Usage:
    from sondeo import memprof

    @memprof
    def allocate_stuff():
        data = [0] * 1_000_000
        return data
"""

from __future__ import annotations

import functools
import time
import tracemalloc
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, TypeVar

from sondeo.profiler import logging as logging_mod

F = TypeVar("F", bound=Callable[..., Any])


@dataclass
class MemoryResult:
    """Result from a memory measurement."""

    name: str
    peak_mb: float
    current_mb: float
    allocated_mb: float
    top_allocations: list[tuple[str, float]]  # (traceback_str, size_mb)
    started_at: float | None = None
    ended_at: float | None = None
    correlation_id: str | None = None

    def __repr__(self) -> str:
        return (
            f"MemoryResult({self.name!r}, peak={self.peak_mb:.2f}MB, "
            f"alloc={self.allocated_mb:.2f}MB)"
        )


_results: list[MemoryResult] = []


def get_results() -> list[MemoryResult]:
    """Get all collected memory results."""
    return list(_results)


def clear_results() -> None:
    """Clear all collected memory results."""
    _results.clear()


def memprof(
    fn: F | None = None,
    *,
    top_n: int = 10,
    trace_depth: int = 3,
) -> F | Callable[[F], F]:
    """Decorator that profiles memory usage of a function.

    Uses tracemalloc to snapshot memory before and after execution.

    Args:
        fn: The function to profile.
        top_n: Number of top allocations to capture.
        trace_depth: Traceback depth for allocation tracking.
    """

    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            parent_cid = logging_mod.current_correlation_id()
            with logging_mod.correlation_scope() as cid:
                was_tracing = tracemalloc.is_tracing()
                if not was_tracing:
                    tracemalloc.start(trace_depth)

                started_at = time.time()
                snapshot_before = tracemalloc.take_snapshot()

                logging_mod.emit_event(
                    event_type="memory",
                    level="debug",
                    module=func.__module__,
                    message=f"{func.__qualname__} memory start",
                    correlation_id=cid,
                    data={
                        "phase": "start",
                        "function": func.__qualname__,
                        "parent_correlation_id": parent_cid,
                    },
                )

                try:
                    result = func(*args, **kwargs)
                finally:
                    snapshot_after = tracemalloc.take_snapshot()
                    current, peak = tracemalloc.get_traced_memory()

                    if not was_tracing:
                        tracemalloc.stop()

                    # Diff snapshots
                    stats = snapshot_after.compare_to(snapshot_before, "lineno")

                    top_allocs: list[tuple[str, float]] = []
                    for stat in stats[:top_n]:
                        top_allocs.append((str(stat), stat.size / (1024 * 1024)))

                    total_allocated = sum(s.size for s in stats if s.size_diff > 0) / (1024 * 1024)

                    mem_result = MemoryResult(
                        name=func.__qualname__,
                        peak_mb=peak / (1024 * 1024),
                        current_mb=current / (1024 * 1024),
                        allocated_mb=total_allocated,
                        started_at=started_at,
                        ended_at=time.time(),
                        correlation_id=cid,
                        top_allocations=top_allocs,
                    )
                    _results.append(mem_result)

                    logging_mod.emit_event(
                        event_type="memory",
                        level="info",
                        module=func.__module__,
                        message=f"{func.__qualname__} memory end",
                        correlation_id=cid,
                        data={
                            "phase": "end",
                            "function": func.__qualname__,
                            "peak_mb": round(mem_result.peak_mb, 3),
                            "allocated_mb": round(mem_result.allocated_mb, 3),
                            "parent_correlation_id": parent_cid,
                        },
                    )

            return result

        return wrapper  # type: ignore[return-value]

    if fn is not None:
        return decorator(fn)
    return decorator
