"""ContainerisedToolRunner class for ReXia.AI"""

import logging
from typing import Any, Dict
import docker
import tempfile
import os

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger(__name__)

class ContainerisedToolRunner:
    """
    A class for executing LLM-generated Python code as tools in isolated Docker containers.
    
    This class provides a secure environment for running potentially untrusted code
    generated by language models.
    """

    def __init__(self, image: str = "python:3.12-slim", timeout: int = 30):
        """
        Initialize a ContainerisedToolRunner instance.

        Args:
            image (str, optional): The Docker image to use for the container. Defaults to "python:3.12-slim".
            timeout (int, optional): The maximum execution time in seconds. Defaults to 30.

        Raises:
            RuntimeError: If the Docker client fails to initialize.
        """
        try:
            self.client = docker.from_env()
            logger.info("Docker client initialized successfully.")
        except Exception as e:
            error_msg = f"Failed to initialize Docker client: {str(e)}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)
        
        self.image = image
        self.timeout = timeout
        logger.info(f"ContainerisedToolRunner initialized with image: {image}, timeout: {timeout}s")

    def execute_code(self, code: str) -> Dict[str, Any]:
        """
        Execute the provided code in a Docker container.

        Args:
            code (str): The Python code to execute.

        Returns:
            Dict[str, Any]: A dictionary containing the execution result.
                            Keys: 'success' (bool), 'output' (str) if successful, 'error' (str) if failed.

        Raises:
            Exception: If an error occurs during code execution.
        """
        try:
            with tempfile.TemporaryDirectory() as tmpdir:
                logger.info(f"Created temporary directory: {tmpdir}")
                self._write_files(tmpdir, code)
                status_code, stdout, stderr = self._run_container(tmpdir)

                if status_code == 0:
                    logger.info("Code execution successful")
                    return {"success": True, "output": stdout}
                else:
                    logger.warning(f"Code execution failed with status code: {status_code}")
                    return {"success": False, "error": stderr}
        except Exception as e:
            error_message = f"An error occurred during code execution: {str(e)}"
            logger.error(error_message)
            logger.exception("Exception details:")
            return {"success": False, "error": error_message}

    def _write_files(self, tmpdir: str, code: str) -> None:
        """
        Write the code to files in the temporary directory.

        Args:
            tmpdir (str): The path to the temporary directory.
            code (str): The Python code to write.

        Raises:
            IOError: If there's an error writing to the files.
        """
        try:
            code_path = os.path.join(tmpdir, "tool.py")
            main_path = os.path.join(tmpdir, "main.py")

            with open(code_path, "w") as f:
                f.write(code)
            logger.info(f"Written code to: {code_path}")

            main_content = """
import tool
import json

if __name__ == '__main__':
    result = tool.main()
    print(json.dumps({"result": result}))
"""
            with open(main_path, "w") as f:
                f.write(main_content)
            logger.info(f"Written main execution logic to: {main_path}")
        except IOError as e:
            logger.error(f"Error writing files: {str(e)}")
            raise IOError(f"Error writing files: {str(e)}")

    def _run_container(self, tmpdir: str) -> tuple:
        """
        Run the Docker container with the provided code.

        Args:
            tmpdir (str): The path to the temporary directory containing the code files.

        Returns:
            tuple: A tuple containing (status_code, stdout, stderr).

        Raises:
            docker.errors.ContainerError: If the container exits with a non-zero status code.
            docker.errors.ImageNotFound: If the specified Docker image is not found.
            docker.errors.APIError: If there's an error in the Docker API.
        """
        try:
            logger.info(f"Starting container with image: {self.image}")
            container = self.client.containers.run(
                self.image,
                command=["python", "/app/main.py"],
                volumes={tmpdir: {"bind": "/app", "mode": "ro"}},
                working_dir="/app",
                detach=True,
                mem_limit="128m",
                cpu_quota=50000,
                network_mode="none",
            )

            logger.info("Container started. Waiting for completion...")
            result = container.wait(timeout=self.timeout)
            stdout = container.logs(stdout=True, stderr=False).decode("utf-8")
            stderr = container.logs(stdout=False, stderr=True).decode("utf-8")
            logger.info(f"Container execution completed with status code: {result['StatusCode']}")
            return result["StatusCode"], stdout, stderr
        except Exception as e:
            logger.error(f"Unexpected error in _run_container: {str(e)}")
            return 1, "", f"Unexpected error in _run_container: {str(e)}"
        finally:
            try:
                container.remove(force=True)
                logger.info("Container removed.")
            except Exception as e:
                logger.error(f"Failed to remove container: {str(e)}")