"""
Memory management module to prevent OOM (Out of Memory) on the host machine.

Monitors Docker container memory usage relative to overall machine available memory
and takes escalating actions:

    1. >= 80% machine RAM usage  ->  Flush Redis containers (FLUSHALL)
    2. >= 90% machine RAM usage  ->  Restart the deployment container using the most RAM
    3. Still >= 90% after restart ->  Kill the top RAM-consuming container
"""

import logging
import os
import re
import subprocess
import threading
import time
from typing import Any, Dict, List, Optional, Tuple

import docker
import psutil
from docker.models.containers import Container as DockerContainer

from matrice_common.utils import log_errors, send_error_log
from matrice_compute.scaling import Scaling

logger = logging.getLogger(__name__)


class MemoryManager:
    """Prevents OOM by monitoring machine memory and managing Docker containers.

    Escalation strategy:
        - >= 80%: Flush all Redis container memory (FLUSHALL / FLUSHDB)
        - >= 90%: Restart the deployment Docker container that consumes the most RAM
        - Still >= 90% after restart grace period: Kill the top RAM-consuming container

    The manager runs in a daemon thread, polling at a configurable interval.
    """

    # ── thresholds (fraction of total machine RAM) ──────────────────────
    FLUSH_THRESHOLD: float = 0.80
    RESTART_THRESHOLD: float = 0.90
    KILL_THRESHOLD: float = 0.90  # applied *after* a restart was already attempted

    # ── timing ──────────────────────────────────────────────────────────
    CHECK_INTERVAL: int = 15          # seconds between each memory check
    POST_ACTION_COOLDOWN: int = 30    # seconds to wait after taking an action
    RESTART_GRACE_PERIOD: int = 45    # seconds to wait after restart before considering kill
    REDIS_FLUSH_TIMEOUT: int = 10     # seconds timeout for redis-cli commands

    # ── Infrastructure action types (must NOT be killed/restarted) ───
    INFRASTRUCTURE_ACTION_TYPES: frozenset = frozenset({
        "database_setup",
        "kafka_setup",
        "redis_setup",
    })

    _OBJECTID_RE = re.compile(r'^[a-f0-9]{24}$')

    def __init__(
        self,
        flush_threshold: float = 0.80,
        restart_threshold: float = 0.90,
        check_interval: int = 15,
        scaling: Optional[Scaling] = None,
    ) -> None:
        """Initialize MemoryManager.

        Args:
            flush_threshold: Machine RAM usage fraction to trigger Redis flush (default 0.80).
            restart_threshold: Machine RAM usage fraction to trigger container restart (default 0.90).
            check_interval: Seconds between memory checks (default 15).
            scaling: Optional Scaling instance for reporting actions to the platform.
        """
        self.FLUSH_THRESHOLD = flush_threshold
        self.RESTART_THRESHOLD = restart_threshold
        self.KILL_THRESHOLD = restart_threshold  # kill uses same threshold as restart
        self.CHECK_INTERVAL = check_interval

        self._stop_event = threading.Event()
        self._monitor_thread: Optional[threading.Thread] = None
        self._is_running = False
        self._docker_client: Optional[docker.DockerClient] = None

        # Track escalation state so we only attempt restart once before escalating to kill
        self._restart_attempted_container_id: Optional[str] = None
        self._restart_attempted_at: float = 0.0

        # Scaling instance for platform reporting (optional)
        self._scaling = scaling

        logger.info(
            "MemoryManager initialised: flush_threshold=%.0f%%, restart_threshold=%.0f%%, interval=%ds, platform_reporting=%s",
            self.FLUSH_THRESHOLD * 100,
            self.RESTART_THRESHOLD * 100,
            self.CHECK_INTERVAL,
            self._scaling is not None,
        )

    # ── Docker client (lazy init) ──────────────────────────────────────
    def _get_docker_client(self) -> docker.DockerClient:
        """Return (and lazily create) the Docker client."""
        if self._docker_client is None:
            self._docker_client = docker.from_env()
        return self._docker_client

    # ── Machine-level metrics ──────────────────────────────────────────
    @staticmethod
    def get_machine_memory_usage() -> Tuple[float, float, float]:
        """Return machine memory metrics.

        Returns:
            Tuple of (usage_fraction, used_gb, total_gb).
        """
        mem = psutil.virtual_memory()
        total_gb = mem.total / (1024 ** 3)
        used_gb = mem.used / (1024 ** 3)
        usage_fraction = mem.percent / 100.0
        return usage_fraction, used_gb, total_gb

    # ── Container memory helpers ───────────────────────────────────────
    @log_errors(default_return=[], raise_exception=False)
    def _get_running_containers(self) -> List[DockerContainer]:
        """Return all running Docker containers."""
        client = self._get_docker_client()
        return client.containers.list(filters={"status": "running"})

    @log_errors(default_return=0.0, raise_exception=False)
    def _get_container_memory_mb(self, container: DockerContainer) -> float:
        """Return memory usage of a single container in MB.

        Uses the Docker stats API (single-shot, non-streaming).
        """
        try:
            stats = container.stats(stream=False)
            memory_usage_bytes = stats.get("memory_stats", {}).get("usage", 0)
            cache_bytes = stats.get("memory_stats", {}).get("stats", {}).get("cache", 0)
            return max(0.0, (memory_usage_bytes - cache_bytes) / (1024 * 1024))
        except Exception as exc:
            logger.debug("Error getting memory stats for container %s: %s", container.name, exc)
            return 0.0

    @log_errors(default_return=[], raise_exception=False)
    def _get_containers_sorted_by_memory(self) -> List[Dict[str, Any]]:
        """Return running containers sorted by memory usage (descending).

        Each entry: {"container": DockerContainer, "name": str, "memory_mb": float}
        """
        containers = self._get_running_containers()
        container_stats: List[Dict[str, Any]] = []

        for container in containers:
            memory_mb = self._get_container_memory_mb(container)
            container_stats.append({
                "container": container,
                "name": container.name,
                "id": container.id,
                "memory_mb": memory_mb,
            })

        container_stats.sort(key=lambda c: c["memory_mb"], reverse=True)
        return container_stats

    # ── Container name parsing and classification ────────────────────
    @staticmethod
    def _parse_container_name(container_name: str) -> Tuple[Optional[str], Optional[str]]:
        """Parse a container name to extract (action_record_id, action_type).

        Container naming conventions:
            - Deployment: ``{action_record_id}_{action_type}``
            - Redis sidecar: ``{action_record_id}_{action_type}_redis_container``

        Args:
            container_name: Docker container name.

        Returns:
            Tuple of (action_record_id, action_type) or (None, None) if unparseable.
        """
        if not container_name:
            return None, None

        parts = container_name.split("_", 1)
        if len(parts) < 2:
            return None, None

        candidate_id = parts[0]
        if not MemoryManager._OBJECTID_RE.match(candidate_id):
            return None, None

        remainder = parts[1]
        redis_suffix = "_redis_container"
        if remainder.endswith(redis_suffix):
            remainder = remainder[: -len(redis_suffix)]

        if not remainder:
            return None, None

        return candidate_id, remainder

    def _is_infrastructure_container(self, container: DockerContainer) -> bool:
        """Determine whether a container is infrastructure that must NOT be killed.

        Infrastructure containers have action_type in INFRASTRUCTURE_ACTION_TYPES.
        Containers with unparseable names are also protected (err on the side of caution).
        """
        container_name = container.name or ""
        _, action_type = self._parse_container_name(container_name)

        if action_type is None:
            return True

        return action_type in self.INFRASTRUCTURE_ACTION_TYPES

    # ── Redis detection and flushing ───────────────────────────────────
    @staticmethod
    def _is_redis_container(container: DockerContainer) -> bool:
        """Determine whether a container is a Redis instance.

        Checks both the image name and the container name for 'redis'.
        """
        image_name = ""
        try:
            image_tags = container.image.tags if container.image else []
            image_name = " ".join(image_tags).lower()
        except Exception:
            pass

        container_name = (container.name or "").lower()
        return "redis" in image_name or "redis" in container_name

    @log_errors(default_return=False, raise_exception=False)
    def _flush_redis_container(self, container: DockerContainer) -> bool:
        """Flush all data in a Redis container using ``redis-cli FLUSHALL``.

        Attempts to detect the password from container environment or command args.

        Returns:
            True if flush succeeded, False otherwise.
        """
        container_name = container.name or container.short_id

        # Try to extract Redis password from container config
        password = self._get_redis_password(container)

        # Build redis-cli command as list to handle passwords with special chars
        flush_cmd = ["redis-cli"]
        if password:
            flush_cmd.extend(["-a", password])
        flush_cmd.extend(["FLUSHALL", "ASYNC"])

        try:
            exec_result = container.exec_run(flush_cmd, demux=True)
            exit_code = exec_result.exit_code
            stdout = (exec_result.output[0] or b"").decode("utf-8", errors="replace").strip() if exec_result.output else ""

            if exit_code == 0 and "OK" in stdout.upper():
                logger.info(
                    "MemoryManager: Successfully flushed Redis container '%s'",
                    container_name,
                )
                return True
            else:
                logger.warning(
                    "MemoryManager: Redis FLUSHALL on '%s' returned exit_code=%s, output='%s'",
                    container_name,
                    exit_code,
                    stdout,
                )
                return False
        except Exception as exc:
            logger.error(
                "MemoryManager: Failed to exec redis-cli FLUSHALL on '%s': %s",
                container_name,
                exc,
            )
            return False

    @staticmethod
    def _get_redis_password(container: DockerContainer) -> Optional[str]:
        """Try to extract the Redis password from container environment / command.

        Looks for --requirepass in CMD, or REDIS_PASSWORD env var.
        """
        try:
            inspect = container.attrs or {}

            # Check environment variables
            env_list = inspect.get("Config", {}).get("Env", []) or []
            for env_entry in env_list:
                if env_entry.startswith("REDIS_PASSWORD="):
                    return env_entry.split("=", 1)[1]

            # Check command line args for --requirepass
            cmd = inspect.get("Config", {}).get("Cmd", []) or []
            for i, arg in enumerate(cmd):
                if arg == "--requirepass" and i + 1 < len(cmd):
                    return cmd[i + 1]

        except Exception as exc:
            logger.debug("Error extracting Redis password: %s", exc)

        return None

    # ── Restart / Kill helpers ─────────────────────────────────────────
    @log_errors(default_return=False, raise_exception=False)
    def _restart_container(self, container: DockerContainer) -> bool:
        """Restart a Docker container with a 30-second timeout.

        Returns:
            True if restart completed without error.
        """
        container_name = container.name or container.short_id
        logger.warning(
            "MemoryManager: Restarting container '%s' (id=%s) to reclaim memory",
            container_name,
            container.short_id,
        )
        try:
            container.restart(timeout=30)
            logger.info(
                "MemoryManager: Container '%s' restarted successfully",
                container_name,
            )
            return True
        except Exception as exc:
            logger.error(
                "MemoryManager: Failed to restart container '%s': %s",
                container_name,
                exc,
            )
            return False

    @log_errors(default_return=False, raise_exception=False)
    def _kill_container(self, container: DockerContainer) -> bool:
        """Force-kill a Docker container (SIGKILL).

        Returns:
            True if kill succeeded.
        """
        container_name = container.name or container.short_id
        logger.critical(
            "MemoryManager: KILLING container '%s' (id=%s) — memory still critical after restart",
            container_name,
            container.short_id,
        )
        try:
            container.kill()
            logger.info(
                "MemoryManager: Container '%s' killed",
                container_name,
            )
            return True
        except Exception as exc:
            logger.error(
                "MemoryManager: Failed to kill container '%s': %s",
                container_name,
                exc,
            )
            return False

    # ── Platform reporting helpers ─────────────────────────────────────
    def _report_container_restart(self, container_name: str) -> None:
        """Report a container restart to the platform via Kafka error_logs.

        Args:
            container_name: Name of the restarted container.
        """
        action_record_id, _ = self._parse_container_name(container_name)

        try:
            send_error_log(
                filename="memory_manager.py",
                function_name="_check_and_act",
                error_message=(
                    f"MemoryManager: Container '{container_name}' restarted due to "
                    f"critical memory pressure (>={self.RESTART_THRESHOLD * 100:.0f}%)"
                ),
                error_type="MEMORY_ERROR",
                service_name="py_compute",
                action_id=action_record_id,
            )
        except Exception as exc:
            logger.error("MemoryManager: Failed to send restart error log: %s", exc)

    def _report_container_kill(self, container_name: str) -> None:
        """Report a container kill (OOM) to the platform and mark the action as ERROR.

        Uses scaling.update_action + scaling.update_action_status to update the
        action to ERROR state, following the same pattern as _check_cuda in
        action_instance.py.

        Args:
            container_name: Name of the killed container.
        """
        action_record_id, action_type = self._parse_container_name(container_name)

        # Send to Kafka error_logs
        try:
            send_error_log(
                filename="memory_manager.py",
                function_name="_check_and_act",
                error_message=(
                    f"MemoryManager: Container '{container_name}' killed due to OOM — "
                    f"memory still critical after restart"
                ),
                error_type="MEMORY_ERROR",
                service_name="py_compute",
                action_id=action_record_id,
            )
        except Exception as exc:
            logger.error("MemoryManager: Failed to send kill error log: %s", exc)

        # Update action status to ERROR on the platform
        if action_record_id and action_type and self._scaling:
            try:
                self._scaling.update_action(
                    id=action_record_id,
                    step_code="ERROR",
                    action_type=action_type,
                    status="ERROR",
                    status_description=(
                        "MemoryManager: Container killed due to OOM — "
                        "memory still critical after restart"
                    ),
                    service="bg-job-scheduler",
                )
                self._scaling.update_action_status(
                    action_record_id=action_record_id,
                    status="ERROR",
                    isRunning=False,
                )
            except Exception as exc:
                logger.error(
                    "MemoryManager: Failed to update action status for %s: %s",
                    action_record_id,
                    exc,
                )

    # ── Top-level escalation logic ─────────────────────────────────────
    @log_errors(raise_exception=False)
    def _check_and_act(self) -> None:
        """Evaluate current memory and take the appropriate escalation action."""
        usage_fraction, used_gb, total_gb = self.get_machine_memory_usage()

        logger.info(
            "MemoryManager: Machine RAM %.1f%% (%.1f / %.1f GB)",
            usage_fraction * 100,
            used_gb,
            total_gb,
        )

        # ── Below 80% — nothing to do ──────────────────────────────────
        if usage_fraction < self.FLUSH_THRESHOLD:
            # Reset escalation state when memory is healthy
            self._restart_attempted_container_id = None
            self._restart_attempted_at = 0.0
            return

        # ── 80%–89% — Flush Redis containers ───────────────────────────
        if usage_fraction < self.RESTART_THRESHOLD:
            logger.warning(
                "MemoryManager: RAM at %.1f%% (>= %.0f%% threshold) — flushing Redis containers",
                usage_fraction * 100,
                self.FLUSH_THRESHOLD * 100,
            )
            self._flush_all_redis_containers()
            return

        # ── >= 90% — Restart or Kill ───────────────────────────────────
        logger.warning(
            "MemoryManager: RAM at %.1f%% (>= %.0f%% threshold) — escalating",
            usage_fraction * 100,
            self.RESTART_THRESHOLD * 100,
        )

        # First, still try to flush Redis if we haven't yet
        self._flush_all_redis_containers()

        # Get the container using the most RAM (deployment only — skip infrastructure)
        sorted_containers = self._get_containers_sorted_by_memory()
        if not sorted_containers:
            logger.warning("MemoryManager: No running containers found to restart/kill")
            return

        deployment_containers = [
            c for c in sorted_containers
            if not self._is_infrastructure_container(c["container"])
        ]

        if not deployment_containers:
            logger.warning(
                "MemoryManager: All top RAM consumers are infrastructure containers — cannot restart/kill"
            )
            return

        top_container_info = deployment_containers[0]
        top_container: DockerContainer = top_container_info["container"]
        top_name = top_container_info["name"]
        top_mem_mb = top_container_info["memory_mb"]

        logger.info(
            "MemoryManager: Top RAM consumer (deployment): '%s' using %.1f MB",
            top_name,
            top_mem_mb,
        )

        # Check if we already tried restarting this container recently
        if (
            self._restart_attempted_container_id == top_container.id
            and (time.time() - self._restart_attempted_at) < self.RESTART_GRACE_PERIOD
        ):
            # Already restarted this container and grace period not elapsed — wait
            logger.info(
                "MemoryManager: Restart grace period active for '%s' (%.0fs remaining)",
                top_name,
                self.RESTART_GRACE_PERIOD - (time.time() - self._restart_attempted_at),
            )
            return

        if (
            self._restart_attempted_container_id == top_container.id
            and (time.time() - self._restart_attempted_at) >= self.RESTART_GRACE_PERIOD
        ):
            # Restarted but still the top consumer after grace — KILL it
            logger.critical(
                "MemoryManager: Container '%s' still top consumer after restart — killing it",
                top_name,
            )
            self._kill_container(top_container)
            self._report_container_kill(top_name)
            # Reset escalation state
            self._restart_attempted_container_id = None
            self._restart_attempted_at = 0.0
            time.sleep(self.POST_ACTION_COOLDOWN)
            return

        # First attempt: restart the top consumer
        success = self._restart_container(top_container)
        if success:
            self._restart_attempted_container_id = top_container.id
            self._restart_attempted_at = time.time()
            self._report_container_restart(top_name)
        else:
            # Restart failed — kill immediately
            logger.warning(
                "MemoryManager: Restart failed for '%s', killing instead",
                top_name,
            )
            self._kill_container(top_container)
            self._report_container_kill(top_name)
            self._restart_attempted_container_id = None
            self._restart_attempted_at = 0.0

        time.sleep(self.POST_ACTION_COOLDOWN)

    @log_errors(raise_exception=False)
    def _flush_all_redis_containers(self) -> int:
        """Find and flush all running Redis containers.

        Returns:
            Number of containers successfully flushed.
        """
        containers = self._get_running_containers()
        flushed = 0

        for container in containers:
            if self._is_redis_container(container):
                logger.info(
                    "MemoryManager: Found Redis container '%s' — flushing",
                    container.name,
                )
                if self._flush_redis_container(container):
                    flushed += 1

        if flushed > 0:
            logger.info("MemoryManager: Flushed %d Redis container(s)", flushed)
        else:
            logger.debug("MemoryManager: No Redis containers found to flush")

        return flushed

    # ── Background monitor thread ──────────────────────────────────────
    def _monitor_worker(self) -> None:
        """Background thread loop that periodically checks memory and acts."""
        logger.info("MemoryManager monitor thread started (interval=%ds)", self.CHECK_INTERVAL)

        while not self._stop_event.is_set():
            try:
                self._check_and_act()
            except Exception as exc:
                logger.error("MemoryManager: Unexpected error in monitor loop: %s", exc)

            # Sleep in small increments so we can stop quickly
            if self._stop_event.wait(self.CHECK_INTERVAL):
                break

        # Cleanup
        if self._docker_client:
            try:
                self._docker_client.close()
            except Exception as exc:
                logger.debug("MemoryManager: Error closing Docker client: %s", exc)

        logger.info("MemoryManager monitor thread stopped")

    # ── Public start / stop API ────────────────────────────────────────
    @log_errors(raise_exception=False, log_error=True)
    def start(self) -> bool:
        """Start the memory manager monitoring thread.

        Returns:
            True if started successfully.
        """
        if self._is_running:
            logger.warning("MemoryManager is already running")
            return False

        self._stop_event.clear()
        self._monitor_thread = threading.Thread(
            target=self._monitor_worker,
            daemon=True,
            name="MemoryManager",
        )
        self._monitor_thread.start()
        self._is_running = True

        logger.info("MemoryManager started")
        return True

    @log_errors(raise_exception=False, log_error=True)
    def stop(self, timeout: int = 10) -> bool:
        """Stop the memory manager monitoring thread gracefully.

        Args:
            timeout: Maximum seconds to wait for the thread to exit.

        Returns:
            True if stopped successfully.
        """
        if not self._is_running:
            logger.warning("MemoryManager is not running")
            return False

        logger.info("Stopping MemoryManager...")
        self._stop_event.set()

        if self._monitor_thread and self._monitor_thread.is_alive():
            self._monitor_thread.join(timeout=timeout)
            if self._monitor_thread.is_alive():
                logger.error("MemoryManager thread did not stop within %ds timeout", timeout)
                return False

        self._is_running = False
        logger.info("MemoryManager stopped")
        return True

    def is_running(self) -> bool:
        """Check if the memory manager is currently running."""
        return self._is_running

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()
