"""
Optimus Research Center - Security & Execution Sandbox Module (Hardened v2.0)
Fulfills Defect 4: Execution Security & Denial Policy.

Separates execution into two tiers:
  A. Trusted internal operations (whitelisted algorithms and benchmarks).
  B. Untrusted / experimental code (requires verified OS-level container isolation).

Key Rules:
  - Canonical path resolution using os.path.commonpath and path components.
  - Subprocess or venv is NOT considered a secure sandbox for untrusted code.
  - If verified OS isolation is not present, untrusted code execution is blocked
    and returns SANDBOX_UNAVAILABLE (proposals preserved as PENDING_SAFE_ENVIRONMENT).
  - Internal runner (ProcessSandbox) executes verified internal scripts with timeout
    and memory limits, verifying path and code safety.
"""

import os
import ast
import time
import json
import subprocess
from pathlib import Path
from typing import Dict, Any, Optional, Set, List

from app.config import settings
from app.core.database import SessionLocal
from app.core.models import AuditLog


class SecurityViolation(Exception):
    """Raised when an operation violates sandbox boundaries or security rules."""
    pass


class SandboxUnavailableError(Exception):
    """Raised when untrusted code execution is requested without real OS sandbox."""
    pass


class CanonicalPathGuard:
    """Validates filesystem paths using canonical component resolution."""

    FORBIDDEN_COMPONENTS: Set[str] = {
        "digital optimus",
        "digitaloptimus",
        "digoptimus",
        "mininova",
        "mini optimus",
        "optimus developer",
        "cerebro"
    }

    @classmethod
    def assert_path_contained(cls, target_path: Path, allowed_root: Optional[Path] = None) -> Path:
        """
        Verifies that target_path is strictly within allowed_root using canonical path components.
        Does NOT rely on naive text substring matching.
        """
        root = (allowed_root or settings.BASE_DIR).resolve()
        target = Path(target_path).resolve()

        # 1. Canonical component check for forbidden system keywords
        for part in target.parts:
            lowered = part.lower().strip()
            if lowered in cls.FORBIDDEN_COMPONENTS:
                cls._log_security_event(
                    "FORBIDDEN_PATH_BLOCKED",
                    f"Path component '{part}' references forbidden system boundary. Target: {target}"
                )
                raise SecurityViolation(f"Access Denied: Path references prohibited component '{part}'")

        # 2. Strict component containment check via os.path.commonpath
        try:
            common = Path(os.path.commonpath([str(root), str(target)])).resolve()
        except ValueError:
            # Different drive letters or invalid roots
            cls._log_security_event("PATH_OUT_OF_BOUNDS", f"Target {target} is on a different drive than root {root}")
            raise SecurityViolation(f"Access Denied: Path {target} is completely outside research environment.")

        if common != root:
            cls._log_security_event("PATH_ESCAPE_BLOCKED", f"Target {target} escapes allowed root {root}")
            raise SecurityViolation(f"Access Denied: Path {target} is outside allowed root {root}")

        return target

    @classmethod
    def ensure_within_allowed(cls, target_path: Path, allowed_roots: Optional[List[Path]] = None) -> Path:
        """Ensures target path is within one of the allowed roots."""
        roots = allowed_roots or [settings.BASE_DIR]
        target = Path(target_path).resolve()

        # Component check first
        for part in target.parts:
            lowered = part.lower().strip()
            if lowered in cls.FORBIDDEN_COMPONENTS:
                cls._log_security_event("FORBIDDEN_PATH_BLOCKED", f"Path component '{part}' forbidden: {target}")
                raise SecurityViolation(f"Access Denied: Path references prohibited component '{part}'")

        for r in roots:
            resolved_root = r.resolve()
            try:
                common = Path(os.path.commonpath([str(resolved_root), str(target)])).resolve()
                if common == resolved_root:
                    return target
            except ValueError:
                continue

        cls._log_security_event("PATH_OUT_OF_BOUNDS", f"Target {target} outside allowed roots")
        raise SecurityViolation(f"Access Denied: Path {target} is outside allowed roots.")

    @classmethod
    def _log_security_event(cls, event_type: str, details: str):
        try:
            db = SessionLocal()
            log = AuditLog(
                event_type=event_type,
                severity="CRITICAL",
                details_json=json.dumps({"details": details, "timestamp": time.time()})
            )
            db.add(log)
            db.commit()
            db.close()
        except Exception:
            pass


