#!/usr/bin/env python3
"""
inference_tripwire.py - Enforce "Inference = Decode Loop" Vocabulary

WO-VOCABULARY-ENFORCE-01: Mechanical enforcement that prevents foreign inference
stacks from being invoked during decode benchmarks.

In helix-cdc:
  - "inference" = helixcode.matmul(X), stream_xw_from_cdna()
  - "decode" = Block-wise X@W from CDNA/HelixCode indices
  - "streaming" = Memory-efficient block streaming

NOT:
  - llama_cpp token generation
  - transformers.generate()
  - Chat token streaming

Usage:
    from helix_cdc.regrow import InferenceTripwire, assert_decode_only

    with InferenceTripwire():
        Y = helixcode.matmul(X)  # OK - this is decode
        # llm = llama_cpp.Llama(...)  # TRIPWIRE! RuntimeError

    # Or use the convenience decorator
    with assert_decode_only():
        Y, stats = hx.matmul_with_stats(X)

To disable for chat mode (explicit escape hatch):
    HELIX_ALLOW_FOREIGN_INFERENCE=1 python my_script.py

Work Order: WO-VOCABULARY-ENFORCE-01
"""

from __future__ import annotations

import os
import subprocess
import sys
from contextlib import contextmanager
from typing import Any, Callable, Dict, List, Optional

# Known foreign inference packages
FOREIGN_INFERENCE_PACKAGES = frozenset([
    'llama_cpp',
    'llama-cpp-python',
    'vllm',
    'ctransformers',
    'litellm',
    'text_generation',
])

# Subprocess args that indicate foreign inference
FOREIGN_INFERENCE_SUBPROCESS_KEYWORDS = frozenset([
    'llama', 'gguf', 'vllm', 'transformers', 'text-generation',
    'llama.cpp', 'llama-cpp', 'ctransformers', 'litellm',
])


def _is_foreign_inference_allowed() -> bool:
    """Check if foreign inference is explicitly allowed via env var."""
    return os.environ.get('HELIX_ALLOW_FOREIGN_INFERENCE', '').lower() in ('1', 'true', 'yes')


