#!/usr/bin/env python3
"""OverlayFS manager: mount root filesystem as overlayfs with tmpfs upper layer
so all writes go to RAM and vanish on reboot. Zero SD card write activity."""

import logging
import os
import shutil
import subprocess
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule
from utils.resource import get_hardware_tier, TIER_OPI_ZERO3, TIER_PI_ZERO, TIER_GENERIC

logger = logging.getLogger(__name__)

# Default overlay paths
OVERLAY_BASE = "/opt/.cache/bb/overlay"
OVERLAY_UPPER = os.path.join(OVERLAY_BASE, "upper")
OVERLAY_WORK = os.path.join(OVERLAY_BASE, "work")
OVERLAY_MERGED = os.path.join(OVERLAY_BASE, "merged")

# Size budgets per tier (MB)
TIER_SIZE_BUDGETS = {
    TIER_OPI_ZERO3: 200,
    TIER_PI_ZERO: 30,
    TIER_GENERIC: None,  # unlimited
}


class OverlayfsManager(BaseModule):
    """Mount an overlayfs with read-only lower (real FS) and tmpfs upper (RAM).
    All writes go to tmpfs -- zero SD card activity. Writes vanish on reboot."""

    name = "overlayfs_manager"
    module_type = "stealth"
    priority = -450
    requires_root = True

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._overlay_base = config.get("overlay_base", OVERLAY_BASE)
        self._upper_dir = config.get("overlay_upper", OVERLAY_UPPER)
        self._work_dir = config.get("overlay_work", OVERLAY_WORK)
        self._merged_dir = config.get("overlay_merged", OVERLAY_MERGED)
        self._lower_dir = config.get("overlay_lower", "/opt/.cache/bb")
        self._overlay_active = False
        self._tmpfs_mounted = False
        self._upper_limit_mb: Optional[int] = None
        self._tier = get_hardware_tier()

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        self._running = True
        self._start_time = time.time()
        self._pid = os.getpid()

        # Determine size budget from hardware tier
        self._upper_limit_mb = TIER_SIZE_BUDGETS.get(self._tier)
        if self._upper_limit_mb is None and self._tier not in TIER_SIZE_BUDGETS:
            # Unknown tier: default conservative
            self._upper_limit_mb = 200

        # Create overlay structure
        if not self._setup_overlay():
            logger.error("OverlayFS setup failed")
            self.state.set_module_status(self.name, "error",
                                         extra={"reason": "setup_failed"})
            return

        self.state.set_module_status(self.name, "running", pid=self._pid)
        self.bus.emit("MODULE_STARTED", {"module": self.name}, source_module=self.name)
        logger.info(
            "OverlayFS active — upper limit: %s MB, tier: %s",
            self._upper_limit_mb or "unlimited", self._tier,
        )

    def stop(self) -> None:
        self._running = False
        self._teardown_overlay()
        self.state.set_module_status(self.name, "stopped")
        self.bus.emit("MODULE_STOPPED", {"module": self.name}, source_module=self.name)
        logger.info("OverlayFS stopped and unmounted")

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "overlay_active": self._overlay_active,
            "tmpfs_mounted": self._tmpfs_mounted,
            "upper_usage_mb": self._get_upper_usage_mb(),
            "upper_limit_mb": self._upper_limit_mb,
            "tier": self._tier,
            "merged_dir": self._merged_dir,
        }

    def configure(self, config: dict) -> None:
        if "overlay_lower" in config:
            self._lower_dir = config["overlay_lower"]

    # ------------------------------------------------------------------
    # Public: overlay path helpers
    # ------------------------------------------------------------------

    @property
    def merged_path(self) -> str:
        """Return the merged overlay mount point (use this as working directory)."""
        return self._merged_dir

    def get_upper_usage_mb(self) -> float:
        """Return current tmpfs upper layer usage in MB."""
        return self._get_upper_usage_mb()

    def is_near_limit(self, threshold_pct: float = 85.0) -> bool:
        """Check if upper layer usage is approaching the size budget."""
        if self._upper_limit_mb is None:
            return False
        usage = self._get_upper_usage_mb()
        return (usage / self._upper_limit_mb * 100) >= threshold_pct

    # ------------------------------------------------------------------
    # Internal: setup and teardown
    # ------------------------------------------------------------------

    def _setup_overlay(self) -> bool:
        """Create tmpfs mount and overlayfs mount."""
        try:
            # Create all directories
            for d in [self._overlay_base, self._upper_dir, self._work_dir, self._merged_dir]:
                os.makedirs(d, mode=0o700, exist_ok=True)

            # Ensure lower directory exists
            os.makedirs(self._lower_dir, mode=0o700, exist_ok=True)

            # Mount tmpfs on the overlay base with size limit
            if not self._is_mounted(self._overlay_base):
                size_opt = f"size={self._upper_limit_mb}m" if self._upper_limit_mb else "size=50%"
                result = subprocess.run(
                    ["mount", "-t", "tmpfs", "-o", f"{size_opt},mode=0700",
                     "tmpfs", self._overlay_base],
                    capture_output=True, timeout=30,
                )
                if result.returncode != 0:
                    logger.error(
                        "tmpfs mount failed: %s",
                        result.stderr.decode(errors="replace"),
                    )
                    return False
                self._tmpfs_mounted = True
                # Recreate subdirectories on fresh tmpfs
                for d in [self._upper_dir, self._work_dir, self._merged_dir]:
                    os.makedirs(d, mode=0o700, exist_ok=True)
                logger.info("tmpfs mounted at %s (%s)", self._overlay_base, size_opt)
            else:
                self._tmpfs_mounted = True

            # Mount overlayfs
            if not self._is_mounted(self._merged_dir):
                mount_opts = (
                    f"lowerdir={self._lower_dir},"
                    f"upperdir={self._upper_dir},"
                    f"workdir={self._work_dir}"
                )
                result = subprocess.run(
                    ["mount", "-t", "overlay", "overlay",
                     "-o", mount_opts, self._merged_dir],
                    capture_output=True, timeout=30,
                )
                if result.returncode != 0:
                    stderr = result.stderr.decode(errors="replace")
                    logger.error("overlayfs mount failed: %s", stderr)
                    # Cleanup tmpfs if overlay fails
                    self._unmount(self._overlay_base)
                    self._tmpfs_mounted = False
                    return False
                self._overlay_active = True
                logger.info("overlayfs mounted at %s", self._merged_dir)
            else:
                self._overlay_active = True

            return True

        except Exception:
            logger.exception("Overlay setup failed")
            return False

    def _teardown_overlay(self) -> None:
        """Sync and unmount overlay, then tmpfs."""
        # Sync any pending writes
        try:
            subprocess.run(["sync"], timeout=10, capture_output=True)
        except Exception:
            pass

        # Unmount overlay first
        if self._overlay_active:
            if self._unmount(self._merged_dir):
                self._overlay_active = False
                logger.info("overlayfs unmounted")

        # Unmount tmpfs (all data in upper layer is destroyed)
        if self._tmpfs_mounted:
            if self._unmount(self._overlay_base):
                self._tmpfs_mounted = False
                logger.info("tmpfs unmounted — all overlay writes destroyed")

    # ------------------------------------------------------------------
    # Internal: mount helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _is_mounted(path: str) -> bool:
        """Check if a path is a mount point."""
        try:
            with open("/proc/mounts", "r") as f:
                for line in f:
                    parts = line.split()
                    if len(parts) >= 2 and parts[1] == path:
                        return True
        except IOError:
            pass
        return False

    @staticmethod
    def _unmount(path: str, lazy: bool = False) -> bool:
        """Unmount a filesystem. Uses lazy unmount as fallback."""
        try:
            cmd = ["umount", path]
            result = subprocess.run(cmd, capture_output=True, timeout=30)
            if result.returncode == 0:
                return True

            if not lazy:
                # Retry with lazy unmount
                cmd = ["umount", "-l", path]
                result = subprocess.run(cmd, capture_output=True, timeout=30)
                if result.returncode == 0:
                    return True

            logger.warning(
                "Failed to unmount %s: %s",
                path, result.stderr.decode(errors="replace"),
            )
            return False
        except Exception:
            logger.exception("Unmount failed: %s", path)
            return False

    # ------------------------------------------------------------------
    # Internal: usage tracking
    # ------------------------------------------------------------------

    def _get_upper_usage_mb(self) -> float:
        """Get current size of the tmpfs upper layer in MB."""
        if not self._tmpfs_mounted:
            return 0.0
        try:
            stat = os.statvfs(self._overlay_base)
            total = stat.f_blocks * stat.f_frsize
            free = stat.f_bavail * stat.f_frsize
            used = total - free
            return used / (1024 * 1024)
        except OSError:
            return 0.0
