#!/usr/bin/env python3
"""
Simplified test CLI for MCP-Forge execution backends.

This is a minimal testing tool that bypasses full security stack.
MCP injection is not supported - this is purely for testing execution.

Usage:
    # Simple (stateless) backend:
    ./simple_test_cli.py
    
    # Jupyter (stateful) backend:
    ./simple_test_cli.py --jupyter
    
    # One-shot execution:
    ./simple_test_cli.py --execute "print(2 + 2)"
"""

import sys
import os
import tempfile
import argparse
import time
import uuid
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from mcp_forge.execution.simple.executor import ExecutionResult
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
from mcp_forge.podman.client import PodmanClient
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
from mcp_forge.security.audit import AuditLogger
from mcp_forge.security.allowlist import OperationValidator
from mcp_forge.security.resource_limits import ResourceLimits
import time
from typing import Optional, Dict


# ANSI colors
GREEN = "\033[92m"
RED = "\033[91m"
GRAY = "\033[90m"
YELLOW = "\033[93m"
RESET = "\033[0m"


class PassthroughValidator(OperationValidator):
    """Dummy validator that allows all operations (for testing only)."""
    
    def __init__(self):
        """Initialize without allowlist manager."""
        # Don't call parent __init__ to avoid needing AllowlistManager
        # But we need to initialize the session_containers set
        self.session_containers = set()
    
    def validate_operation(self, operation: str, target: str) -> tuple:
        """Allow all operations without validation."""
        return (True, None)
    
    def validate_container_create(self, image: str, params: dict, session_id: str = None) -> None:
        """Allow all container creations."""
        pass
    
    def register_session_container(self, container_id: str) -> None:
        """Track session containers."""
        self.session_containers.add(container_id)
    
    def unregister_session_container(self, container_id: str) -> None:
        """Untrack session containers."""
        self.session_containers.discard(container_id)


class SimpleExecutorWrapper:
    """Wrapper for CodeExecutor that bypasses resource limits."""
    
    def __init__(self, container_manager: SecureContainerManager, image: str):
        self.container_manager = container_manager
        self.image = image
    
    def execute(self, code: str, timeout: int = 30) -> ExecutionResult:
        """Execute code without resource limits."""
        start_time = time.time()
        container_id = None
        
        try:
            # Prepare code with result extraction
            wrapped_code = f"""
import sys
import json

try:
    # Execute code
    exec_globals = {{'__builtins__': __builtins__}}
    exec({repr(code)}, exec_globals)
    
    # Get result (last expression value if any)
    result = exec_globals.get('_', None)
except Exception as e:
    import traceback
    print(traceback.format_exc(), file=sys.stderr)
    sys.exit(1)
"""
            
            # Create container configuration WITHOUT resource limits
            config = ContainerConfig(
                image=self.image,
                command=["python", "-c", wrapped_code],
                resource_limits=None  # Skip resource limits for rootless testing
            )
            
            # Create and start container
            session_id = str(uuid.uuid4())
            container_id = self.container_manager.create_container(config, session_id=session_id)
            self.container_manager.start_container(container_id)
            
            # Wait for completion
            exit_code = self.container_manager.wait_for_container(container_id, timeout=timeout)
            
            # Get logs (returns tuple of (stdout, stderr))
            stdout, stderr = self.container_manager.get_container_logs(container_id)
            
            execution_time = time.time() - start_time
            
            return ExecutionResult(
                success=(exit_code == 0),
                stdout=stdout,
                stderr=stderr,
                exit_code=exit_code,
                execution_time=execution_time,
                result=None,  # Don't try to extract result for simplicity
                error=stderr if exit_code != 0 else None
            )
            
        finally:
            # Cleanup
            if container_id:
                try:
                    self.container_manager.remove_container(container_id)
                except Exception:
                    pass