class InferenceTripwire:
    """
    Context manager that fails if foreign inference stacks are invoked.

    Monkeypatches llama_cpp, transformers, and subprocess to raise RuntimeError
    if called during decode benchmarks.

    This provides HARD PROOF that decode benchmarks use only the
    HelixCode/regrow decode loop, not foreign inference stacks.

    Args:
        enabled: If False, tripwire is disabled (no-op). Default True.
                 Also disabled if HELIX_ALLOW_FOREIGN_INFERENCE=1 env var is set.
        trap_subprocess: If True, also trap subprocess calls that look like
                         foreign inference invocations. Default True.
    """

    TRIPWIRE_MESSAGE = (
        "TRIPWIRE: Foreign inference stack invoked in decode benchmark!\n"
        "  In helix-cdc, 'inference = decode loop = HelixCode/regrow'.\n"
        "  DO NOT use llama_cpp, transformers.generate in decode benchmarks.\n"
        "  See WO-VOCABULARY-ENFORCE-01 for vocabulary enforcement.\n"
        "  To disable: set HELIX_ALLOW_FOREIGN_INFERENCE=1"
    )

    def __init__(self, enabled: bool = True, trap_subprocess: bool = True):
        self._enabled = enabled and not _is_foreign_inference_allowed()
        self._trap_subprocess = trap_subprocess
        self._original_funcs: Dict[str, Any] = {}
        self._modules_patched: Dict[str, bool] = {}

    def _make_tripwire_class(self, name: str):
        """Create a class that raises on instantiation."""
        message = self.TRIPWIRE_MESSAGE

        class TripwireClass:
            def __init__(self_inner, *args, **kwargs):
                raise RuntimeError(f"{message}\n  Attempted to instantiate: {name}")

        return TripwireClass

    def _make_tripwire_func(self, name: str):
        """Create a function that raises when called."""
        message = self.TRIPWIRE_MESSAGE

        def tripwire_func(*args, **kwargs):
            raise RuntimeError(f"{message}\n  Attempted to call: {name}")

        return tripwire_func

    def _make_tripwire_method(self, name: str, original: Callable):
        """Create a method that raises when called (for instance methods)."""
        message = self.TRIPWIRE_MESSAGE

        def tripwire_method(self_inner, *args, **kwargs):
            raise RuntimeError(f"{message}\n  Attempted to call: {name}")

        return tripwire_method

    def _make_subprocess_trap(self, original_func: Callable, func_name: str):
        """Create a subprocess wrapper that traps foreign inference calls."""
        message = self.TRIPWIRE_MESSAGE
        keywords = FOREIGN_INFERENCE_SUBPROCESS_KEYWORDS

        def trapped_subprocess(*args, **kwargs):
            # Check args for foreign inference indicators
            cmd = args[0] if args else kwargs.get('args', [])
            if isinstance(cmd, str):
                cmd_lower = cmd.lower()
                for keyword in keywords:
                    if keyword in cmd_lower:
                        raise RuntimeError(
                            f"{message}\n  Attempted subprocess call with foreign inference: {func_name}\n"
                            f"  Command: {cmd[:100]}..."
                        )
            elif isinstance(cmd, (list, tuple)):
                cmd_str = ' '.join(str(c) for c in cmd).lower()
                for keyword in keywords:
                    if keyword in cmd_str:
                        raise RuntimeError(
                            f"{message}\n  Attempted subprocess call with foreign inference: {func_name}\n"
                            f"  Command: {cmd_str[:100]}..."
                        )
            return original_func(*args, **kwargs)

        return trapped_subprocess

    def __enter__(self):
        if not self._enabled:
            return self

        # =========================================================
        # llama_cpp patching
        # =========================================================
        try:
            import llama_cpp
            self._original_funcs['llama_cpp.Llama'] = llama_cpp.Llama
            llama_cpp.Llama = self._make_tripwire_class('llama_cpp.Llama')
            self._modules_patched['llama_cpp.Llama'] = True
        except ImportError:
            self._modules_patched['llama_cpp.Llama'] = False

        # =========================================================
        # transformers patching
        # =========================================================

        # transformers.pipeline
        try:
            import transformers
            self._original_funcs['transformers.pipeline'] = transformers.pipeline
            transformers.pipeline = self._make_tripwire_func('transformers.pipeline')
            self._modules_patched['transformers.pipeline'] = True
        except (ImportError, AttributeError):
            self._modules_patched['transformers.pipeline'] = False

        # transformers.AutoModelForCausalLM.generate
        try:
            from transformers import AutoModelForCausalLM
            if hasattr(AutoModelForCausalLM, 'generate'):
                self._original_funcs['AutoModelForCausalLM.generate'] = AutoModelForCausalLM.generate
                AutoModelForCausalLM.generate = self._make_tripwire_func('AutoModelForCausalLM.generate')
                self._modules_patched['AutoModelForCausalLM.generate'] = True
            else:
                self._modules_patched['AutoModelForCausalLM.generate'] = False
        except ImportError:
            self._modules_patched['AutoModelForCausalLM.generate'] = False

        # transformers.GenerationMixin.generate (catches most model.generate() calls)
        try:
            from transformers.generation import GenerationMixin
            if hasattr(GenerationMixin, 'generate'):
                self._original_funcs['GenerationMixin.generate'] = GenerationMixin.generate
                GenerationMixin.generate = self._make_tripwire_func('GenerationMixin.generate')
                self._modules_patched['GenerationMixin.generate'] = True
            else:
                self._modules_patched['GenerationMixin.generate'] = False
        except ImportError:
            self._modules_patched['GenerationMixin.generate'] = False

        # transformers.TextGenerationPipeline.__call__
        try:
            from transformers import TextGenerationPipeline
            if hasattr(TextGenerationPipeline, '__call__'):
                self._original_funcs['TextGenerationPipeline.__call__'] = TextGenerationPipeline.__call__
                TextGenerationPipeline.__call__ = self._make_tripwire_method(
                    'TextGenerationPipeline.__call__',
                    TextGenerationPipeline.__call__
                )
                self._modules_patched['TextGenerationPipeline.__call__'] = True
            else:
                self._modules_patched['TextGenerationPipeline.__call__'] = False
        except ImportError:
            self._modules_patched['TextGenerationPipeline.__call__'] = False

        # transformers.Pipeline.__call__ (base class)
        try:
            from transformers.pipelines.base import Pipeline
            if hasattr(Pipeline, '__call__'):
                self._original_funcs['Pipeline.__call__'] = Pipeline.__call__
                Pipeline.__call__ = self._make_tripwire_method('Pipeline.__call__', Pipeline.__call__)
                self._modules_patched['Pipeline.__call__'] = True
            else:
                self._modules_patched['Pipeline.__call__'] = False
        except ImportError:
            self._modules_patched['Pipeline.__call__'] = False

        # =========================================================
        # vllm patching
        # =========================================================
        try:
            import vllm
            if hasattr(vllm, 'LLM'):
                self._original_funcs['vllm.LLM'] = vllm.LLM
                vllm.LLM = self._make_tripwire_class('vllm.LLM')
                self._modules_patched['vllm.LLM'] = True
            else:
                self._modules_patched['vllm.LLM'] = False
        except ImportError:
            self._modules_patched['vllm.LLM'] = False

        # =========================================================
        # ctransformers patching
        # =========================================================
        try:
            import ctransformers
            if hasattr(ctransformers, 'AutoModelForCausalLM'):
                self._original_funcs['ctransformers.AutoModelForCausalLM'] = ctransformers.AutoModelForCausalLM
                ctransformers.AutoModelForCausalLM = self._make_tripwire_class('ctransformers.AutoModelForCausalLM')
                self._modules_patched['ctransformers.AutoModelForCausalLM'] = True
            else:
                self._modules_patched['ctransformers.AutoModelForCausalLM'] = False
        except ImportError:
            self._modules_patched['ctransformers.AutoModelForCausalLM'] = False

        # =========================================================
        # subprocess patching (optional, catches shell-outs to inference binaries)
        # =========================================================
        if self._trap_subprocess:
            # subprocess.run
            self._original_funcs['subprocess.run'] = subprocess.run
            subprocess.run = self._make_subprocess_trap(subprocess.run, 'subprocess.run')
            self._modules_patched['subprocess.run'] = True

            # subprocess.Popen
            self._original_funcs['subprocess.Popen'] = subprocess.Popen
            # For Popen, we need a class wrapper
            original_popen = subprocess.Popen
            message = self.TRIPWIRE_MESSAGE
            keywords = FOREIGN_INFERENCE_SUBPROCESS_KEYWORDS

            class TrapPopen(original_popen):
                def __init__(self_inner, args, *pargs, **kwargs):
                    cmd = args
                    if isinstance(cmd, str):
                        cmd_lower = cmd.lower()
                        for keyword in keywords:
                            if keyword in cmd_lower:
                                raise RuntimeError(
                                    f"{message}\n  Attempted subprocess.Popen with foreign inference\n"
                                    f"  Command: {cmd[:100]}..."
                                )
                    elif isinstance(cmd, (list, tuple)):
                        cmd_str = ' '.join(str(c) for c in cmd).lower()
                        for keyword in keywords:
                            if keyword in cmd_str:
                                raise RuntimeError(
                                    f"{message}\n  Attempted subprocess.Popen with foreign inference\n"
                                    f"  Command: {cmd_str[:100]}..."
                                )
                    super().__init__(args, *pargs, **kwargs)

            subprocess.Popen = TrapPopen
            self._modules_patched['subprocess.Popen'] = True

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self._enabled:
            return False

        # Restore all patched functions/classes
        # llama_cpp
        if 'llama_cpp.Llama' in self._original_funcs:
            try:
                import llama_cpp
                llama_cpp.Llama = self._original_funcs['llama_cpp.Llama']
            except ImportError:
                pass

        # transformers
        if 'transformers.pipeline' in self._original_funcs:
            try:
                import transformers
                transformers.pipeline = self._original_funcs['transformers.pipeline']
            except ImportError:
                pass

        if 'AutoModelForCausalLM.generate' in self._original_funcs:
            try:
                from transformers import AutoModelForCausalLM
                AutoModelForCausalLM.generate = self._original_funcs['AutoModelForCausalLM.generate']
            except ImportError:
                pass

        if 'GenerationMixin.generate' in self._original_funcs:
            try:
                from transformers.generation import GenerationMixin
                GenerationMixin.generate = self._original_funcs['GenerationMixin.generate']
            except ImportError:
                pass

        if 'TextGenerationPipeline.__call__' in self._original_funcs:
            try:
                from transformers import TextGenerationPipeline
                TextGenerationPipeline.__call__ = self._original_funcs['TextGenerationPipeline.__call__']
            except ImportError:
                pass

        if 'Pipeline.__call__' in self._original_funcs:
            try:
                from transformers.pipelines.base import Pipeline
                Pipeline.__call__ = self._original_funcs['Pipeline.__call__']
            except ImportError:
                pass

        # vllm
        if 'vllm.LLM' in self._original_funcs:
            try:
                import vllm
                vllm.LLM = self._original_funcs['vllm.LLM']
            except ImportError:
                pass

        # ctransformers
        if 'ctransformers.AutoModelForCausalLM' in self._original_funcs:
            try:
                import ctransformers
                ctransformers.AutoModelForCausalLM = self._original_funcs['ctransformers.AutoModelForCausalLM']
            except ImportError:
                pass

        # subprocess
        if 'subprocess.run' in self._original_funcs:
            subprocess.run = self._original_funcs['subprocess.run']

        if 'subprocess.Popen' in self._original_funcs:
            subprocess.Popen = self._original_funcs['subprocess.Popen']

        return False  # Don't suppress exceptions

    def get_patched_modules(self) -> Dict[str, bool]:
        """Return which modules were successfully patched."""
        return self._modules_patched.copy()

    @property
    def enabled(self) -> bool:
        """Return whether the tripwire is enabled."""
        return self._enabled


