import hashlib
import os
import platform
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import traceback
import venv

from clarifai_grpc.grpc.api import resources_pb2, service_pb2

from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.utils.constants import MIN_REQUIRED_PYTHON_VERSION
from clarifai.utils.logging import logger


class ModelRunLocally:
    def __init__(self, model_path, model_builder: ModelBuilder = None):
        """
        Initialize a helper to run a Clarifai model locally in an isolated environment.

        Parameters
        ----------
        model_path : str
            Filesystem path to the root directory of the model. This directory is expected
            to contain the model code and a ``requirements.txt`` file describing its
            Python dependencies.
        model_builder : ModelBuilder, optional
            An existing :class:`ModelBuilder` instance to use for interacting with the
            model. Pass this when you already have a configured builder you want to
            reuse. If ``None`` (the default), a new ``ModelBuilder`` is created for
            ``model_path`` with ``download_validation_only=True``.
        """
        self.model_path = os.path.abspath(model_path)
        self.requirements_file = os.path.join(self.model_path, "requirements.txt")

        # ModelBuilder contains multiple useful methods to interact with the model
        self.builder = (
            model_builder
            if model_builder
            else ModelBuilder(self.model_path, download_validation_only=True)
        )
        self.config = self.builder.config

    def _get_method_signatures(self):
        """Get the method signatures for the model."""
        return self.builder.get_method_signatures()

    def _requirements_hash(self):
        """Generate a hash of the requirements file."""
        with open(self.requirements_file, "r") as f:
            return hashlib.md5(f.read().encode('utf-8')).hexdigest()

    def _get_env_executable(self):
        """Get the python executable from the virtual environment."""
        # Depending on the platform, venv scripts are placed in either "Scripts" (Windows) or "bin" (Linux/Mac)
        if platform.system().lower().startswith("win"):
            scripts_folder = "Scripts"
            python_exe = "python.exe"
            pip_exe = "pip.exe"
        else:
            scripts_folder = "bin"
            python_exe = "python"
            pip_exe = "pip"

        self.python_executable = os.path.join(self.venv_dir, scripts_folder, python_exe)
        self.pip_executable = os.path.join(self.venv_dir, scripts_folder, pip_exe)

        return self.python_executable, self.pip_executable

    def create_temp_venv(self):
        """Create a temporary virtual environment."""
        requirements_hash = self._requirements_hash()

        temp_dir = os.path.join(tempfile.gettempdir(), str(requirements_hash))
        venv_dir = os.path.join(temp_dir, "venv")

        if os.path.exists(temp_dir):
            logger.info(f"Using previous virtual environment at {temp_dir}")
            use_existing_venv = True
        else:
            logger.info("Creating temporary virtual environment...")
            use_existing_venv = False
            venv.create(venv_dir, with_pip=True)
            logger.info(f"Created temporary virtual environment at {venv_dir}")

        self.venv_dir = venv_dir
        self.temp_dir = temp_dir
        self.python_executable, self.pip_executable = self._get_env_executable()

        return use_existing_venv

    def install_requirements(self):
        """Install the dependencies from requirements.txt and Clarifai."""
        _, pip_executable = self._get_env_executable()
        try:
            logger.info(
                f"Installing requirements from {self.requirements_file}... in the virtual environment {self.venv_dir}"
            )
            subprocess.check_call([pip_executable, "install", "-r", self.requirements_file])
            logger.info("Installing Clarifai package...")
            subprocess.check_call([pip_executable, "install", "clarifai"])
            logger.info("Requirements installed successfully!")
        except subprocess.CalledProcessError as e:
            logger.error(f"Error installing requirements: {e}")
            self.clean_up()
            sys.exit(1)

    def _build_request(self):
        """Create a mock inference request for testing the model."""

        model_version_proto = self.builder.get_model_version_proto()
        model_version_proto.id = "model_version"

        return service_pb2.PostModelOutputsRequest(
            model=resources_pb2.Model(model_version=model_version_proto),
            inputs=[
                resources_pb2.Input(
                    data=resources_pb2.Data(
                        text=resources_pb2.Text(raw="How many people live in new york?"),
                        image=resources_pb2.Image(
                            url="https://samples.clarifai.com/metro-north.jpg"
                        ),
                        audio=resources_pb2.Audio(
                            url="https://samples.clarifai.com/GoodMorning.wav"
                        ),
                        video=resources_pb2.Video(url="https://samples.clarifai.com/beer.mp4"),
                    )
                )
            ],
        )

    def _run_test(self):
        """Test the model locally by making a prediction."""
        # Create the model
        model = self.builder.create_model_instance()
        # call its test method, if it has one
        if hasattr(model, "test"):
            try:
                model.test()
                logger.info("Model tested successfully!")
            except Exception as e:
                logger.error(f"Error occurred while testing the model: {e}")
                traceback.print_exc()

    def test_model(self):
        """Test the model by running it locally in the virtual environment."""

        import_path = repr(os.path.dirname(os.path.abspath(__file__)))
        model_path = repr(self.model_path)

        command_string = (
            f"import sys; "
            f"sys.path.append({import_path}); "
            f"from model_run_locally import ModelRunLocally; "
            f"ModelRunLocally({model_path})._run_test()"
        )

        command = [self.python_executable, "-c", command_string]
        process = None
        try:
            logger.info("Testing the model locally...")
            # Set PYTHONDONTWRITEBYTECODE=1 to prevent __pycache__ folder generation
            env = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1"}
            process = subprocess.Popen(command, env=env)
            # Wait for the process to complete
            process.wait()
            if process.returncode == 0:
                logger.info("Model tested successfully!")
            if process.returncode != 0:
                raise subprocess.CalledProcessError(process.returncode, command)
        except subprocess.CalledProcessError as e:
            logger.error(f"Error testing the model: {e}")
            sys.exit(1)
        except Exception as e:
            logger.error(f"Unexpected error: {e}")
            sys.exit(1)
        finally:
            # After the function runs, check if the process is still running
            if process and process.poll() is None:
                logger.info("Process is still running. Terminating process.")
                process.terminate()
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    logger.info("Process did not terminate gracefully. Killing process.")
                    # Kill the process if it doesn't terminate after 5 seconds
                    process.kill()

    # run the model server
    def run_model_server(self, port=8080, grpc=True, **kwargs):
        """Run the Clarifai Runner's model server.

        Args:
            port: Port to run the server on.
            grpc: If True, start a standalone gRPC server. If False, start a
                runner that connects to the Clarifai API (requires kwargs like
                user_id, compute_cluster_id, nodepool_id, runner_id, base_url, pat).
            **kwargs: Additional arguments passed to clarifai.runners.server
                (e.g. user_id, compute_cluster_id, nodepool_id, runner_id,
                base_url, pat, num_threads, pool_size).
        """

        command = [
            self.python_executable,
            "-m",
            "clarifai.runners.server",
            "--model_path",
            self.model_path,
            "--port",
            str(port),
        ]
        if grpc:
            command.append("--grpc")
        # Pass additional runner args
        for key, value in kwargs.items():
            if value is None:
                continue
            command.extend([f"--{key}", str(value)])

        try:
            logger.info(
                f"Starting model server at localhost:{port} with the model at {self.model_path}..."
            )
            # Set PYTHONDONTWRITEBYTECODE=1 to prevent __pycache__ folder generation
            env = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1"}
            subprocess.check_call(command, env=env)
            logger.info("Model server started successfully and running at localhost:{port}")
        except subprocess.CalledProcessError as e:
            logger.error(f"Error running model server: {e}")
            self.clean_up()
            sys.exit(1)

    def _docker_hash(self):
        """Generate a hash of the combined requirements file and Dockefile"""
        with open(self.requirements_file, "r") as f:
            requirements_hash = hashlib.md5(f.read().encode('utf-8')).hexdigest()
        with open(os.path.join(self.model_path, "Dockerfile"), "r") as f:
            dockerfile_hash = hashlib.md5(f.read().encode('utf-8')).hexdigest()

        return hashlib.md5(f"{requirements_hash}{dockerfile_hash}".encode('utf-8')).hexdigest()

    def is_docker_installed(self):
        """Checks if Docker is installed on the system."""
        try:
            logger.info("Checking if Docker is installed...")
            subprocess.run(["docker", "--version"], check=True)
            return True
        except subprocess.CalledProcessError:
            logger.error(
                "Docker is not installed! Please install Docker to run the model in a container."
            )
            return False

    def build_docker_image(
        self,
        image_name="model_image",
    ):
        """Build the docker image using the Dockerfile in the model directory."""
        try:
            logger.info(f"Building docker image from Dockerfile in {self.model_path}...")

            # since we don't want to copy the model directory into the container, we need to modify the Dockerfile and comment out the COPY instruction
            dockerfile_path = os.path.join(self.model_path, "Dockerfile")
            # Read the Dockerfile
            with open(dockerfile_path, 'r') as file:
                lines = file.readlines()

            # Comment out lines not needed for local container builds:
            # - download-checkpoints: checkpoints are mounted at runtime, not baked in
            # - downloader/unused.yaml: only exists in cloud build context
            modified_lines = []
            for line in lines:
                stripped = line.strip()
                if 'download-checkpoints' in line and '/home/nonroot/main' in line:
                    modified_lines.append(f'# {line}')
                elif stripped.startswith('COPY') and 'downloader/' in line:
                    modified_lines.append(f'# {line}')
                else:
                    modified_lines.append(line)

            # Create a temporary directory to store the modified Dockerfile
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_dockerfile_path = os.path.join(temp_dir, "Dockerfile.temp")

                # Write the modified Dockerfile to the temporary file
                with open(temp_dockerfile_path, 'w') as file:
                    file.writelines(modified_lines)

                # Build the Docker image using the temporary Dockerfile
                subprocess.check_call(
                    [
                        'docker',
                        'build',
                        '-t',
                        image_name,
                        '-f',
                        temp_dockerfile_path,
                        self.model_path,
                    ]
                )
            logger.info(f"Docker image '{image_name}' built successfully!")
        except subprocess.CalledProcessError as e:
            logger.info(f"Error occurred while building the Docker image: {e}")
            sys.exit(1)

    def docker_image_exists(self, image_name):
        """Check if the Docker image exists."""
        try:
            logger.info(f"Checking if Docker image '{image_name}' exists...")
            subprocess.run(["docker", "inspect", image_name], check=True)
            logger.info(f"Docker image '{image_name}' exists!")
            return True
        except subprocess.CalledProcessError:
            logger.info(f"Docker image '{image_name}' does not exist!")
            return False

    def _gpu_is_available(self):
        """
        Checks if nvidia-smi is available, indicating a GPU is likely accessible.
        """
        return shutil.which("nvidia-smi") is not None

    def _validate_test_environment(self):
        """
        Validate that the current environment supports model testing.
        Provides immediate feedback for unsupported configurations.
        This function runs only during CLI commands:
            1. clarifai model serve
            2. clarifai model local-test
        """
        warnings = []

        # Check Python version compatibility
        current_python = sys.version_info[:2]
        if current_python < MIN_REQUIRED_PYTHON_VERSION:
            logger.error(
                "❌ Environment validation failed: "
                f"Python {current_python[0]}.{current_python[1]} is not supported. "
                f"Please use Python {MIN_REQUIRED_PYTHON_VERSION[0]}.{MIN_REQUIRED_PYTHON_VERSION[1]} or higher for model testing."
            )
            sys.exit(1)
        elif current_python == MIN_REQUIRED_PYTHON_VERSION:
            warnings.append(
                f"Python {current_python[0]}.{current_python[1]} may have limited support. "
                f"Consider upgrading to Python {MIN_REQUIRED_PYTHON_VERSION[0]}.{MIN_REQUIRED_PYTHON_VERSION[1]} or higher for optimal compatibility."
            )

        # Check for GPU-related configurations on unsupported platforms
        current_environment = platform.system().lower()
        if current_environment == "darwin":
            current_environment = "darwin (macOS)"
        if not self._gpu_is_available():
            # Check if model requires GPU acceleration
            requires_gpu = False
            if hasattr(self, 'config') and self.config:
                inference_compute_info = self.config.get('inference_compute_info', {})
                num_accelerators = inference_compute_info.get('num_accelerators', 0)
                requires_gpu = (
                    num_accelerators != 0
                )  # Check if GPU is required based on num_accelerators

            if requires_gpu:
                logger.error(
                    "❌ Environment validation failed: "
                    f"Your current environment '{current_environment}' does not have NVIDIA GPU support. "
                    "Your model configuration requires GPU support. "
                    "Please test on an environment with NVIDIA GPU support, "
                    "or modify your model configuration to use CPU-only inference."
                )
                sys.exit(1)
            else:
                warnings.append(
                    f"Running on '{current_environment}' without NVIDIA GPU support. "
                    "If your model requires GPU support, testing may fail. "
                    "Consider testing on an environment with NVIDIA GPU support."
                )

        # Check for Docker availability when needed
        if hasattr(self, 'config') and self.config:
            # This is not a hard requirement but good to warn about
            if not shutil.which("docker"):
                warnings.append(
                    "Docker is not available. Container-based testing will not work. "
                    "Install Docker if you plan to test in containers."
                )

        # Log warnings
        for warning in warnings:
            logger.warning(f"⚠️ Environment Warning: {warning}")

        if warnings:
            logger.info("✅ Environment validation passed with warnings.")
        else:
            logger.info("✅ Environment validation passed.")

        return True

    def run_docker_container(
        self,
        image_name,
        container_name="clarifai-model-container",
        port=8080,
        env_vars=None,
        is_local_runner: bool = False,
        **kwargs,
    ):
        """Runs a Docker container from the specified image."""
        try:
            logger.info(
                f"Running Docker container '{container_name}' from image '{image_name}'..."
            )
            # Base docker run command
            cmd = ["docker", "run", "--name", container_name, '--rm', "--network", "host"]
            if self._gpu_is_available():
                cmd.extend(["--gpus", "all"])
            # Add volume mappings (use --mount to handle paths with colons, e.g. "gemma3:1b")
            cmd.extend(
                [
                    "--mount",
                    f"type=bind,source={self.model_path},target=/home/nonroot/main",
                ]
            )
            # Add environment variables
            if env_vars:
                for key, value in env_vars.items():
                    cmd.extend(["-e", f"{key}={value}"])
            # Set PYTHONDONTWRITEBYTECODE=1 to prevent __pycache__ folder generation
            cmd.extend(["-e", "PYTHONDONTWRITEBYTECODE=1"])
            # Add the image name
            cmd.append(image_name)
            # Override CMD to run the server (ENTRYPOINT is tini, so we need the full command)
            server_cmd = [
                "python",
                "-m",
                "clarifai.runners.server",
                "--model_path",
                "/home/nonroot/main",
            ]
            if is_local_runner:
                kwargs.pop("pool_size", None)  # remove pool_size if exists
                server_cmd.extend(
                    [
                        "--compute_cluster_id",
                        str(kwargs.get("compute_cluster_id", None)),
                        "--user_id",
                        str(kwargs.get("user_id", None)),
                        "--nodepool_id",
                        str(kwargs.get("nodepool_id", None)),
                        "--runner_id",
                        str(kwargs.get("runner_id", None)),
                        "--base_url",
                        str(kwargs.get("base_url", None)),
                        "--pat",
                        str(kwargs.get("pat", None)),
                        "--num_threads",
                        str(kwargs.get("num_threads", 0)),
                    ]
                )
                # Only pass health_check_port if non-zero (avoids compat issues with
                # older clarifai versions inside the container that lack the flag)
                hcp = kwargs.get("health_check_port", 8080)
                if hcp and hcp > 0:
                    server_cmd.extend(["--health_check_port", str(hcp)])
            else:
                server_cmd.extend(["--grpc", "--port", str(port)])
            cmd.extend(server_cmd)
            # Run the container
            logger.info(f"Running docker commands: {cmd}")
            process = subprocess.Popen(
                cmd,
            )
            logger.info(
                f"Docker container '{container_name}' is running successfully! access the model at http://localhost:{port}"
            )

            # Function to handle Ctrl+C (SIGINT) gracefully
            def signal_handler(sig, frame):
                logger.info(f"Stopping Docker container '{container_name}'...")
                subprocess.run(["docker", "stop", container_name], check=True)
                process.terminate()
                logger.info(f"Docker container '{container_name}' stopped successfully!")
                time.sleep(1)
                sys.exit(0)

            # Register the signal handler for SIGINT (Ctrl+C)
            signal.signal(signal.SIGINT, signal_handler)
            # Wait for the process to finish (keeps the container running until it's stopped)
            process.wait()
        except subprocess.CalledProcessError as e:
            logger.info(f"Error occurred while running the Docker container: {e}")
            sys.exit(1)
        except Exception as e:
            logger.info(f"Error occurred while running the Docker container: {e}")
            sys.exit(1)

    def test_model_container(
        self, image_name, container_name="clarifai-model-container", env_vars=None
    ):
        """Test the model inside the Docker container."""
        try:
            logger.info("Testing the model inside the Docker container...")
            # Base docker run command
            cmd = ["docker", "run", "--name", container_name, '--rm', "--network", "host"]
            if self._gpu_is_available():
                cmd.extend(["--gpus", "all"])
            # update the entrypoint for testing the model
            cmd.extend(["--entrypoint", "python"])
            # Add volume mappings (use --mount to handle paths with colons, e.g. "gemma3:1b")
            cmd.extend(
                [
                    "--mount",
                    f"type=bind,source={self.model_path},target=/home/nonroot/main",
                ]
            )
            # Add environment variables
            if env_vars:
                for key, value in env_vars.items():
                    cmd.extend(["-e", f"{key}={value}"])
            # Set PYTHONDONTWRITEBYTECODE=1 to prevent __pycache__ folder generation
            cmd.extend(["-e", "PYTHONDONTWRITEBYTECODE=1"])
            # Add the image name
            cmd.append(image_name)
            # update the CMD to test the model inside the container
            cmd.extend(
                [
                    "-c",
                    "from clarifai.runners.models.model_run_locally import ModelRunLocally; ModelRunLocally('/home/nonroot/main')._run_test()",
                ]
            )
            # Run the container
            subprocess.check_call(cmd)
            logger.info("Model tested successfully!")
        except subprocess.CalledProcessError as e:
            logger.error(f"Error testing the model inside the Docker container: {e}")
            sys.exit(1)

    def container_exists(self, container_name="clarifai-model-container"):
        """Check if the Docker container exists."""
        try:
            # Run docker ps -a to list all containers (running and stopped)
            result = subprocess.run(
                [
                    "docker",
                    "ps",
                    "-a",
                    "--filter",
                    f"name={container_name}",
                    "--format",
                    "{{.Names}}",
                ],
                check=True,
                capture_output=True,
                text=True,
            )
            # If the container name is returned, it exists
            if result.stdout.strip() == container_name:
                logger.info(f"Docker container '{container_name}' exists.")
                return True
            else:
                return False
        except subprocess.CalledProcessError as e:
            logger.error(f"Error occurred while checking if container exists: {e}")
            return False

    def stop_docker_container(self, container_name="clarifai-model-container"):
        """Stop the Docker container if it's running."""
        try:
            # Check if the container is running
            result = subprocess.run(
                ["docker", "ps", "--filter", f"name={container_name}", "--format", "{{.Names}}"],
                check=True,
                capture_output=True,
                text=True,
            )
            if result.stdout.strip() == container_name:
                logger.info(f"Docker container '{container_name}' is running. Stopping it...")
                subprocess.run(["docker", "stop", container_name], check=True)
                logger.info(f"Docker container '{container_name}' stopped successfully!")
        except subprocess.CalledProcessError as e:
            logger.error(f"Error occurred while stopping the Docker container: {e}")

    def remove_docker_container(self, container_name="clarifai-model-container"):
        """Remove the Docker container."""
        try:
            logger.info(f"Removing Docker container '{container_name}'...")
            subprocess.run(["docker", "rm", container_name], check=True)
            logger.info(f"Docker container '{container_name}' removed successfully!")
        except subprocess.CalledProcessError as e:
            logger.error(f"Error occurred while removing the Docker container: {e}")

    def remove_docker_image(self, image_name):
        """Remove the Docker image."""
        try:
            logger.info(f"Removing Docker image '{image_name}'...")
            subprocess.run(["docker", "rmi", image_name], check=True)
            logger.info(f"Docker image '{image_name}' removed successfully!")
        except subprocess.CalledProcessError as e:
            logger.error(f"Error occurred while removing the Docker image: {e}")

    def clean_up(self):
        """Clean up the temporary virtual environment."""
        if os.path.exists(self.temp_dir):
            logger.info("Cleaning up temporary virtual environment...")
            shutil.rmtree(self.temp_dir)


