import asyncio
import json
import tempfile
import threading
import time
import urllib.parse
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Union

import websockets

from kubetorch.globals import config, LoggingConfig, service_url, service_url_async
from kubetorch.logger import get_logger
from kubetorch.provisioning.constants import DEFAULT_K8S_SERVICE_PORT
from kubetorch.provisioning.utils import has_k8s_credentials, KubernetesCredentialsError
from kubetorch.resources.callables.utils import get_names_for_reload_fallbacks, locate_working_dir
from kubetorch.resources.compute.utils import (
    ControllerRequestError,
    delete_resources_for_services,
    print_byo_deletion_warning,
    ServiceTimeoutError,
    VersionMismatchError,
)
from kubetorch.serving.http_client import HTTPClient
from kubetorch.serving.utils import clean_and_validate_k8s_name, generate_unique_request_id, is_running_in_kubernetes
from kubetorch.utils import (
    ColoredFormatter,
    extract_host_port,
    get_kt_install_url,
    iso_timestamp_to_nanoseconds,
    ServerLogsFormatter,
)

logger = get_logger(__name__)


class Module:
    MODULE_TYPE = None

    def __init__(
        self,
        name: str,
        pointers: tuple,
    ):
        self._compute = None
        self._service_config = None
        self._http_client = None
        self._get_if_exists = True
        self._reload_prefixes = None
        self._serialization = "json"  # Default serialization format
        self._async = False
        self._remote_pointers = None
        self._container_project_root = None
        self._service_name = None

        self.pointers = pointers
        self.name = clean_and_validate_k8s_name(name, allow_full_length=False) if name else None

    @property
    def module_name(self):
        """Name of the function or class."""
        return self.pointers[2]

    @property
    def reload_prefixes(self):
        return self._reload_prefixes or []

    @reload_prefixes.setter
    def reload_prefixes(self, value: Union[str, List[str]]):
        """Set the reload_prefixes property."""
        if isinstance(value, (list)):
            self._reload_prefixes = value
        elif isinstance(value, str):
            self._reload_prefixes = [value]
        else:
            raise ValueError("`reload_prefixes` must be a string or a list.")

    @property
    def namespace(self):
        """Namespace where the service is deployed."""
        if self.compute is not None:
            return self.compute.namespace
        return config.namespace

    @property
    def service_name(self):
        """Name of the knative service, formatted according to k8s regex rules."""
        if self._service_name:
            return self._service_name

        service_name = self.name

        if config.username and not self.reload_prefixes and not service_name.startswith(config.username + "-"):
            service_name = f"{config.username}-{service_name}"

        self._service_name = clean_and_validate_k8s_name(service_name, allow_full_length=True)
        return self._service_name

    @service_name.setter
    def service_name(self, value: str):
        self._service_name = clean_and_validate_k8s_name(value, allow_full_length=True)

    @property
    def compute(self):
        """Compute object corresponding to the module."""
        return self._compute

    @compute.setter
    def compute(self, compute: "Compute"):
        self._compute = compute

    @property
    def remote_pointers(self):
        if self._remote_pointers:
            return self._remote_pointers

        source_dir, _, _ = locate_working_dir(self.pointers[0])
        relative_module_path = Path(self.pointers[0]).expanduser().relative_to(source_dir)
        source_dir_name = Path(source_dir).name
        if self.compute.working_dir is not None:
            container_module_path = str(Path(self.compute.working_dir) / source_dir_name / relative_module_path)
            self._container_project_root = str(Path(self.compute.working_dir) / source_dir_name)
        else:
            # Leave it as relative path
            container_module_path = str(Path(source_dir_name) / relative_module_path)
            self._container_project_root = source_dir_name
        self._remote_pointers = (
            container_module_path,
            self.pointers[1],
            self.pointers[2],
        )
        return self._remote_pointers

    @property
    def container_project_root(self):
        """Returns the project root path in the container."""
        if self._container_project_root is None:
            # Trigger computation via remote_pointers
            _ = self.remote_pointers
        return self._container_project_root

    @property
    def service_config(self) -> dict:
        """Knative service configuration loaded from Kubernetes API."""
        return self._service_config

    @service_config.setter
    def service_config(self, value: dict):
        self._service_config = value

    @property
    def base_endpoint(self):
        """Endpoint for the module."""
        if is_running_in_kubernetes():
            if not self._compute.endpoint:
                return self._compute._wait_for_endpoint()
            return self._compute.endpoint

        if self._compute._endpoint_config and self._compute._endpoint_config.url:
            return self._compute._endpoint_config.get_proxied_url(self._compute.client_port())

        # URL format when using the NGINX proxy: /{namespace}/{service}:{port}/{path}
        return f"http://localhost:{self._compute.client_port()}/{self.namespace}/{self.service_name}:{DEFAULT_K8S_SERVICE_PORT}"

    @property
    def request_headers(self):
        return {}

    @property
    def serialization(self):
        """Default serialization format for this module.
        More info in the `Call Modes Guide <https://www.run.house/kubetorch/concepts/call-modes>`__."""
        return self._serialization

    @serialization.setter
    def serialization(self, value: str):
        """Set the default serialization format for this module."""
        if value not in ["json", "pickle"]:
            raise ValueError("Serialization must be 'json' or 'pickle'")
        self._serialization = value

    @property
    def async_(self):
        """Whether to run the function or class methods in async mode."""
        return self._async

    @async_.setter
    def async_(self, value: bool):
        if not isinstance(value, bool):
            raise ValueError("`async_` must be a boolean")
        self._async = value

    @property
    def logging_config(self) -> LoggingConfig:
        """Get the logging configuration for this module from its compute."""
        if self.compute:
            return self.compute.logging_config
        return LoggingConfig()

    @property
    def stream_logs(self) -> bool:
        """Determine if log streaming should be enabled.

        Fallback chain:
        1. logging_config.stream_logs if explicitly set (not None)
        2. Global config.stream_logs if explicitly set (not None)
        3. Default to True
        """
        if self.logging_config.stream_logs is not None:
            return self.logging_config.stream_logs
        return config.stream_logs  # Defaults to True if not set

    @classmethod
    def from_name(
        cls,
        name: str,
        namespace: str = None,
        reload_prefixes: Union[str, List[str]] = [],
    ):
        """Reload an existing callable by its service name."""
        import kubetorch as kt
        from kubetorch.provisioning.service_manager import ServiceManager

        controller_client = kt.globals.controller_client()
        namespace = namespace or config.namespace
        if isinstance(reload_prefixes, str):
            reload_prefixes = [reload_prefixes]
        potential_names = get_names_for_reload_fallbacks(name=name, prefixes=reload_prefixes)

        all_services = ServiceManager.discover_services(namespace=namespace)

        # Create name-to-service lookup for efficient searching
        # Prefer non-selector services over selector pools (which don't have env vars)
        service_dict = {}
        for svc in all_services:
            svc_name = svc["name"]
            if svc_name not in service_dict or service_dict[svc_name].get("template_type") == "selector":
                # Add if new, or replace selector pool with actual K8s resource
                service_dict[svc_name] = svc

        # Try to find the first matching service across all service types
        for candidate in potential_names:
            service_info = service_dict.get(candidate)
            if service_info is None:
                continue

            # Skip selector-based pools - they don't have template for reload
            if service_info.get("template_type") == "selector":
                continue

            compute = kt.Compute.from_template(service_info)

            pods_result = controller_client.list_pods(
                namespace=namespace,
                label_selector=f"kubetorch.com/service={candidate}",
            )
            volumes = []

            # TODO: handle case where service is scaled to 0?
            pod_items = pods_result.get("items", [])
            if pod_items:
                # Use runtime Pod spec
                pod = pod_items[0]
                volumes_list = pod.get("spec", {}).get("volumes", [])
                for v in volumes_list:
                    if v.get("persistentVolumeClaim"):
                        existing_volume = kt.Volume.from_name(name=v.get("name"))
                        volumes.append(existing_volume)

            # Get module info from controller's pool database
            pool_info = controller_client.get_pool(namespace, candidate)
            if not pool_info or not pool_info.get("module"):
                raise ValueError(f"No module info found for pool {candidate}")

            module_info = pool_info["module"]
            module_pointers = module_info.get("pointers", {})
            pointers = (
                module_pointers.get("file_path"),
                module_pointers.get("module_name"),
                module_pointers.get("cls_or_fn_name"),
            )

            callable_type = module_info.get("type", "fn")
            if callable_type == "cls":
                init_args = module_pointers.get("init_args") or {}
                reloaded_module = kt.Cls(name=candidate, pointers=pointers, init_args=init_args)
            elif callable_type == "fn":
                reloaded_module = kt.Fn(name=candidate, pointers=pointers)
            else:
                raise ValueError(f"Unknown module type: {callable_type}")

            reloaded_module.service_name = candidate
            reloaded_module.compute = compute
            return reloaded_module

        raise ValueError(
            f"Service '{name}' not found in namespace '{namespace}' with reload_prefixes={reload_prefixes}"
        )

    def _client(self, *args, **kwargs):
        """Return the client through which to interact with the remote Module.
        If compute is not yet set, attempt to reload it.
        """
        if self._http_client is not None:
            return self._http_client

        if self.compute is None or self.service_config is None:
            namespace = self.namespace
            # When rebuilding the http client on reload, need to know whether to look for a prefix
            reload_prefixes = self.reload_prefixes
            logger.debug(
                f"Attempting to reload service '{self.service_name}' in namespace '{namespace}' with "
                f"reload_prefixes={reload_prefixes}"
            )
            reloaded_module = Module.from_name(
                name=self.service_name,
                namespace=namespace,
                reload_prefixes=reload_prefixes,
            )

            # Update settable attributes with reloaded module values
            self.compute = self.compute or reloaded_module.compute
            self.service_config = reloaded_module.service_config
            self.pointers = reloaded_module.pointers
            self.name = reloaded_module.name
            self.service_name = reloaded_module.service_name

        self._http_client = HTTPClient(
            base_url=self.endpoint(*args, **kwargs),
            compute=self.compute,
            service_name=self.service_name,
        )

        return self._http_client

    def endpoint(self, method_name: str = None):
        if not hasattr(self, "init_args"):
            return f"{self.base_endpoint}/{self.module_name}"
        else:
            return f"{self.base_endpoint}/{self.module_name}/{method_name}"

    def deploy(self):
        """
        Helper method to deploy modules specified by the @compute decorator. Used by `kt deploy` CLI command.
        Deploys the module to the specified compute.
        """
        if self.compute is None:
            raise ValueError("Compute must be set before deploying the module.")
        return self.to(self.compute, init_args=getattr(self, "init_args", None))

    async def deploy_async(self):
        """
        Async helper method to deploy modules specified by the @compute decorator. Used by `kt deploy` CLI command
        when multiple modules are present. Deploys the module to the specified compute asynchronously.
        """
        if self.compute is None:
            raise ValueError("Compute must be set before deploying the module.")
        return await self.to_async(self.compute, init_args=getattr(self, "init_args", None))

    def to(
        self,
        compute: "Compute",
        init_args: Dict = None,
        stream_logs: Union[bool, None] = None,
        get_if_exists: bool = False,
        reload_prefixes: Union[str, List[str]] = [],
        dryrun: bool = False,
    ):
        """
        Send the function or class to the specified compute.

        Args:
            compute (Compute): The compute to send the function or class to.
            init_args (Dict, optional): Initialization arguments, which may be relevant for a class.
            stream_logs (bool, optional): Whether to stream logs during service launch. If None, uses
                the compute's logging_config.stream_logs setting, or falls back to global config.stream_logs.
            get_if_exists (Union[bool, List[str]], optional): Controls how service lookup is performed to determine
                whether to send the service to the compute.

                - If False (default): Do not attempt to reload the service.
                - If True: Attempt to find an existing service using a standard fallback order
                  (e.g., username, git branch, then prod). If found, re-use that existing service.
            reload_prefixes (Union[str, List[str]], optional): A list of prefixes to use when reloading the function
                (e.g., ["qa", "prod", "git-branch-name"]). If not provided, will use the current username,
                git branch, and prod.
            dryrun (bool, optional): Whether to setup and return the object as a dryrun (``True``),
                or to actually launch the compute and service (``False``).
        Returns:
            Module: The module instance.

        Example:

        .. code-block:: python

            import kubetorch as kt

            remote_cls = kt.cls(SlowNumpyArray, name=name).to(
                kt.Compute(cpus=".1"),
                init_args={"size": 10},
                stream_logs=True
            )
        """
        if not has_k8s_credentials():
            raise KubernetesCredentialsError(
                "Kubernetes credentials not found. Please ensure you are running in a Kubernetes cluster or have a valid kubeconfig file."
            )

        if compute.service_name and compute.service_name != self.service_name:
            logger.info(f"Renaming service to match compute service name {compute.service_name}")
            self.service_name = compute.service_name

        if get_if_exists:
            try:
                existing_service = self._get_existing_service(reload_prefixes)
                if existing_service:
                    logger.debug(f"Reusing existing service: {existing_service.service_name}")
                    return existing_service
            except Exception as e:
                logger.debug(
                    f"Service {self.service_name} not found in namespace {compute.namespace} "
                    f"with reload_prefixes={reload_prefixes}: {str(e)}"
                )

        self.compute = compute
        self.compute.service_name = self.service_name

        if hasattr(self, "init_args"):
            self.init_args = init_args

        # We need the deployment timestamp at the start of the update so we know that artifacts deployed **after**
        # this time are part of the current deployment. We actually set it at the end to ensure that the deployment is
        # successful.
        logger.debug(f"Deploying module: {self.service_name}")
        deployment_timestamp = datetime.now(timezone.utc).isoformat()
        install_url, use_editable = get_kt_install_url(self.compute.freeze)

        if not dryrun and not self.compute.freeze:
            self._rsync_repo_and_image_patches(install_url, use_editable, init_args)

        self._launch_service(
            install_url,
            use_editable,
            init_args,
            deployment_timestamp,
            stream_logs,
            dryrun,
        )

        return self

    async def to_async(
        self,
        compute: "Compute",
        init_args: Dict = None,
        stream_logs: Union[bool, None] = None,
        get_if_exists: bool = False,
        reload_prefixes: Union[str, List[str]] = [],
        dryrun: bool = False,
    ):
        """
        Async version of the `.to` method. Send the function or class to the specified compute asynchronously.

        Args:
            compute (Compute): The compute to send the function or class to.
            init_args (Dict, optional): Initialization arguments, which may be relevant for a class.
            stream_logs (bool, optional): Whether to stream logs during service launch. If None, uses
                the compute's logging_config.stream_logs setting, or falls back to global config.stream_logs.
            get_if_exists (Union[bool, List[str]], optional): Controls how service lookup is performed to determine
                whether to send the service to the compute.

                - If False (default): Do not attempt to reload the service.
                - If True: Attempt to find an existing service using a standard fallback order
                  (e.g., username, git branch, then prod). If found, re-use that existing service.
            reload_prefixes (Union[str, List[str]], optional): A list of prefixes to use when reloading the function
                (e.g., ["qa", "prod", "git-branch-name"]). If not provided, will use the current username,
                git branch, and prod.
            dryrun (bool, optional): Whether to setup and return the object as a dryrun (``True``),
                or to actually launch the compute and service (``False``).
        Returns:
            Module: The module instance.

        Example:

        .. code-block:: python

            import kubetorch as kt

            remote_cls = await kt.cls(SlowNumpyArray, name=name).to_async(
                kt.Compute(cpus=".1"),
                init_args={"size": 10},
                stream_logs=True
            )
        """
        if compute.service_name and compute.service_name != self.service_name:
            logger.info(f"Renaming service to match compute service name {compute.service_name}")
            self.service_name = compute.service_name

        if get_if_exists:
            try:
                existing_service = await self._get_existing_service_async(reload_prefixes)
                if existing_service:
                    logger.debug(f"Reusing existing service: {existing_service.service_name}")
                    return existing_service
            except Exception as e:
                logger.info(
                    f"Service {self.service_name} not found in namespace {compute.namespace} "
                    f"with reload_prefixes={reload_prefixes}: {str(e)}"
                )

        self.compute = compute
        self.compute.service_name = self.service_name

        if hasattr(self, "init_args"):
            self.init_args = init_args

        logger.debug(f"Deploying module: {self.service_name}")
        deployment_timestamp = datetime.now(timezone.utc).isoformat()
        install_url, use_editable = get_kt_install_url(self.compute.freeze)

        if not dryrun and not self.compute.freeze:
            await self._rsync_repo_and_image_patches_async(install_url, use_editable, init_args)

        await self._launch_service_async(
            install_url,
            use_editable,
            init_args,
            deployment_timestamp,
            stream_logs,
            dryrun,
        )

        return self

    def _get_existing_service(self, reload_prefixes):
        try:
            existing_service = Module.from_name(
                self.service_name,
                namespace=self.namespace,
                reload_prefixes=reload_prefixes,
            )
            if existing_service:
                if self.compute:
                    # Replace the compute object, if the user has already constructed it locally
                    existing_service.compute = self.compute
                logger.info(
                    f"Existing service '{self.service_name}' found in namespace '{self.namespace}', not "
                    f"redeploying."
                )
                return existing_service
        except Exception as e:
            raise ValueError(
                f"Failed to reload service {self.service_name} in namespace {self.namespace} "
                f"and reload_prefixes={reload_prefixes}: {str(e)}"
            )

    async def _get_existing_service_async(self, reload_prefixes):
        try:
            existing_service = Module.from_name(
                self.service_name,
                namespace=self.namespace,
                reload_prefixes=reload_prefixes,
            )
            if existing_service:
                if self.compute:
                    # Replace the compute object, if the user has already constructed it locally
                    existing_service.compute = self.compute
                logger.info(
                    f"Existing service '{self.service_name}' found in namespace '{self.namespace}', not "
                    f"redeploying."
                )
                return existing_service
        except Exception as e:
            raise ValueError(
                f"Failed to reload service {self.service_name} in namespace {self.namespace} "
                f"and reload_prefixes={reload_prefixes}: {str(e)}"
            )

    def _get_rsync_dirs_and_dockerfile(self, install_url, use_editable, init_args):
        source_dir, has_kt_dir, matched_file = locate_working_dir(self.pointers[0])
        rsync_dirs = [str(source_dir)]
        if not has_kt_dir:
            # Use the source file (.py) instead of directory
            source_file = Path(f"{self.pointers[0]}/{self.pointers[1]}.py")
            rsync_dirs = [str(source_file)]
            logger.info(f"No project markers found, syncing file {source_file}")
        else:
            source_dir_name = Path(source_dir).name
            if self.compute.working_dir:
                remote_path = f"{self.compute.working_dir}/{source_dir_name}"
                logger.info(
                    f"Syncing project directory {source_dir} -> {remote_path} (working directory detected via {matched_file})"
                )
            else:
                logger.info(
                    f"Syncing project directory {source_dir} -> ./{source_dir_name} (working directory detected via {matched_file})"
                )

        if install_url.endswith(".whl") or (use_editable and install_url != str(source_dir)):
            rsync_dirs.append(install_url)

        # Module metadata is now sent via controller WebSocket, not baked into dockerfile
        service_dockerfile = self._get_service_dockerfile()
        return rsync_dirs, service_dockerfile

    def _rsync_repo_and_image_patches(self, install_url, use_editable, init_args):
        logger.debug("Rsyncing data to the rsync pod")
        rsync_dirs, service_dockerfile = self._get_rsync_dirs_and_dockerfile(install_url, use_editable, init_args)
        self._construct_and_rsync_files(rsync_dirs, service_dockerfile)
        logger.debug(f"Rsync completed for service {self.service_name}")

    async def _rsync_repo_and_image_patches_async(self, install_url, use_editable, init_args):
        logger.debug("Rsyncing data to the rsync pod")
        rsync_dirs, service_dockerfile = self._get_rsync_dirs_and_dockerfile(install_url, use_editable, init_args)
        await self._construct_and_rsync_files_async(rsync_dirs, service_dockerfile)
        logger.debug(f"Rsync completed for service {self.service_name}")

    def _launch_service(
        self,
        install_url,
        use_editable,
        init_args,
        deployment_timestamp,
        stream_logs,
        dryrun,
    ):
        # Start log streaming if enabled
        stop_event = threading.Event()
        log_thread = None

        # Resolve stream_logs using module's property if not explicitly set
        stream_logs = stream_logs if stream_logs is not None else self.stream_logs

        launch_request_id = "-"
        if stream_logs and not dryrun:
            # Create a unique request ID for this launch sequence
            launch_request_id = f"launch_{generate_unique_request_id('launch', deployment_timestamp)}"

            # Start log streaming in a separate thread
            log_thread = threading.Thread(
                target=self._stream_launch_logs,
                args=(
                    launch_request_id,
                    stop_event,
                    deployment_timestamp,
                ),
            )
            log_thread.daemon = True
            log_thread.start()

        try:
            startup_rsync_command = self._startup_rsync_command(use_editable, install_url, dryrun)

            # Generate dockerfile (module metadata sent via controller WebSocket)
            dockerfile = self._get_service_dockerfile()

            # Build module spec for pool registration
            dispatch = self.compute.dispatch_method

            module_metadata = {
                "type": self.MODULE_TYPE,
                "pointers": {
                    "file_path": self.remote_pointers[0],
                    "module_name": self.remote_pointers[1],
                    "cls_or_fn_name": self.remote_pointers[2],
                    "project_root": self.container_project_root,
                    "init_args": init_args,
                },
                "dispatch": dispatch,
                "procs": 1,
            }

            # Launch the compute in the form of a service with the requested resources
            # Note: module metadata (pointers, init_args) is now sent via controller WebSocket
            service_config = self.compute._launch(
                service_name=self.compute.service_name,
                install_url=install_url if not use_editable else None,
                module_name=self.remote_pointers[1],  # module_name for labels
                startup_rsync_command=startup_rsync_command,
                deployment_timestamp=deployment_timestamp,
                dryrun=dryrun,
                dockerfile=dockerfile,
                module=module_metadata,
            )
            self.service_config = service_config

            if not dryrun:
                self.compute._check_service_ready()
                # Additional health check to ensure HTTP server is ready
                self._wait_for_http_health()
        finally:
            # Stop log streaming
            if log_thread:
                stop_event.set()

    async def _launch_service_async(
        self,
        install_url,
        use_editable,
        init_args,
        deployment_timestamp,
        stream_logs,
        dryrun,
    ):
        # Start log streaming if enabled
        stop_event = asyncio.Event()
        log_task = None

        # Resolve stream_logs using module's property if not explicitly set
        stream_logs = stream_logs if stream_logs is not None else self.stream_logs

        launch_request_id = "-"
        if stream_logs and not dryrun:
            # Create a unique request ID for this launch sequence
            launch_request_id = f"launch_{generate_unique_request_id('launch', deployment_timestamp)}"

            # Start log streaming as an async task
            log_task = asyncio.create_task(
                self._stream_launch_logs_async(
                    launch_request_id,
                    stop_event,
                    deployment_timestamp,
                )
            )

        try:
            startup_rsync_command = self._startup_rsync_command(use_editable, install_url, dryrun)

            # Generate dockerfile (module metadata sent via controller WebSocket)
            dockerfile = self._get_service_dockerfile()

            # Build module spec for pool registration
            dispatch = self.compute.dispatch_method

            module_metadata = {
                "type": self.MODULE_TYPE,
                "pointers": {
                    "file_path": self.remote_pointers[0],
                    "module_name": self.remote_pointers[1],
                    "cls_or_fn_name": self.remote_pointers[2],
                    "project_root": self.container_project_root,
                    "init_args": init_args,
                },
                "dispatch": dispatch,
                "procs": 1,
            }

            # Launch the compute in the form of a service with the requested resources
            # Note: module metadata (pointers, init_args) is now sent via controller WebSocket
            service_config = await self.compute._launch_async(
                service_name=self.compute.service_name,
                install_url=install_url if not use_editable else None,
                module_name=self.remote_pointers[1],  # module_name for labels
                startup_rsync_command=startup_rsync_command,
                deployment_timestamp=deployment_timestamp,
                dryrun=dryrun,
                dockerfile=dockerfile,
                module=module_metadata,
            )
            self.service_config = service_config

            if not dryrun:
                await self.compute._check_service_ready_async()
                await self._wait_for_http_health_async()
        finally:
            # Stop log streaming
            if log_task:
                stop_event.set()
                try:
                    await asyncio.wait_for(log_task, timeout=2.0)
                except asyncio.TimeoutError:
                    log_task.cancel()
                    try:
                        await log_task
                    except asyncio.CancelledError:
                        pass

    def _get_service_dockerfile(self):
        # Module metadata is now sent via controller WebSocket, not baked into dockerfile
        image_instructions = self.compute._image_setup_and_instructions()
        logger.debug(f"Generated Dockerfile for service {self.service_name}:\n{image_instructions}")
        return image_instructions

    def _construct_and_rsync_files(self, rsync_dirs, service_dockerfile):
        with tempfile.TemporaryDirectory() as tmpdir:
            temp_file = Path(tmpdir) / ".kt" / "image.dockerfile"
            temp_file.parent.mkdir(parents=True, exist_ok=True)
            temp_file.write_text(service_dockerfile)

            # Add .kt directory to the list of directories to sync
            source_dir = str(Path(tmpdir) / ".kt")
            all_dirs = rsync_dirs + [source_dir]

            logger.debug(f"Rsyncing directories: {all_dirs}")
            # Use DataStoreClient KV interface - files go to rsync pod, then sync to service pods at startup
            # Using contents=False preserves directory structure - each directory becomes a subdirectory
            # This matches the old behavior where python_client/ and .kt/ were both preserved as subdirectories
            from kubetorch import data_store

            dt_client = data_store.DataStoreClient(namespace=self.compute.namespace)
            # Use absolute key path (starting with /) to prevent auto-prepending of caller's service name
            # This is critical when a parent service creates a child service - we want files at
            # /data/{namespace}/{child_service_name}/, not /data/{namespace}/{parent_service_name}/{child_service_name}/
            dt_client.put(key=f"/{self.compute.service_name}", src=all_dirs, contents=False)

    async def _construct_and_rsync_files_async(self, rsync_dirs, service_dockerfile):
        with tempfile.TemporaryDirectory() as tmpdir:
            temp_file = Path(tmpdir) / ".kt" / "image.dockerfile"
            temp_file.parent.mkdir(parents=True, exist_ok=True)
            temp_file.write_text(service_dockerfile)

            # Add .kt directory to the list of directories to sync
            source_dir = str(Path(tmpdir) / ".kt")
            all_dirs = rsync_dirs + [source_dir]

            logger.debug(f"Rsyncing directories: {all_dirs}")
            # Use DataStoreClient KV interface - files go to rsync pod, then sync to service pods at startup
            # Using contents=False preserves directory structure - each directory becomes a subdirectory
            # This matches the old behavior where python_client/ and .kt/ were both preserved as subdirectories
            from kubetorch import data_store

            dt_client = data_store.DataStoreClient(namespace=self.compute.namespace)
            # Use absolute key path (starting with /) to prevent auto-prepending of caller's service name
            # This is critical when a parent service creates a child service - we want files at
            # /data/{namespace}/{child_service_name}/, not /data/{namespace}/{parent_service_name}/{child_service_name}/
            dt_client.put(key=f"/{self.compute.service_name}", src=all_dirs, contents=False)

    def _startup_rsync_command(self, use_editable, install_url, dryrun):
        if dryrun:
            return None

        if use_editable or (install_url and install_url.endswith(".whl")):
            # rsync from the rsync pod's file system directly
            startup_cmd = self.compute._rsync_svc_url()
            cmd = f"rsync -av {startup_cmd} ."
            return cmd

        return None

    def teardown(self):
        """Delete the service and its associated resources."""
        logger.debug(f"Deleting service {self.service_name}")

        try:
            delete_result = delete_resources_for_services(
                services=[self.service_name],
                namespace=self.namespace,
                force=True,
                exact_match=True,
            )
            msg = f"Successfully force deleted {self.service_name}"

            # BYO compute: kubetorch does not own the K8s resource lifecycle
            byo_deleted_services = delete_result.pop("byo_deleted_services")
            if byo_deleted_services:
                print_byo_deletion_warning(byo_deleted_services)
                msg_suffix = "."
            else:
                msg_suffix = " and its associated resources."
            logger.info(f"{msg}{msg_suffix}")
        except ControllerRequestError as e:
            error_msg = str(e).split(":")[-1]
            logger.error(f"Failed to delete {self.service_name}: {error_msg}")

    def _stream_launch_logs(
        self,
        request_id: str,
        stop_event: threading.Event,
        deployment_timestamp: str,
    ):
        """Stream logs and events during service launch sequence.

        Args:
            request_id: Unique ID for this launch sequence
            stop_event: Event to signal when to stop streaming
            deployment_timestamp: Timestamp to filter logs after
        """
        try:
            # Query using labels set by LogCapture (service, namespace)
            # Use regex for service name to capture nested/child service logs (e.g., parent-service-child)
            # Child services have their own request_ids, so we don't filter by request_id here
            pod_query = f'{{service=~"{self.service_name}.*", namespace="{self.namespace}"}}'
            # Event query for K8s events pushed by controller's event watcher
            # Include service name pattern AND actual pod names for selector-only mode
            name_patterns = [f"{self.service_name}.*"]
            name_regex = "|".join(name_patterns)
            event_query = f'{{job="kubetorch-events", namespace="{self.namespace}", name=~"{name_regex}"}}'

            encoded_pod_query = urllib.parse.quote_plus(pod_query)
            encoded_event_query = urllib.parse.quote_plus(event_query)
            logger.debug(f"Streaming launch logs and events for service {self.service_name}")

            def start_log_threads(host, port):
                def run_pod_logs():
                    self._run_log_stream(
                        request_id,
                        stop_event,
                        host,
                        port,
                        encoded_pod_query,
                        deployment_timestamp,
                        namespace=self.namespace,
                        dedup=True,
                    )

                def run_event_logs():
                    self._run_log_stream(
                        request_id,
                        stop_event,
                        host,
                        port,
                        encoded_event_query,
                        deployment_timestamp,
                        namespace=self.namespace,
                    )

                pod_thread = threading.Thread(target=run_pod_logs, daemon=True)

                pod_thread.start()

                # Only start event log thread if include_events is enabled
                if self.logging_config.include_events:
                    event_thread = threading.Thread(target=run_event_logs, daemon=True)
                    event_thread.start()
                    event_thread.join(timeout=1.0)

                # Don't block indefinitely on joins - use short timeouts
                pod_thread.join(timeout=1.0)

            base_url = service_url()
            host, port = extract_host_port(base_url)
            logger.debug(f"Streaming launch logs with url={base_url} host={host} and local port {port}")
            start_log_threads(host, port)

        except Exception as e:
            logger.error(f"Failed to stream launch logs: {e}")
            raise e

    async def _stream_launch_logs_async(
        self,
        request_id: str,
        stop_event: asyncio.Event,
        deployment_timestamp: str,
    ):
        """Async version of _stream_launch_logs. Stream logs and events during service launch sequence.

        Args:
            request_id: Unique ID for this launch sequence
            stop_event: Event to signal when to stop streaming
            deployment_timestamp: Timestamp to filter logs after
        """
        try:
            # Query using labels set by LogCapture (service, namespace)
            # Use regex for service name to capture nested/child service logs (e.g., parent-service-child)
            # Child services have their own request_ids, so we don't filter by request_id here
            pod_query = f'{{service=~"{self.service_name}.*", namespace="{self.namespace}"}}'
            # Event query for K8s events pushed by controller's event watcher
            # Include service name pattern AND actual pod names for selector-only mode
            name_patterns = [f"{self.service_name}.*"]
            name_regex = "|".join(name_patterns)
            event_query = f'{{job="kubetorch-events", namespace="{self.namespace}", name=~"{name_regex}"}}'

            encoded_pod_query = urllib.parse.quote_plus(pod_query)
            encoded_event_query = urllib.parse.quote_plus(event_query)
            logger.debug(f"Streaming launch logs and events for service {self.service_name}")

            base_url = await service_url_async()
            host, port = extract_host_port(base_url)
            logger.debug(f"Streaming launch logs with url={base_url} host={host} and local port {port}")

            # Create async tasks for log streams
            tasks = []

            pod_task = asyncio.create_task(
                self._stream_logs_websocket(
                    request_id,
                    stop_event,
                    host=host,
                    port=port,
                    query=encoded_pod_query,
                    deployment_timestamp=deployment_timestamp,
                    namespace=self.namespace,
                    dedup=True,
                )
            )
            tasks.append(pod_task)

            # Only create event task if include_events is enabled
            if self.logging_config.include_events:
                event_task = asyncio.create_task(
                    self._stream_logs_websocket(
                        request_id,
                        stop_event,
                        host=host,
                        port=port,
                        query=encoded_event_query,
                        deployment_timestamp=deployment_timestamp,
                        namespace=self.namespace,
                    )
                )
                tasks.append(event_task)

            # Wait for tasks to complete or be cancelled
            try:
                await asyncio.gather(*tasks, return_exceptions=True)
            except Exception as e:
                logger.error(f"Error in async log streaming: {e}")

        except Exception as e:
            logger.error(f"Failed to stream launch logs: {e}")
            raise e

    def _run_log_stream(
        self,
        request_id: str,
        stop_event: threading.Event,
        host: str,
        port: int,
        query: str,
        deployment_timestamp: str,
        namespace: str,
        dedup: bool = False,
    ):
        """Helper to run log streaming in an event loop"""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(
                self._stream_logs_websocket(
                    request_id,
                    stop_event,
                    host=host,
                    port=port,
                    query=query,
                    deployment_timestamp=deployment_timestamp,
                    namespace=namespace,
                    dedup=dedup,
                )
            )
        finally:
            # Cancel all pending tasks to prevent "Task was destroyed but it is pending!" warnings
            pending = asyncio.all_tasks(loop)
            for task in pending:
                task.cancel()
            if pending:
                loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
            loop.close()

    async def _run_log_stream_async(
        self,
        request_id: str,
        stop_event: asyncio.Event,
        host: str,
        port: int,
        query: str,
        deployment_timestamp: str,
        namespace: str,
        dedup: bool = False,
    ):
        """Async helper to run log streaming directly in the current event loop"""
        await self._stream_logs_websocket(
            request_id,
            stop_event,
            host=host,
            port=port,
            query=query,
            deployment_timestamp=deployment_timestamp,
            namespace=namespace,
            dedup=dedup,
        )

    async def _stream_logs_websocket(
        self,
        request_id: str,
        stop_event: Union[threading.Event, asyncio.Event],
        host: str,
        port: int,
        query: str,
        deployment_timestamp: str,
        namespace: str,
        dedup: bool = False,
    ):
        """Stream logs and events using Loki's websocket tail endpoint.

        Args:
            request_id: Unique ID for this launch sequence
            stop_event: Event to signal when to stop streaming
            host: WebSocket host
            port: WebSocket port
            query: Loki query string
            deployment_timestamp: Timestamp to filter logs after
            namespace: Namespace to stream logs from
            dedup: Whether to deduplicate log messages
        """
        # Map logging_config.level to filtering behavior
        # "debug" -> show all, "info" -> hide debug, "warning"/"error" -> only errors
        log_level = self.logging_config.level.lower() if self.logging_config.level else "info"

        try:
            # Track most recent deployment timestamp to filter out old logs / events
            start_timestamp = iso_timestamp_to_nanoseconds(deployment_timestamp)

            # Namespace-aware Loki URL - routes to data store in the target namespace
            # Include start timestamp to avoid showing old events/logs
            uri = f"ws://{host}:{port}/loki/{namespace}/api/v1/tail?query={query}&start={start_timestamp}"

            # Track the last timestamp we've seen to avoid duplicates
            last_timestamp = None

            # Track when we should stop
            stop_time = None

            shown_event_messages = set()

            # Track seen log messages for deduplication
            seen_log_messages = set() if dedup else None

            # For formatting the server setup logs
            formatters = {}
            base_formatter = ServerLogsFormatter()
            websocket = None
            try:
                # Add timeout to prevent hanging connections
                websocket = await websockets.connect(
                    uri,
                    close_timeout=10,  # Max time to wait for close handshake
                    ping_interval=20,  # Send ping every 20 seconds
                    ping_timeout=10,  # Wait 10 seconds for pong
                )
                while True:
                    # If stop event is set, start counting down
                    # Handle both threading.Event and asyncio.Event
                    is_stop_set = stop_event.is_set() if hasattr(stop_event, "is_set") else stop_event.is_set()
                    if is_stop_set and stop_time is None:
                        stop_time = time.time() + self.logging_config.grace_period

                    # If we're past the grace period, exit
                    if stop_time is not None and time.time() > stop_time:
                        break

                    try:
                        # Use shorter timeout during grace period
                        timeout = (
                            self.logging_config.grace_poll_timeout
                            if stop_time is not None
                            else self.logging_config.poll_timeout
                        )
                        message = await asyncio.wait_for(websocket.recv(), timeout=timeout)
                        data = json.loads(message)

                        if data.get("streams"):
                            for stream in data["streams"]:
                                labels = stream.get("stream", {})
                                # Detect events by job label (pushed by controller's event watcher)
                                is_event = labels.get("job") == "kubetorch-events"
                                for value in stream["values"]:
                                    ts_ns = int(value[0])
                                    if start_timestamp is not None and ts_ns < start_timestamp:
                                        continue
                                    log_line = value[1]
                                    if is_event:
                                        event_type = labels.get("event_type", "Normal")
                                        reason = labels.get("reason", "")

                                        # Skip Normal events when log level is warning or error
                                        if log_level in ["warning", "error"] and event_type == "Normal":
                                            continue

                                        # Parse the event message from JSON payload
                                        try:
                                            event_data = json.loads(log_line)
                                            msg = event_data.get("message", log_line)
                                        except json.JSONDecodeError:
                                            msg = log_line

                                        # Skip expected probe failures during setup
                                        if reason == "Unhealthy" and (
                                            "HTTP probe failed with statuscode: 503" in msg
                                            or "Startup probe failed" in msg
                                        ):
                                            continue

                                        # Ignore noisy events
                                        ignore_patterns = (
                                            "queue-proxy",
                                            "resolving reference: address not set for kind = service",
                                            "failed to get private k8s service endpoints:",
                                        )
                                        if any(pattern in msg.lower() for pattern in ignore_patterns):
                                            continue

                                        # Only show unique event messages
                                        if msg in shown_event_messages:
                                            continue
                                        shown_event_messages.add(msg)

                                        # Format timestamp from event
                                        from datetime import datetime

                                        try:
                                            event_ts = datetime.fromtimestamp(ts_ns / 1e9).strftime("%Y-%m-%d %H:%M:%S")
                                        except Exception:
                                            event_ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

                                        # Use yellow for warnings, green for normal
                                        if event_type == "Warning":
                                            color = ColoredFormatter.get_color("yellow")
                                        else:
                                            color = ColoredFormatter.get_color("green")
                                        reset = ColoredFormatter.get_color("reset")

                                        # Format like metrics: ({service} events) timestamp | reason: message
                                        service_name = self.compute.service_name
                                        prefix = f"({service_name} events)"
                                        if event_type == "Normal":
                                            if log_level in ["debug", "info"]:
                                                print(f"{color}{prefix} {event_ts} | {reason}: {msg}{reset}")
                                        else:
                                            print(f"{color}{prefix} {event_ts} | {reason}: {msg}{reset}")
                                        continue

                                    if last_timestamp is not None and value[0] <= last_timestamp:
                                        continue
                                    last_timestamp = value[0]

                                    log_request_id = labels.get("request_id", "-")
                                    if log_request_id not in ("-", request_id):
                                        continue

                                    if log_level in ["debug", "info"]:
                                        try:
                                            log_dict = json.loads(log_line)
                                        except json.JSONDecodeError:
                                            # setup steps pre server start are not JSON formatted
                                            log_dict = None

                                        if log_dict is not None:
                                            # Use pod from Loki stream labels (set by LogCapture)
                                            pod_name = labels.get("pod", request_id)
                                            levelname = log_dict.get("levelname", "INFO")
                                            ts = log_dict.get("asctime")
                                            message = log_dict.get("message", "")

                                            # Filter by log level
                                            # warning/error level: only show ERROR and CRITICAL
                                            # info level: skip DEBUG
                                            # debug level: show all
                                            if log_level in ["warning", "error"] and levelname not in [
                                                "ERROR",
                                                "CRITICAL",
                                            ]:
                                                continue
                                            if log_level == "info" and levelname == "DEBUG":
                                                continue

                                            log_line = f"{levelname} | {ts} | {message}"
                                            if pod_name not in formatters:
                                                formatters[pod_name] = ServerLogsFormatter(pod_name)
                                            formatter = formatters[pod_name]
                                        else:
                                            # streaming pre server setup logs, before we have the pod name
                                            formatter = base_formatter
                                            message = log_line  # Use raw log line for dedup

                                        # Check for duplicates if dedup is enabled
                                        if seen_log_messages is not None:
                                            if message in seen_log_messages:
                                                continue
                                            seen_log_messages.add(message)

                                        # Add service name prefix if configured
                                        prefix = f"({self.service_name}) " if self.logging_config.include_name else ""
                                        formatted_line = (
                                            f"{formatter.start_color}{prefix}{log_line}{formatter.reset_color}"
                                        )

                                        print(formatted_line, flush=True)
                    except asyncio.TimeoutError:
                        # Timeout is expected, just continue the loop
                        continue
                    except websockets.exceptions.ConnectionClosed as e:
                        logger.debug(f"WebSocket connection closed: {str(e)}")
                        break
            finally:
                if websocket:
                    try:
                        # Use wait_for to prevent hanging on close
                        await asyncio.wait_for(websocket.close(), timeout=1.0)
                    except (asyncio.TimeoutError, Exception):
                        pass
        except Exception as e:
            logger.error(f"Error in websocket stream: {e}")
            raise e
        finally:
            # Ensure websocket is closed even if we didn't enter the try block
            if websocket:
                try:
                    # Use wait_for to prevent hanging on close
                    await asyncio.wait_for(websocket.close(), timeout=1.0)
                except (asyncio.TimeoutError, Exception):
                    pass

    def _wait_for_http_health(self, timeout=60, retry_interval=0.2, backoff=1.5, max_interval=2):
        """Wait for the HTTP server to be ready by checking the /health and /ready endpoints.

        Args:
            timeout: Maximum time to wait in seconds
            retry_interval: Time between health check attempts in seconds
        """
        import time

        logger.info(f"Polling {self.service_name} service health endpoint")
        start_time = time.time()
        health_ok = False

        while time.time() - start_time < timeout:
            try:
                client = self._client()

                # First check /health (server is up)
                if not health_ok:
                    response = client.get(
                        endpoint=f"{self.base_endpoint}/health",
                        headers=self.request_headers,
                        timeout=5,
                    )
                    if response.status_code == 200:
                        health_ok = True
                        logger.debug(f"Service {self.service_name} health check passed, checking readiness...")
                    else:
                        logger.debug(f"Health check returned status {response.status_code}, retrying...")
                        time.sleep(retry_interval)
                        retry_interval = min(retry_interval * backoff, max_interval)
                        continue

                # Then check /ready (callable is loaded)
                response = client.get(
                    endpoint=f"{self.base_endpoint}/ready",
                    headers=self.request_headers,
                    timeout=5,
                )
                if response.status_code == 200:
                    logger.info(f"Service {self.service_name} ready")
                    return
                else:
                    logger.debug(f"Readiness check returned status {response.status_code}, retrying...")

            except VersionMismatchError as e:
                raise e

            except Exception as e:
                # Don't log 502/503 errors - they're expected during startup
                if "502" not in str(e) and "503" not in str(e):
                    logger.debug(f"Health/readiness check failed: {e}, retrying...")

            time.sleep(retry_interval)
            retry_interval = min(retry_interval * backoff, max_interval)

        # If we get here, we've timed out
        raise ServiceTimeoutError(f"Service {self.service_name} not ready after {timeout}s")

    async def _wait_for_http_health_async(self, timeout=60, retry_interval=0.2, backoff=1.5, max_interval=2):
        """Async version of _wait_for_http_health. Wait for the HTTP server to be ready by checking the /health and /ready endpoints.

        Args:
            timeout: Maximum time to wait in seconds
            retry_interval: Time between health check attempts in seconds
        """
        import asyncio

        logger.debug(f"Waiting for HTTP server to be ready for service {self.service_name}")
        start_time = time.time()
        health_ok = False

        while time.time() - start_time < timeout:
            try:
                client = self._client()

                # First check /health (server is up)
                if not health_ok:
                    response = client.get(
                        endpoint=f"{self.base_endpoint}/health",
                        headers=self.request_headers,
                        timeout=5,
                    )
                    if response.status_code == 200:
                        health_ok = True
                        logger.debug(f"Service {self.service_name} health check passed, checking readiness...")
                    else:
                        logger.debug(f"Health check returned status {response.status_code}, retrying...")
                        await asyncio.sleep(retry_interval)
                        retry_interval = min(retry_interval * backoff, max_interval)
                        continue

                # Then check /ready (callable is loaded)
                response = client.get(
                    endpoint=f"{self.base_endpoint}/ready",
                    headers=self.request_headers,
                    timeout=5,
                )
                if response.status_code == 200:
                    logger.info(f"Service {self.service_name} ready")
                    return
                else:
                    logger.debug(f"Readiness check returned status {response.status_code}, retrying...")

            except Exception as e:
                # Don't log 502/503 errors - they're expected during startup
                if "502" not in str(e) and "503" not in str(e):
                    logger.debug(f"Health/readiness check failed: {e}, retrying...")

            await asyncio.sleep(retry_interval)
            retry_interval = min(retry_interval * backoff, max_interval)

        # If we get here, we've timed out
        raise ServiceTimeoutError(f"Service {self.service_name} not ready after {timeout}s")

    def __getstate__(self):
        """Remove local stateful values before pickle serialization."""
        state = self.__dict__.copy()
        # Remove local stateful values that shouldn't be serialized
        state["_http_client"] = None
        state["_service_config"] = None
        state["_remote_pointers"] = None
        # Pointers need to be converted to not be absolute paths if we're passing
        # the service elsewhere, e.g. into another service
        state["pointers"] = self.remote_pointers
        return state

    def __setstate__(self, state):
        """Restore state after pickle deserialization."""
        self.__dict__.update(state)
        # Reset local stateful values to None to ensure clean initialization
        self._http_client = None
        self._service_config = None
        self._remote_pointers = None

    def __del__(self):
        if hasattr(self, "_http_client") and self._http_client is not None:
            try:
                self._http_client.close()
            except Exception as e:
                logger.debug(f"Error closing HTTPClient in Module deletion: {e}")
            finally:
                self._http_client = None