@contextmanager
def assert_decode_only(enabled: bool = True, trap_subprocess: bool = True):
    """
    Convenience context manager for decode-only assertions.

    Equivalent to InferenceTripwire() but with a more descriptive name.

    Args:
        enabled: If False, tripwire is disabled (no-op). Default True.
        trap_subprocess: If True, also trap subprocess calls. Default True.

    Usage:
        with assert_decode_only():
            Y, stats = hx.matmul_with_stats(X)
            # Any call to llama_cpp or transformers.generate raises RuntimeError
    """
    with InferenceTripwire(enabled=enabled, trap_subprocess=trap_subprocess) as tripwire:
        yield tripwire


if __name__ == "__main__":
    print("=== InferenceTripwire Test ===")
    print()
    print("This module provides vocabulary enforcement for helix-cdc.")
    print()
    print("In helix-cdc:")
    print("  - 'inference' = helixcode.matmul(X), stream_xw_from_cdna()")
    print("  - 'decode' = Block-wise X@W from CDNA/HelixCode indices")
    print("  - 'streaming' = Memory-efficient block streaming")
    print()
    print("NOT:")
    print("  - llama_cpp token generation")
    print("  - transformers.generate()")
    print("  - Chat token streaming")
    print()
    print("Usage:")
    print("  from helix_cdc.regrow import InferenceTripwire, assert_decode_only")
    print()
    print("  with InferenceTripwire():")
    print("      Y = helixcode.matmul(X)  # OK")
    print("      # llama_cpp.Llama(...)   # TRIPWIRE!")
    print()
    print("To disable for chat mode:")
    print("  HELIX_ALLOW_FOREIGN_INFERENCE=1 python my_script.py")
    print()

    # Test env var
    print(f"HELIX_ALLOW_FOREIGN_INFERENCE currently: {_is_foreign_inference_allowed()}")