def main(
    model_path,
    run_model_server=False,
    inside_container=False,
    port=8080,
    keep_env=False,
    keep_image=False,
    skip_dockerfile: bool = False,
):
    manager = ModelRunLocally(model_path)

    # Validate test environment early to provide immediate feedback
    logger.info("Validating test environment...")
    manager._validate_test_environment()

    # get whatever stage is in config.yaml to force download now
    # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
    if inside_container:
        if not manager.is_docker_installed():
            sys.exit(1)
        if not skip_dockerfile:
            manager.builder.create_dockerfile()
        image_tag = manager._docker_hash()
        model_id = manager.config['model']['id'].lower()
        # must be in lowercase
        image_name = f"{model_id}:{image_tag}"
        container_name = model_id
        if not manager.docker_image_exists(image_name):
            manager.build_docker_image(image_name=image_name)
        try:
            if run_model_server:
                manager.run_docker_container(
                    image_name=image_name, container_name=container_name, port=port
                )
            else:
                manager.test_model_container(image_name=image_name, container_name=container_name)
        finally:
            if manager.container_exists(container_name):
                manager.stop_docker_container(container_name)
                manager.remove_docker_container(container_name=container_name)
            if not keep_image:
                manager.remove_docker_image(image_name)

    else:
        try:
            use_existing_env = manager.create_temp_venv()
            if not use_existing_env:
                manager.install_requirements()
            if run_model_server:
                manager.run_model_server(port)
            else:
                manager.test_model()
        finally:
            if not keep_env:
                manager.clean_up()