class SandboxSecurityGuard:
    """High-level security guard for paths and code."""

    @classmethod
    def assert_path_safe(cls, target_path: Path, operation: str = "READ") -> Path:
        return CanonicalPathGuard.assert_path_contained(target_path)

    @classmethod
    def assert_code_safe(cls, code_string: str) -> bool:
        """Audits code string for hazardous imports, trading APIs, or forbidden keywords."""
        lowered = code_string.lower()

        # Keyword checks
        for forbidden in CanonicalPathGuard.FORBIDDEN_COMPONENTS:
            if forbidden in lowered:
                CanonicalPathGuard._log_security_event("FORBIDDEN_CODE_BLOCKED", f"Reference to '{forbidden}' in code.")
                raise SecurityViolation(f"Code references prohibited system keyword: '{forbidden}'")

        # Prohibited libraries and trading operations
        prohibited_tokens = [
            "os.system", "shutil.rmtree", "subprocess.popen", "format c:",
            "ccxt", "binance", "coinbase", "alpaca_trade_api", "create_order",
            "open_socket", "eval(", "exec("
        ]
        for token in prohibited_tokens:
            if token in lowered:
                CanonicalPathGuard._log_security_event("HAZARDOUS_CODE_BLOCKED", f"Prohibited token '{token}' detected.")
                raise SecurityViolation(f"Code contains prohibited operation: '{token}'")

        # AST syntax check
        try:
            tree = ast.parse(code_string)
            for node in ast.walk(tree):
                if isinstance(node, ast.Name) and node.id in {"eval", "exec", "__import__", "compile"}:
                    raise SecurityViolation(f"Prohibited dynamic AST symbol: '{node.id}'")
        except SyntaxError as e:
            raise SecurityViolation(f"Invalid Python syntax: {e}")

        return True


class ExecutionPolicy:
    """
    Separates execution into two distinct tiers:
    - CATEGORY A: Trusted internal algorithms and built-in pipelines.
    - CATEGORY B: Untrusted or dynamically submitted experimental code.
    """

    TRUSTED_OPERATIONS: Set[str] = {
        "ALGORITHM_BENCHMARK_OPTIMIZATION",
        "ALGORITHM_BENCHMARK_PREDICTIVE",
        "FINANCIAL_QUANT_RESEARCH",
        "TECH_LITERATURE_ANALYSIS",
        "TECH_EMPIRICAL_BENCHMARK",
        "AI_TRAINING_AND_EVALUATION",
        "AI_ARCHITECTURE_INVESTIGATION",
        "ROBOTICS_SIMULATION",
        "CROSS_LAB_COLLABORATION_EVALUATION"
    }

    @classmethod
    def is_trusted_operation(cls, op_type: str) -> bool:
        return op_type in cls.TRUSTED_OPERATIONS

    @classmethod
    def validate_category_b_code(cls, code_string: str) -> Dict[str, Any]:
        """Safely parses Category B untrusted code using Python AST without executing it."""
        try:
            tree = ast.parse(code_string)
            node_count = sum(1 for _ in ast.walk(tree))
        except SyntaxError as e:
            raise SecurityViolation(f"Invalid Python code syntax: {e}")

        for node in ast.walk(tree):
            if isinstance(node, ast.Name):
                if node.id in {"exec", "eval", "__import__", "compile"}:
                    raise SecurityViolation(f"Prohibited dynamic AST symbol: '{node.id}'")

        return {
            "syntax_valid": True,
            "ast_nodes": node_count,
            "status": "PENDING_SAFE_ENVIRONMENT",
            "message": "Code parsed and verified safe for storage. Execution requires verified OS container sandbox."
        }

    @classmethod
    def execute_untrusted_code(
        cls,
        code_string: str,
        action_name: str = "untrusted_execution",
        raise_exception: bool = False
    ) -> Dict[str, Any]:
        """
        STRICT POLICY: Category B untrusted code execution requires verified OS-level container isolation.
        Because a Windows Container / AppContainer sandbox is not currently provisioned,
        this operation is UNCONDITIONALLY BLOCKED and reports SANDBOX_UNAVAILABLE.
        The proposal is safely preserved as PENDING_SAFE_ENVIRONMENT.
        """
        ast_info = cls.validate_category_b_code(code_string)

        CanonicalPathGuard._log_security_event(
            "UNTRUSTED_CODE_EXECUTION_BLOCKED",
            f"Action '{action_name}': Untrusted Category B code blocked due to SANDBOX_UNAVAILABLE policy."
        )

        result = {
            "status": "SANDBOX_UNAVAILABLE",
            "action_name": action_name,
            "reason": (
                "SANDBOX_UNAVAILABLE: Real OS-level isolation (Windows Containers/AppContainer) is not provisioned. "
                "Untrusted experimental code execution is strictly blocked by policy. "
                "Proposal safely preserved as PENDING_SAFE_ENVIRONMENT."
            ),
            "candidate_status": "PENDING_SAFE_ENVIRONMENT",
            "ast_audit": ast_info
        }

        if raise_exception:
            raise SandboxUnavailableError(result["reason"])

        return result


