#!/usr/bin/env python3
"""tmpfs manager — RAM-backed mounts for sensitive operations.

Creates tmpfs mounts for working directories, runtime state, and SQLite WAL
files.  Power loss = instant evidence destruction since tmpfs is RAM-only.
Size limits are tier-aware (OPi Zero 3: 200MB, Pi Zero: 30MB, generic: 512MB).
"""

import logging
import os
import subprocess
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

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 tmpfs size limits per tier (MB)
_TIER_TMPFS_LIMITS = {
    TIER_OPI_ZERO3: 200,
    TIER_PI_ZERO: 30,
    TIER_GENERIC: 512,
}


class TmpfsManager(BaseModule):
    """Manage tmpfs mounts for RAM-only sensitive storage."""

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

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._mounts: Dict[str, dict] = {}  # mount_point -> {size_mb, purpose}
        self._install_path: str = ""
        self._tier_limit_mb: int = 200

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

    def start(self) -> None:
        if self._running:
            return

        self._install_path = self.config.get("device", {}).get(
            "install_path", "/opt/.cache/bb"
        )

        # Determine tier-based tmpfs limit
        tier = get_hardware_tier()
        hw_tiers = self.config.get("hardware_tiers", {})
        tier_cfg = hw_tiers.get(tier, {})
        self._tier_limit_mb = tier_cfg.get(
            "overlayfs_mb",
            _TIER_TMPFS_LIMITS.get(tier, 200),
        ) or _TIER_TMPFS_LIMITS.get(tier, 200)

        # Create and mount tmpfs volumes
        tmp_path = os.path.join(self._install_path, "tmp")
        run_path = os.path.join(self._install_path, "run")

        # Split budget: 70% for tmp (working data), 30% for run (runtime state)
        tmp_size = int(self._tier_limit_mb * 0.7)
        run_size = int(self._tier_limit_mb * 0.3)

        self._mount_tmpfs(tmp_path, tmp_size, "general tmpfs working directory")
        self._mount_tmpfs(run_path, run_size, "runtime state")

        # Create standard subdirectories
        for subdir in ("wal", "work", "modules"):
            os.makedirs(os.path.join(tmp_path, subdir), exist_ok=True)
        for subdir in ("pids", "sockets", "locks"):
            os.makedirs(os.path.join(run_path, subdir), exist_ok=True)

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self.state.set_module_status(self.name, "running", pid=os.getpid())

        logger.info(
            "TmpfsManager active — %d mounts, tier=%s, budget=%dMB",
            len(self._mounts), tier, self._tier_limit_mb,
        )

    def stop(self) -> None:
        if not self._running:
            return

        # Unmount in reverse order (most recently mounted first)
        for mount_point in reversed(list(self._mounts.keys())):
            self._unmount_tmpfs(mount_point)

        self._mounts.clear()
        self._running = False
        self.state.set_module_status(self.name, "stopped")
        logger.info("TmpfsManager stopped — all tmpfs mounts removed")

    def status(self) -> dict:
        mounts_info = []
        for mount_point, info in self._mounts.items():
            entry = {
                "mount_point": mount_point,
                "size_mb": info["size_mb"],
                "purpose": info["purpose"],
            }
            # Get current usage
            try:
                stat = os.statvfs(mount_point)
                total = stat.f_blocks * stat.f_frsize
                free = stat.f_bavail * stat.f_frsize
                used = total - free
                entry["used_mb"] = round(used / (1024 * 1024), 1)
                entry["free_mb"] = round(free / (1024 * 1024), 1)
                entry["used_pct"] = round((used / total) * 100, 1) if total > 0 else 0
            except OSError:
                entry["used_mb"] = 0
                entry["free_mb"] = 0
                entry["used_pct"] = 0
            mounts_info.append(entry)

        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "tier_limit_mb": self._tier_limit_mb,
            "mount_count": len(self._mounts),
            "mounts": mounts_info,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def get_tmp_path(self) -> Optional[str]:
        """Return the general tmpfs working directory, or None if not mounted."""
        path = os.path.join(self._install_path, "tmp")
        return path if path in self._mounts else None

    def get_run_path(self) -> Optional[str]:
        """Return the runtime state tmpfs path, or None if not mounted."""
        path = os.path.join(self._install_path, "run")
        return path if path in self._mounts else None

    def get_wal_path(self) -> Optional[str]:
        """Return the WAL file tmpfs directory for SQLite databases."""
        tmp = self.get_tmp_path()
        if tmp:
            wal_dir = os.path.join(tmp, "wal")
            return wal_dir if os.path.isdir(wal_dir) else None
        return None

    def get_module_workdir(self, module_name: str) -> Optional[str]:
        """Get or create a per-module working directory on tmpfs."""
        tmp = self.get_tmp_path()
        if not tmp:
            return None
        workdir = os.path.join(tmp, "modules", module_name)
        os.makedirs(workdir, exist_ok=True)
        return workdir

    def symlink_wal(self, db_path: str) -> bool:
        """Symlink a SQLite WAL file from its storage location to tmpfs.

        Call this after creating/opening a SQLite database so the WAL file
        lives in RAM instead of on disk.

        Args:
            db_path: Path to the .db file (the WAL will be at db_path + '-wal')

        Returns:
            True if the symlink was created successfully.
        """
        wal_dir = self.get_wal_path()
        if not wal_dir:
            return False

        wal_source = db_path + "-wal"
        db_name = os.path.basename(db_path)
        wal_target = os.path.join(wal_dir, db_name + "-wal")

        try:
            # Remove existing WAL if present
            if os.path.exists(wal_source) and not os.path.islink(wal_source):
                os.unlink(wal_source)

            # Create symlink: storage/foo.db-wal -> tmpfs/wal/foo.db-wal
            if not os.path.islink(wal_source):
                os.symlink(wal_target, wal_source)

            # Also handle the -shm file
            shm_source = db_path + "-shm"
            shm_target = os.path.join(wal_dir, db_name + "-shm")
            if os.path.exists(shm_source) and not os.path.islink(shm_source):
                os.unlink(shm_source)
            if not os.path.islink(shm_source):
                os.symlink(shm_target, shm_source)

            logger.debug("WAL symlinked: %s -> %s", wal_source, wal_target)
            return True
        except (IOError, OSError) as exc:
            logger.warning("Failed to symlink WAL for %s: %s", db_path, exc)
            return False

    # ------------------------------------------------------------------
    # Internal mount management
    # ------------------------------------------------------------------

    def _mount_tmpfs(self, mount_point: str, size_mb: int, purpose: str) -> bool:
        """Create a tmpfs mount at the given path."""
        try:
            os.makedirs(mount_point, exist_ok=True)

            # Check if already mounted
            if self._is_mounted(mount_point):
                logger.debug("Already mounted: %s", mount_point)
                self._mounts[mount_point] = {"size_mb": size_mb, "purpose": purpose}
                return True

            subprocess.run(
                ["mount", "-t", "tmpfs", "-o", f"size={size_mb}m,mode=0700,nodev,nosuid",
                 "tmpfs", mount_point],
                check=True, capture_output=True,
            )

            self._mounts[mount_point] = {"size_mb": size_mb, "purpose": purpose}
            logger.debug("Mounted tmpfs: %s (%dMB) — %s", mount_point, size_mb, purpose)
            return True
        except (subprocess.CalledProcessError, OSError) as exc:
            logger.error("Failed to mount tmpfs at %s: %s", mount_point, exc)
            return False

    def _unmount_tmpfs(self, mount_point: str) -> bool:
        """Unmount a tmpfs mount."""
        try:
            if self._is_mounted(mount_point):
                subprocess.run(
                    ["umount", "-l", mount_point],  # lazy unmount for safety
                    check=True, capture_output=True,
                )
                logger.debug("Unmounted tmpfs: %s", mount_point)

            # Clean up the directory
            try:
                os.rmdir(mount_point)
            except OSError:
                pass  # Directory may not be empty or may not exist

            return True
        except subprocess.CalledProcessError as exc:
            logger.warning("Failed to unmount %s: %s", mount_point, exc)
            return False

    @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:
                    fields = line.split()
                    if len(fields) >= 2 and fields[1] == path:
                        return True
        except IOError:
            pass
        return False