class JupyterKernelWrapper:
    """Wrapper for JupyterKernelManager that bypasses resource limits."""
    
    def __init__(self, container_manager: SecureContainerManager, image: str, connection_dir: Path):
        # Create minimal resource limits
        minimal_limits = ResourceLimits(
            memory="512m",
            cpu_quota=100000,
            storage="1g",
            timeout=30
        )
        self.manager = JupyterKernelManager(
            container_manager=container_manager,
            image=image,
            resource_limits=minimal_limits
        )
        # Override resource_limits to None after construction to avoid cgroup issues
        self.manager.resource_limits = None
        self.connection_dir = connection_dir
    
    def start_kernel(self, session_id: str) -> str:
        """Start kernel."""
        return self.manager.start_kernel(session_id)
    
    def execute_code(self, kernel_id: str, code: str) -> ExecutionResult:
        """Execute code in kernel."""
        return self.manager.execute_code(kernel_id, code)
    
    def shutdown_kernel(self, kernel_id: str) -> None:
        """Shutdown kernel."""
        self.manager.shutdown_kernel(kernel_id)


class SimpleCLI:
    """Minimal CLI for testing executors without full security stack."""

    def __init__(self, use_jupyter: bool = False):
        """Initialize CLI with minimal setup."""
        self.use_jupyter = use_jupyter

        # Create temp directory for logs and connection files
        self.temp_dir = Path(tempfile.mkdtemp(prefix="mcp_forge_cli_"))
        print(f"{GRAY}Using temp dir: {self.temp_dir}{RESET}")

        # Setup audit logger
        log_file = self.temp_dir / "audit.log"
        self.audit_logger = AuditLogger(log_file)

        # Setup Podman client with passthrough validator
        # Use user socket (rootless Podman)
        runtime_dir = os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")
        socket_path = Path(runtime_dir) / "podman" / "podman.sock"
        
        self.validator = PassthroughValidator()
        self.podman_client = PodmanClient(
            socket_path=socket_path,
            validator=self.validator,
            audit_logger=self.audit_logger
        )

        # Setup container manager
        self.container_manager = SecureContainerManager(
            podman_client=self.podman_client,
            validator=self.validator,
            audit_logger=self.audit_logger
        )

        # Setup executor or kernel manager
        if use_jupyter:
            print(f"{GRAY}Using Jupyter (stateful) backend{RESET}")
            self.kernel_manager = JupyterKernelWrapper(
                container_manager=self.container_manager,
                image="mcp-forge/jupyter:latest",
                connection_dir=self.temp_dir
            )
            self.session_id = str(uuid.uuid4())
            self.kernel_id = None
        else:
            print(f"{GRAY}Using Simple (stateless) backend{RESET}")
            self.executor = SimpleExecutorWrapper(
                container_manager=self.container_manager,
                image="mcp-forge/python:3.12"
            )

    def execute_code(self, code: str) -> dict:
        """Execute code using the configured backend."""
        if self.use_jupyter:
            # Start kernel if not already started
            if not self.kernel_id:
                print(f"{GRAY}Starting Jupyter kernel...{RESET}")
                self.kernel_id = self.kernel_manager.start_kernel(
                    session_id=self.session_id
                )
                print(f"{GRAY}Kernel started: {self.kernel_id}{RESET}")

            # Execute code
            result = self.kernel_manager.execute_code(self.kernel_id, code)
            return {
                "output": result.stdout + result.stderr,
                "result": result.result,
                "success": result.exit_code == 0,
                "execution_time": result.execution_time,
                "exit_code": result.exit_code
            }
        else:
            # Simple executor
            result = self.executor.execute(code=code)
            return {
                "output": result.stdout + result.stderr,
                "result": result.result,
                "success": result.exit_code == 0,
                "execution_time": result.execution_time,
                "exit_code": result.exit_code
            }

    def repl(self):
        """Run interactive REPL."""
        print(f"\n{GREEN}MCP-Forge Interactive Shell{RESET}")
        if self.use_jupyter:
            print(f"{YELLOW}Note: Variables persist between executions in Jupyter mode{RESET}")
        else:
            print(f"{YELLOW}Note: Each execution is isolated in Simple mode{RESET}")
        print(f"{GRAY}Type '.exit' or '.quit' to exit, '.help' for help{RESET}\n")

        while True:
            try:
                # Get code input (support multi-line with empty prompt on continuation)
                code_lines = []
                while True:
                    if not code_lines:
                        line = input(">>> ")
                    else:
                        line = input("... ")

                    code_lines.append(line)

                    # Check if more lines are needed
                    code = "\n".join(code_lines)
                    if not line.strip() or not self._needs_more_lines(code):
                        break

                code = code.strip()
                if not code:
                    continue

                # Check for special commands
                if code in [".exit", ".quit"]:
                    print("Goodbye!")
                    break
                elif code == ".help":
                    self._show_help()
                    continue
                elif code == ".restart":
                    if self.use_jupyter and self.kernel_id:
                        print(f"{GRAY}Restarting kernel...{RESET}")
                        self.kernel_manager.shutdown_kernel(self.kernel_id)
                        self.kernel_id = None
                        print(f"{GREEN}Kernel will restart on next execution{RESET}")
                    else:
                        print(f"{RED}Restart only available in Jupyter mode{RESET}")
                    continue
                elif code == ".vars":
                    if self.use_jupyter:
                        # Show variables by running dir()
                        result = self.execute_code("print([v for v in dir() if not v.startswith('_')])")
                        if result["success"]:
                            print(f"{GREEN}{result['output']}{RESET}")
                    else:
                        print(f"{RED}Variables only available in Jupyter mode (stateless backend){RESET}")
                    continue
                elif code == ".clear":
                    os.system('clear' if os.name == 'posix' else 'cls')
                    continue

                # Execute code
                start_time = time.time()
                result = self.execute_code(code)
                elapsed = time.time() - start_time

                # Display result
                if result["success"]:
                    # Show output (stdout/stderr)
                    if result["output"].strip():
                        print(result["output"])
                    # Show result value if different from output
                    if result["result"] and result["result"] not in [None, "None"]:
                        print(f"{GREEN}→ {result['result']}{RESET}")
                else:
                    # Show error
                    print(f"{RED}{result['output']}{RESET}")

                print(f"{GRAY}({elapsed:.2f}s){RESET}")

            except KeyboardInterrupt:
                print("\n(Use .exit or .quit to exit)")
            except EOFError:
                print("\nGoodbye!")
                break
            except Exception as e:
                print(f"{RED}CLI Error: {e}{RESET}")
                import traceback
                traceback.print_exc()

    def cleanup(self):
        """Clean up resources."""
        if self.use_jupyter and self.kernel_id:
            try:
                print(f"\n{GRAY}Shutting down kernel...{RESET}")
                self.kernel_manager.shutdown_kernel(self.kernel_id)
            except Exception as e:
                print(f"{GRAY}Warning: Error shutting down kernel: {e}{RESET}")

    def _needs_more_lines(self, code: str) -> bool:
        """Check if code needs more lines (simple heuristic)."""
        # Check for unclosed brackets/parens
        opens = code.count('(') + code.count('[') + code.count('{')
        closes = code.count(')') + code.count(']') + code.count('}')
        if opens > closes:
            return True

        # Check for continuation indicators
        if code.rstrip().endswith((':', '\\')):
            return True

        return False

    def _show_help(self):
        """Show help message."""
        print(f"""
{GREEN}Special commands:{RESET}
  .exit, .quit  - Exit the shell
  .help         - Show this help
  .restart      - Restart Jupyter kernel (Jupyter mode only)
  .vars         - Show variables (Jupyter mode only)
  .clear        - Clear screen

{GREEN}Tips:{RESET}
  - Multi-line input: Press Enter on a blank line to execute
  - Jupyter mode: Variables persist between executions
  - Simple mode: Each execution is in a fresh container
        """.strip())


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Simple CLI for testing MCP-Forge backends"
    )
    parser.add_argument(
        "--jupyter", "-j",
        action="store_true",
        help="Use Jupyter (stateful) backend instead of simple executor"
    )
    parser.add_argument(
        "--execute", "-e",
        type=str,
        help="Execute code and exit (non-interactive)"
    )

    args = parser.parse_args()

    # Create CLI
    cli = SimpleCLI(use_jupyter=args.jupyter)

    try:
        if args.execute:
            # One-shot execution
            result = cli.execute_code(args.execute)
            if result["success"]:
                # Print output
                if result["output"].strip():
                    print(result["output"])
                # Print result if available
                if result["result"] and result["result"] not in [None, "None"]:
                    print(result["result"])
                sys.exit(0)
            else:
                print(result["output"], file=sys.stderr)
                sys.exit(1)
        else:
            # Interactive REPL
            cli.repl()
    finally:
        cli.cleanup()


if __name__ == "__main__":
    main()