class ProcessSandbox:
    """
    Subprocess execution harness for internal, whitelisted scripts (Category A).
    Applies timeout, peak memory limit enforcement, environment scrubbing, and path containment checks.
    NOTE: This is an internal subprocess harness, NOT an OS-level container sandbox for untrusted code.
    """

    SENSITIVE_ENV_TOKENS = (
        "API_KEY", "SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "BROKER", "PRIVATE_KEY"
    )

    def __init__(self, timeout_seconds: int = 10, max_memory_mb: int = 512):
        self.timeout_seconds = timeout_seconds
        self.max_memory_mb = max_memory_mb

    @classmethod
    def build_sanitized_env(cls) -> Dict[str, str]:
        """Constructs a sanitized environment dictionary stripped of credentials and external tokens."""
        clean_env: Dict[str, str] = {}
        for k, v in os.environ.items():
            k_upper = k.upper()
            if any(tok in k_upper for tok in cls.SENSITIVE_ENV_TOKENS):
                continue
            clean_env[k] = v
        clean_env["OPTIMUS_SANDBOX_SUBPROCESS"] = "1"
        clean_env["ALLOW_EXTERNAL_API_CALLS"] = "False"
        clean_env["ALLOW_LIVE_TRADING"] = "False"
        clean_env["ALLOW_PHYSICAL_ACTUATION"] = "False"
        return clean_env

    def execute_in_sandbox(self, script_path: Path, args: Optional[List[str]] = None) -> Dict[str, Any]:
        """Executes an internal script with path validation, env scrubbing, memory monitoring, and timeout boundaries."""
        import psutil

        # 1. Assert path safety
        safe_path = CanonicalPathGuard.assert_path_contained(script_path)

        # 2. Check script contents
        content = safe_path.read_text(encoding="utf-8", errors="replace")
        SandboxSecurityGuard.assert_code_safe(content)

        sandbox_cwd = settings.sandbox_tmp_path.resolve()
        sandbox_cwd.mkdir(parents=True, exist_ok=True)
        sanitized_env = self.build_sanitized_env()

        cmd = [os.sys.executable, str(safe_path)] + (args or [])
        start_time = time.time()
        peak_memory_mb = 0.0

        try:
            proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                cwd=str(sandbox_cwd),
                env=sanitized_env
            )
            ps_proc = None
            try:
                ps_proc = psutil.Process(proc.pid)
            except Exception:
                ps_proc = None

            while proc.poll() is None:
                elapsed = time.time() - start_time
                if ps_proc is not None:
                    try:
                        total_rss_bytes = ps_proc.memory_info().rss
                        children = ps_proc.children(recursive=True)
                        for ch in children:
                            try:
                                total_rss_bytes += ch.memory_info().rss
                            except (psutil.NoSuchProcess, psutil.AccessDenied):
                                pass
                        rss_mb = total_rss_bytes / (1024 * 1024)
                        if rss_mb > peak_memory_mb:
                            peak_memory_mb = rss_mb
                        if rss_mb > self.max_memory_mb:
                            for ch in children:
                                try:
                                    ch.kill()
                                except (psutil.NoSuchProcess, psutil.AccessDenied):
                                    pass
                            proc.kill()
                            stdout_out, stderr_out = proc.communicate(timeout=2)
                            return {
                                "success": False,
                                "stdout": stdout_out or "",
                                "stderr": stderr_out or "",
                                "exit_code": -2,
                                "wall_time_seconds": time.time() - start_time,
                                "peak_memory_mb": round(peak_memory_mb, 2),
                                "max_memory_mb": self.max_memory_mb,
                                "env_scrubbed": True,
                                "isolation_level": "SUBPROCESS_HARNESS_NOT_OS_CONTAINER",
                                "error": f"Execution exceeded memory limit ({rss_mb:.1f} MB > {self.max_memory_mb} MB)."
                            }
                    except (psutil.NoSuchProcess, psutil.AccessDenied):
                        pass

                if elapsed > self.timeout_seconds:
                    if ps_proc is not None:
                        for ch in ps_proc.children(recursive=True):
                            try:
                                ch.kill()
                            except (psutil.NoSuchProcess, psutil.AccessDenied):
                                pass
                    proc.kill()
                    stdout_out, stderr_out = proc.communicate(timeout=2)
                    return {
                        "success": False,
                        "stdout": stdout_out or "",
                        "stderr": stderr_out or "",
                        "exit_code": -1,
                        "wall_time_seconds": elapsed,
                        "peak_memory_mb": round(peak_memory_mb, 2),
                        "max_memory_mb": self.max_memory_mb,
                        "env_scrubbed": True,
                        "isolation_level": "SUBPROCESS_HARNESS_NOT_OS_CONTAINER",
                        "error": f"Execution timed out after {self.timeout_seconds} seconds."
                    }
                time.sleep(0.02)

            stdout_out, stderr_out = proc.communicate(timeout=2)
            wall_time = time.time() - start_time
            return {
                "success": proc.returncode == 0,
                "stdout": stdout_out,
                "stderr": stderr_out,
                "exit_code": proc.returncode,
                "wall_time_seconds": wall_time,
                "peak_memory_mb": round(peak_memory_mb, 2),
                "max_memory_mb": self.max_memory_mb,
                "env_scrubbed": True,
                "isolation_level": "SUBPROCESS_HARNESS_NOT_OS_CONTAINER"
            }
        except Exception as e:
            return {
                "success": False,
                "stdout": "",
                "stderr": str(e),
                "exit_code": -1,
                "wall_time_seconds": time.time() - start_time,
                "peak_memory_mb": round(peak_memory_mb, 2),
                "max_memory_mb": self.max_memory_mb,
                "env_scrubbed": True,
                "isolation_level": "SUBPROCESS_HARNESS_NOT_OS_CONTAINER",
                "error": str(e)
            }