"""Jupyter backend adapter for MCP-Forge.

Wraps pod_executor.JupyterBackend with MCP-Forge configuration.
"""

from typing import Optional, List, Dict, Any
import logging

from pod_executor import JupyterBackend as PodJupyterBackend, ExecutionResult
from pod_executor.containers.manager import SecureContainerManager
from pod_executor.jupyter.sessions import SessionState, SessionError

from ..config.schema import ForgeConfig
from ..security.audit import AuditLogger

logger = logging.getLogger(__name__)


# Re-export for compatibility
__all__ = ["JupyterBackend", "SessionState", "SessionError"]


class JupyterBackend:
    """Adapter for stateful Python code execution using Jupyter kernels.
    
    This wraps pod_executor.JupyterBackend and adapts it to MCP-Forge's
    configuration system.
    """
    
    def __init__(
        self,
        container_manager: SecureContainerManager,
        audit_logger: AuditLogger,
        config: ForgeConfig
    ):
        """
        Initialize Jupyter backend adapter.
        
        Args:
            container_manager: Container lifecycle manager
            audit_logger: Audit logging instance
            config: MCP-Forge configuration
        """
        self.container_manager = container_manager
        self.audit_logger = audit_logger
        self.config = config
        
        # Create pod_executor backend with config parameters
        self.backend = PodJupyterBackend(
            container_manager=container_manager,
            image=config.images.jupyter,
            default_timeout=config.execution.default_timeout,
            default_memory=config.execution.default_memory,
            default_cpu_quota=config.execution.default_cpu_quota,
            max_timeout=config.execution.max_timeout,
            max_memory=config.execution.max_memory,
            max_cpu_quota=config.execution.max_cpu_quota,
            max_sessions=config.sessions.max_concurrent,
            idle_timeout=config.sessions.idle_timeout,
            audit_logger=audit_logger
        )
        
        logger.debug(f"JupyterBackend initialized with image={config.images.jupyter}")
    
    def execute(
        self,
        code: str,
        session_id: str,
        timeout: Optional[int] = None,
        memory: Optional[str] = None,
        cpu_quota: Optional[int] = None
    ) -> ExecutionResult:
        """
        Execute code in a stateful Jupyter session.
        
        Args:
            code: Python code to execute
            session_id: Session identifier
            timeout: Optional timeout override (seconds)
            memory: Optional memory limit override
            cpu_quota: Optional CPU quota override
            
        Returns:
            ExecutionResult with stdout, stderr, result, etc.
        """
        return self.backend.execute(
            code=code,
            session_id=session_id,
            timeout=timeout,
            memory=memory,
            cpu_quota=cpu_quota
        )
    
    def list_sessions(self) -> List[Dict[str, Any]]:
        """List all active sessions."""
        return self.backend.list_sessions()
    
    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Get session information."""
        return self.backend.get_session(session_id)
    
    def destroy_session(self, session_id: str) -> bool:
        """Destroy a session."""
        return self.backend.destroy_session(session_id)
    
    def cleanup_idle_sessions(self) -> int:
        """Clean up idle sessions."""
        return self.backend.cleanup_idle_sessions()