# Copyright (C) 2025-2026 AKIOUD AI, SAS <contact@akioud.ai>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""
Sandbox quotas enforcement for AKIOS

Enforce hard resource limits using kernel-level cgroups v2.
Real cgroup-based resource isolation for strong security.
"""

import os
import resource
from pathlib import Path
from typing import Dict, Any, Optional

from ...config import get_settings

_docker_notice_emitted = False
_limits_notice_emitted = False


class QuotaViolationError(Exception):
    """Raised when resource quotas are exceeded"""

    def __init__(self, resource_type: str, actual_value: float, limit_value: float, pid: int = None):
        self.resource_type = resource_type
        self.actual_value = actual_value
        self.limit_value = limit_value
        self.pid = pid

        message = f"{resource_type.capitalize()} quota exceeded: {actual_value:.2f} > {limit_value:.2f}"
        if pid:
            message += f" (PID: {pid})"

        super().__init__(message)


class ResourceQuotas:
    """
    Resource quota enforcement using cgroups v2

    Provides hard limits on CPU, memory, disk I/O, and network bandwidth.
    Uses real cgroups v2 for kernel-level resource isolation.
    """

    def __init__(self):
        self.settings = get_settings()
        self.applied_limits = {}
        self.cgroup_base = Path('/sys/fs/cgroup')
        self.cgroup_mount = self._find_cgroup_mount()

    def _find_cgroup_mount(self) -> Optional[Path]:
        """Find the cgroup v2 mount point"""
        try:
            # Check if cgroup v2 is mounted
            with open('/proc/mounts', 'r') as f:
                for line in f:
                    parts = line.split()
                    if len(parts) >= 3 and parts[2] == 'cgroup2':
                        return Path(parts[1])
        except Exception:
            pass
        return None

    def _cgroups_available(self) -> bool:
        """Check if cgroups v2 is available and accessible"""
        if not self.cgroup_mount:
            return False

        # Check if we can write to cgroup filesystem
        try:
            test_file = self.cgroup_base / '.akios_test'
            test_file.write_text('test')
            test_file.unlink()
            return True
        except (OSError, PermissionError):
            return False

    def _create_cgroup(self, name: str) -> Path:
        """Create a new cgroup"""
        cgroup_path = self.cgroup_base / name
        try:
            cgroup_path.mkdir(parents=True, exist_ok=True)
            return cgroup_path
        except (OSError, PermissionError) as e:
            raise QuotaViolationError("Failed to create cgroup", 0, 0) from e

    def _write_cgroup_file(self, cgroup_path: Path, filename: str, value: str) -> None:
        """Write to a cgroup control file"""
        file_path = cgroup_path / filename
        try:
            file_path.write_text(value)
        except (OSError, PermissionError) as e:
            raise QuotaViolationError(f"Failed to write cgroup file {filename}", 0, 0) from e

    def _move_process_to_cgroup(self, pid: int, cgroup_path: Path) -> None:
        """Move a process to a cgroup"""
        try:
            procs_file = cgroup_path / 'cgroup.procs'
            procs_file.write_text(str(pid))
        except (OSError, PermissionError) as e:
            raise QuotaViolationError(f"Failed to move PID {pid} to cgroup", 0, 0, pid) from e

    def apply_cpu_quota(self) -> None:
        """Apply CPU time limits"""
        if not self.settings.sandbox_enabled:
            return

        cpu_seconds = 300  # 5 minutes max CPU time
        try:
            # Get current limits with threading-based timeout to avoid hangs in cgroups
            import threading
            
            result = [None]
            exception = [None]
            
            def get_rlimit():
                try:
                    result[0] = resource.getrlimit(resource.RLIMIT_CPU)
                except Exception as e:
                    exception[0] = e
            
            thread = threading.Thread(target=get_rlimit)
            thread.start()
            thread.join(timeout=1.0)  # 1 second timeout
            
            if thread.is_alive():
                # Thread is still running, rlimit call is hanging
                # Skip setting limits to avoid hanging
                self.applied_limits['cpu'] = float('inf')
                return
            
            if exception[0]:
                raise exception[0]
            
            current_soft, current_hard = result[0]

            # Only apply if our limit is more restrictive than current
            if cpu_seconds < current_soft:
                resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, min(cpu_seconds * 2, current_hard)))
                self.applied_limits['cpu'] = cpu_seconds
            else:
                # Keep track of effective limit (current limit)
                self.applied_limits['cpu'] = current_soft
        except (OSError, ValueError) as e:
            # If we can't set limits, log but don't fail - some systems don't allow this
            self.applied_limits['cpu'] = float('inf')  # Indicate unlimited

    def apply_memory_quota(self) -> None:
        """Apply memory limits"""
        if not self.settings.sandbox_enabled:
            return

        memory_bytes = self.settings.memory_limit_mb * 1024 * 1024
        try:
            # Get current limits with threading-based timeout to avoid hangs in cgroups
            import threading
            
            result = [None]
            exception = [None]
            
            def get_rlimit():
                try:
                    result[0] = resource.getrlimit(resource.RLIMIT_AS)
                except Exception as e:
                    exception[0] = e
            
            thread = threading.Thread(target=get_rlimit)
            thread.start()
            thread.join(timeout=1.0)  # 1 second timeout
            
            if thread.is_alive():
                # Thread is still running, rlimit call is hanging
                # Skip setting limits to avoid hanging
                self.applied_limits['memory'] = float('inf')
                return
            
            if exception[0]:
                raise exception[0]
            
            current_soft, current_hard = result[0]

            # Only apply if our limit is more restrictive than current
            if memory_bytes < current_soft:
                resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, min(memory_bytes, current_hard)))
                self.applied_limits['memory'] = memory_bytes
            else:
                # Keep track of effective limit (current limit)
                self.applied_limits['memory'] = current_soft
        except (OSError, ValueError) as e:
            # If we can't set limits, log but don't fail - some systems don't allow this
            self.applied_limits['memory'] = float('inf')  # Indicate unlimited

    def apply_file_descriptor_quota(self) -> None:
        """Apply file descriptor limits"""
        if not self.settings.sandbox_enabled:
            return

        max_fds = self.settings.max_open_files
        try:
            # Get current limits with threading-based timeout to avoid hangs in cgroups
            import threading
            
            result = [None]
            exception = [None]
            
            def get_rlimit():
                try:
                    result[0] = resource.getrlimit(resource.RLIMIT_NOFILE)
                except Exception as e:
                    exception[0] = e
            
            thread = threading.Thread(target=get_rlimit)
            thread.start()
            thread.join(timeout=1.0)  # 1 second timeout
            
            if thread.is_alive():
                # Thread is still running, rlimit call is hanging
                # Skip setting limits to avoid hanging
                self.applied_limits['file_descriptors'] = float('inf')
                return
            
            if exception[0]:
                raise exception[0]
            
            current_soft, current_hard = result[0]

            # Only apply if our limit is more restrictive than current
            if max_fds < current_soft:
                resource.setrlimit(resource.RLIMIT_NOFILE, (max_fds, min(max_fds, current_hard)))
                self.applied_limits['file_descriptors'] = max_fds
            else:
                # Keep track of effective limit (current limit)
                self.applied_limits['file_descriptors'] = current_soft
        except (OSError, ValueError) as e:
            # If we can't set limits, log but don't fail - some systems don't allow this
            self.applied_limits['file_descriptors'] = float('inf')  # Indicate unlimited

    def apply_to_process(self, pid: int, cpu_percent: float = 100, memory_mb: float = 1024,
                        max_files: int = 1024) -> str:
        """
        Apply resource quotas to a process using cgroups v2

        Args:
            pid: Process ID to apply limits to
            cpu_percent: CPU usage limit as percentage (0-100)
            memory_mb: Memory limit in MB
            max_files: Maximum number of open files

        Returns:
            Cgroup name that was created

        Raises:
            QuotaViolationError: If cgroup operations fail
        """
        if not self._cgroups_available():
            # Fallback to POSIX limits if cgroups not available
            self._apply_posix_limits(pid, cpu_percent, memory_mb, max_files)
            return "posix_fallback"

        cgroup_name = f"akios_process_{pid}"
        cgroup_path = self._create_cgroup(cgroup_name)

        try:
            # Move process to cgroup
            self._move_process_to_cgroup(pid, cgroup_path)

            # Apply CPU limits
            if cpu_percent < 100:
                # CPU limit as percentage: write to cpu.max
                # Format: $MAX $PERIOD
                # For 50% CPU: "50000 100000" (50ms per 100ms)
                cpu_quota = int(cpu_percent * 1000)  # Convert to microseconds per 1 second
                period = 1000000  # 1 second in microseconds
                cpu_max_value = f"{cpu_quota} {period}"
                self._write_cgroup_file(cgroup_path, 'cpu.max', cpu_max_value)

            # Apply memory limits
            memory_bytes = int(memory_mb * 1024 * 1024)
            self._write_cgroup_file(cgroup_path, 'memory.max', str(memory_bytes))

            # Apply file descriptor limits (approximated via pids.max and memory limits)
            # cgroups v2 doesn't directly limit file descriptors, but we can limit pids
            # which indirectly controls resource usage
            self._write_cgroup_file(cgroup_path, 'pids.max', str(min(max_files // 10, 1024)))

            # Store applied limits for tracking
            self.applied_limits[pid] = {
                'cpu_percent': cpu_percent,
                'memory_mb': memory_mb,
                'max_files': max_files,
                'cgroup': cgroup_name
            }

            return cgroup_name

        except Exception as e:
            # Cleanup on failure
            try:
                import shutil
                shutil.rmtree(cgroup_path, ignore_errors=True)
            except Exception:
                pass
            raise QuotaViolationError(f"Failed to apply cgroup limits to PID {pid}", 0, 0, pid) from e

    def _apply_posix_limits(self, pid: int, cpu_percent: float, memory_mb: float, max_files: int) -> None:
        """Fallback to POSIX resource limits if cgroups not available"""
        # Skip POSIX limits in Docker containers - Docker provides equivalent isolation
        # Also check DOCKER_CONTAINER env var for simulation/testing
        if os.path.exists('/.dockerenv') or os.getenv("DOCKER_CONTAINER") == "1":
            from ...core.ui.semantic_colors import print_info
            global _docker_notice_emitted
            if not _docker_notice_emitted:
                if os.getenv("AKIOS_DEBUG_ENABLED") == "1":
                    print_info("Docker container detected - using Docker's built-in resource limits")
                _docker_notice_emitted = True
            return

        try:
            # Set CPU time limit (seconds)
            cpu_seconds = int((cpu_percent / 100) * 300)  # Max 5 minutes scaled by percentage
            resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds))

            # Set memory limit
            memory_bytes = int(memory_mb * 1024 * 1024)
            resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes))

            # Set file descriptor limit
            resource.setrlimit(resource.RLIMIT_NOFILE, (max_files, max_files))

            # Disable core dumps
            resource.setrlimit(resource.RLIMIT_CORE, (0, 0))

        except (OSError, ValueError) as e:
            raise QuotaViolationError(f"Failed to apply POSIX limits to PID {pid}", 0, 0, pid) from e

    def apply_all_quotas(self) -> None:
        """Apply all resource quotas to the current process.

        Legacy method kept for backward compatibility — delegates to
        apply_to_process() with the current PID.
        """
        import os
        import logging
        logger = logging.getLogger(__name__)
        logger.debug("apply_all_quotas: delegating to apply_to_process(pid=%d)", os.getpid())
        try:
            self.apply_to_process(os.getpid())
        except Exception as e:
            logger.warning("apply_all_quotas: quota enforcement failed: %s", e)

    def check_resource_usage(self) -> Dict[str, Any]:
        """
        Check current resource usage against limits

        Returns:
            Dict with current usage and limit status
        """
        try:
            import psutil
            process = psutil.Process()
            memory_usage = process.memory_info().rss
            cpu_times = process.cpu_times()

            return {
                'memory_usage_bytes': memory_usage,
                'memory_limit_bytes': self.applied_limits.get('memory', 0),
                'cpu_time_user': cpu_times.user,
                'cpu_time_system': cpu_times.system,
                'cpu_limit_seconds': self.applied_limits.get('cpu', 0),
                'file_descriptors': len(process.open_files()),
                'file_descriptor_limit': self.applied_limits.get('file_descriptors', 0)
            }
        except (psutil.Error, OSError):
            # If we can't check resources, return empty dict
            return {}

    def enforce_quotas(self) -> None:
        """
        Enforce resource quotas - fail hard if exceeded

        This implements the "fail if exceeded" requirement from scope.
        Optimized for Docker environments where psutil checks can be slow.
        """
        # In Docker containers, skip expensive resource checks to avoid performance issues
        # Docker already provides process isolation, so we rely on cgroup limits instead
        if os.path.exists('/.dockerenv'):
            return  # Skip resource checking in Docker for performance
        usage = self.check_resource_usage()

        if not usage:
            return  # Can't check, assume OK

        # Check memory usage
        memory_limit = usage.get('memory_limit_bytes', 0)
        if memory_limit > 0 and usage.get('memory_usage_bytes', 0) > memory_limit:
            raise QuotaViolationError(
                f"Memory quota exceeded: {usage['memory_usage_bytes']} > {memory_limit}"
            )

        # Check CPU time (approximate)
        cpu_limit = usage.get('cpu_limit_seconds', 0)
        total_cpu = usage.get('cpu_time_user', 0) + usage.get('cpu_time_system', 0)
        if cpu_limit > 0 and total_cpu > cpu_limit:
            raise QuotaViolationError(
                f"CPU quota exceeded: {total_cpu}s > {cpu_limit}s"
            )

        # Check file descriptors
        fd_limit = usage.get('file_descriptor_limit', 0)
        if fd_limit > 0 and usage.get('file_descriptors', 0) > fd_limit:
            raise QuotaViolationError(
                f"File descriptor quota exceeded: {usage['file_descriptors']} > {fd_limit}"
            )


def enforce_hard_limits(pid: int = None, cpu_percent: float = 100, memory_mb: float = 1024,
                       max_files: int = 1024) -> None:
    """
    Enforce hard resource limits on current process using cgroups v2

    This is the main entry point for quota enforcement.
    Called by sandbox manager to set up process constraints.

    Args:
        pid: Process ID (default: current process)
        cpu_percent: CPU limit as percentage (0-100)
        memory_mb: Memory limit in MB
        max_files: Maximum open files
    """
    if pid is None:
        pid = os.getpid()

    quotas = ResourceQuotas()
    cgroup_name = quotas.apply_to_process(pid, cpu_percent, memory_mb, max_files)

    global _limits_notice_emitted
    if not _limits_notice_emitted:
        from ...core.ui.semantic_colors import print_success
        if os.getenv("AKIOS_DEBUG_ENABLED") == "1":
            print_success(f"Applied resource limits to PID {pid} using cgroup '{cgroup_name}'")
        _limits_notice_emitted = True

    # Immediately enforce limits
    quotas.enforce_quotas()
