import asyncio
import atexit
import io
import os
import shlex
import subprocess
import sys
import tarfile
from pathlib import Path
from typing import Optional

from kubernetes import client as k8s_client
from kubernetes import config as k8s_config
from kubernetes.client.rest import ApiException
from kubernetes.stream import stream
from tenacity import retry, stop_after_attempt, wait_exponential

from harbor.environments.base import BaseEnvironment, ExecResult
from harbor.models.environment_type import EnvironmentType
from harbor.models.task.config import EnvironmentConfig
from harbor.models.trial.paths import EnvironmentPaths, TrialPaths
from harbor.utils.logger import logger


class KubernetesClientManager:
    """
    Singleton manager for the Kubernetes client.

    Ensures a single shared client instance across all GKEEnvironment instances,
    with proper cleanup at program termination.
    """

    _instance: "KubernetesClientManager | None" = None
    _lock = asyncio.Lock()

    def __init__(self):
        self._core_api = None
        self._reference_count = 0
        self._client_lock = asyncio.Lock()
        self._initialized = False
        self._cleanup_registered = False
        self._logger = logger.getChild(__name__)
        # Store cluster config to validate consistency across calls
        self._cluster_name: str | None = None
        self._region: str | None = None
        self._project_id: str | None = None

    @classmethod
    async def get_instance(cls) -> "KubernetesClientManager":
        """Get or create the singleton instance."""
        if cls._instance is None:
            async with cls._lock:
                if cls._instance is None:
                    cls._instance = cls()

        assert cls._instance is not None

        return cls._instance

    def _init_client(self, cluster_name: str, region: str, project_id: str):
        """Initialize Kubernetes client with GKE credentials."""
        if self._initialized:
            return

        try:
            k8s_config.load_kube_config()
            self._core_api = k8s_client.CoreV1Api()
            self._initialized = True
        except k8s_config.ConfigException:
            get_creds_cmd = [
                "gcloud",
                "container",
                "clusters",
                "get-credentials",
                cluster_name,
                "--region",
                region,
                "--project",
                project_id,
            ]

            result = subprocess.run(get_creds_cmd, capture_output=True, text=True)
            if result.returncode != 0:
                raise RuntimeError(
                    f"Failed to get GKE credentials: {result.stderr}\n"
                    f"Ensure cluster {cluster_name} exists in {region}"
                )

            k8s_config.load_kube_config()
            self._core_api = k8s_client.CoreV1Api()
            self._initialized = True

        # Store cluster config for validation
        self._cluster_name = cluster_name
        self._region = region
        self._project_id = project_id

    async def get_client(self, cluster_name: str, region: str, project_id: str):
        """
        Get the shared Kubernetes CoreV1Api client, creating it if necessary.
        Also increments the reference count.

        Note: This manager assumes all GKEEnvironment instances in a process
        connect to the same cluster. If a different cluster is requested after
        initialization, a ValueError is raised.
        """
        async with self._client_lock:
            if not self._initialized:
                self._logger.debug("Creating new Kubernetes client")
                await asyncio.to_thread(
                    self._init_client, cluster_name, region, project_id
                )

                if not self._cleanup_registered:
                    atexit.register(self._cleanup_sync)
                    self._cleanup_registered = True
            else:
                # Validate cluster config matches
                if (
                    self._cluster_name != cluster_name
                    or self._region != region
                    or self._project_id != project_id
                ):
                    raise ValueError(
                        f"KubernetesClientManager already initialized for cluster "
                        f"'{self._cluster_name}' in {self._region} (project: {self._project_id}). "
                        f"Cannot connect to cluster '{cluster_name}' in {region} "
                        f"(project: {project_id}). Use separate processes for different clusters."
                    )

            self._reference_count += 1
            self._logger.debug(
                f"Kubernetes client reference count incremented to {self._reference_count}"
            )
            return self._core_api

    async def release_client(self):
        """
        Decrement the reference count for the client.
        Note: Actual cleanup happens at program exit via atexit.
        """
        async with self._client_lock:
            if self._reference_count > 0:
                self._reference_count -= 1
                self._logger.debug(
                    f"Kubernetes client reference count decremented to {self._reference_count}"
                )

    def _cleanup_sync(self):
        """Synchronous cleanup wrapper for atexit."""
        try:
            asyncio.run(self._cleanup())
        except Exception as e:
            print(f"Error during Kubernetes client cleanup: {e}", file=sys.stderr)

    async def _cleanup(self):
        """Clean up the Kubernetes client if it exists."""
        async with self._client_lock:
            if self._initialized:
                try:
                    self._logger.debug("Cleaning up Kubernetes client at program exit")
                    self._core_api = None
                    self._initialized = False
                    self._logger.debug("Kubernetes client cleaned up successfully")
                except Exception as e:
                    self._logger.error(f"Error cleaning up Kubernetes client: {e}")


class GKEEnvironment(BaseEnvironment):
    """
    GKE implementation for Harbor sandboxes.

    Supports both Standard GKE and Autopilot clusters.
    """

    def __init__(
        self,
        environment_dir: Path,
        environment_name: str,
        session_id: str,
        trial_paths: TrialPaths,
        task_env_config: EnvironmentConfig,
        cluster_name: str,
        region: str,
        namespace: str,
        registry_location: str,
        registry_name: str,
        project_id: Optional[str] = None,
        memory_limit_multiplier: float | None = None,
        cloud_build_machine_type: str = "E2_HIGHCPU_8",
        cloud_build_disk_size_gb: int = 200,
        **kwargs,
    ):
        """
        Initialize GKE environment.

        Args:
            environment_dir: Path to the environment directory containing Dockerfile
            environment_name: Name of the environment (e.g., sb__hello-world)
            session_id: Session ID for this trial
            trial_paths: Trial paths for logs and output
            task_env_config: Task environment configuration (includes cpus, memory_mb, storage_mb)
            project_id: GCP project ID (defaults to env var or gcloud config)
            cluster_name: GKE cluster name
            region: GCP region
            namespace: Kubernetes namespace
            registry_location: Artifact Registry location (e.g., us-central1)
            registry_name: Artifact Registry repository name
            memory_limit_multiplier: Optional multiplier for memory limits.
                None or 0: no limits (Burstable QoS, can use all node memory)
                1: limits == requests (Guaranteed QoS, strict enforcement)
                >1: limits = multiplier * requests (Burstable with ceiling)
            cloud_build_machine_type: Cloud Build machine type (default: E2_HIGHCPU_8)
            cloud_build_disk_size_gb: Cloud Build disk size in GB (default: 200)
        """
        super().__init__(
            environment_dir=environment_dir,
            environment_name=environment_name,
            session_id=session_id,
            trial_paths=trial_paths,
            task_env_config=task_env_config,
            **kwargs,
        )

        # GKE configuration
        self.project_id = project_id or self._get_default_project()
        self.cluster_name = cluster_name
        self.region = region
        self.namespace = namespace

        # Resource configuration from task_env_config
        self.cpu_request = str(task_env_config.cpus)
        # Use Mi directly to avoid precision loss from integer division
        self.memory_request = f"{task_env_config.memory_mb}Mi"
        # Use Mi for ephemeral storage as well
        self.ephemeral_storage_request = f"{task_env_config.storage_mb}Mi"

        # Optional memory limit control
        if memory_limit_multiplier is not None and memory_limit_multiplier > 0:
            limit_memory_mb = int(task_env_config.memory_mb * memory_limit_multiplier)
            self.memory_limit = f"{limit_memory_mb}Mi"
        else:
            self.memory_limit = None

        # Cloud Build configuration
        self.cloud_build_machine_type = cloud_build_machine_type
        self.cloud_build_disk_size_gb = cloud_build_disk_size_gb

        self.registry_location = registry_location
        self.registry_name = registry_name

        # Pod naming - GKE compatible (max 63 chars, lowercase, alphanumeric + hyphens)
        self.pod_name = f"{session_id.lower().replace('_', '-')}"[:63]

        # Client manager for shared Kubernetes client
        self._client_manager: KubernetesClientManager | None = None
        self._core_api: k8s_client.CoreV1Api | None = None

    def _get_default_project(self) -> str:
        """Get default GCP project from environment or gcloud config."""
        # Check environment variable first
        project = os.environ.get("GCP_PROJECT") or os.environ.get(
            "GOOGLE_CLOUD_PROJECT"
        )
        if project:
            return project

        # Fall back to gcloud config
        try:
            result = subprocess.run(
                ["gcloud", "config", "get-value", "project"],
                capture_output=True,
                text=True,
                check=True,
            )
            return result.stdout.strip()
        except (subprocess.CalledProcessError, FileNotFoundError):
            raise ValueError(
                "No GCP project specified. Set project_id parameter, "
                "GCP_PROJECT environment variable, or configure gcloud default project."
            )

    @property
    def _api(self) -> k8s_client.CoreV1Api:
        """Return the Kubernetes API client, raising if not initialized."""
        if self._core_api is None:
            raise RuntimeError(
                "Kubernetes client not initialized. Call _ensure_client() first."
            )
        return self._core_api

    async def _ensure_client(self):
        """Ensure Kubernetes client is initialized via the singleton manager."""
        if self._client_manager is None:
            self._client_manager = await KubernetesClientManager.get_instance()
        if self._core_api is None:
            self._core_api = await self._client_manager.get_client(
                self.cluster_name, self.region, self.project_id
            )

    @staticmethod
    def type() -> EnvironmentType:
        return EnvironmentType.GKE

    @property
    def is_mounted(self) -> bool:
        """Cloud environments don't mount directories."""
        return False

    @property
    def supports_gpus(self) -> bool:
        return False

    @property
    def can_disable_internet(self) -> bool:
        return False

    @property
    def _environment_definition_path(self) -> Path:
        return self.environment_dir / "Dockerfile"

    def _validate_definition(self):
        if not self._environment_definition_path.exists():
            raise FileNotFoundError(
                f"{self._environment_definition_path} not found. Please ensure the "
                "file exists."
            )

    def _get_image_url(self) -> str:
        """Get the container image URL in Artifact Registry."""
        return f"{self.registry_location}-docker.pkg.dev/{self.project_id}/{self.registry_name}/{self.environment_name}:latest"

    async def _image_exists(self) -> bool:
        """Check if image already exists in Artifact Registry."""
        image_name = self.environment_name

        check_cmd = [
            "gcloud",
            "artifacts",
            "docker",
            "images",
            "describe",
            f"{self.registry_location}-docker.pkg.dev/{self.project_id}/{self.registry_name}/{image_name}:latest",
            "--project",
            self.project_id,
        ]

        try:
            result = await asyncio.create_subprocess_exec(
                *check_cmd,
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
            )
            await result.wait()
            return result.returncode == 0
        except Exception as e:
            self.logger.warning(
                f"Failed to check for existing image, will attempt to build. Error: {e}"
            )
            return False

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=2, min=5, max=60),
        reraise=True,
    )
    async def _build_and_push_image(self):
        """Build and push image using Cloud Build with retries."""
        image_url = self._get_image_url()
        self.logger.debug(f"Building and pushing image: {image_url}")

        build_cmd = [
            "gcloud",
            "builds",
            "submit",
            "--tag",
            image_url,
            "--project",
            self.project_id,
            "--region",
            self.region,
            "--timeout",
            "2400",  # 40 minutes for large images
            "--machine-type",
            self.cloud_build_machine_type,
            "--disk-size",
            str(self.cloud_build_disk_size_gb),
            "--quiet",
            str(self.environment_dir),
        ]

        result = await asyncio.create_subprocess_exec(
            *build_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        stdout, stderr = await result.communicate()

        if result.returncode != 0:
            error_msg = stderr.decode()
            raise RuntimeError(
                f"Image build failed: {error_msg}\nStdout: {stdout.decode()}"
            )

        self.logger.debug(f"Successfully built and pushed: {image_url}")

    async def start(self, force_build: bool):
        """Start a pod in GKE."""
        # Initialize Kubernetes client via singleton manager
        await self._ensure_client()

        # Hybrid build approach: build only if needed
        if force_build:
            await self._build_and_push_image()
        else:
            if not await self._image_exists():
                self.logger.debug(
                    f"Image {self._get_image_url()} not found, building..."
                )
                await self._build_and_push_image()
            else:
                self.logger.debug(f"Using existing image: {self._get_image_url()}")

        # Build resource requests
        requests = {
            "cpu": self.cpu_request,
            "memory": self.memory_request,
        }
        if self.ephemeral_storage_request:
            requests["ephemeral-storage"] = self.ephemeral_storage_request

        # Build resource limits (optional)
        limits = {}
        if self.memory_limit:
            limits["memory"] = self.memory_limit

        # Create Pod specification
        pod = k8s_client.V1Pod(
            api_version="v1",
            kind="Pod",
            metadata=k8s_client.V1ObjectMeta(
                name=self.pod_name,
                namespace=self.namespace,
                labels={
                    "app": "sandbox",
                    "session": self.session_id,
                    "environment": self.environment_name,
                },
            ),
            spec=k8s_client.V1PodSpec(
                containers=[
                    k8s_client.V1Container(
                        name="main",
                        image=self._get_image_url(),
                        command=["sleep", "infinity"],
                        resources=k8s_client.V1ResourceRequirements(
                            requests=requests,
                            limits=limits or None,
                        ),
                        volume_mounts=[],
                    )
                ],
                restart_policy="Never",
            ),
        )

        # Create the pod
        try:
            await asyncio.to_thread(
                self._api.create_namespaced_pod,
                namespace=self.namespace,
                body=pod,
            )
        except ApiException as e:
            if e.status == 409:  # Already exists
                self.logger.debug(f"Pod {self.pod_name} already exists, recreating...")
                # Delete existing pod inline (don't call stop() as it releases the client)
                try:
                    await asyncio.to_thread(
                        self._api.delete_namespaced_pod,
                        name=self.pod_name,
                        namespace=self.namespace,
                        body=k8s_client.V1DeleteOptions(
                            grace_period_seconds=0, propagation_policy="Foreground"
                        ),
                    )
                    # Wait for deletion
                    for _ in range(60):
                        try:
                            await asyncio.to_thread(
                                self._api.read_namespaced_pod,
                                name=self.pod_name,
                                namespace=self.namespace,
                            )
                            await asyncio.sleep(1)
                        except ApiException as del_e:
                            if del_e.status == 404:
                                break
                    else:
                        raise RuntimeError(
                            f"Pod {self.pod_name} was not deleted in time."
                        )
                except ApiException as del_e:
                    if del_e.status != 404:
                        raise RuntimeError(f"Failed to delete existing pod: {del_e}")

                await asyncio.to_thread(
                    self._api.create_namespaced_pod,
                    namespace=self.namespace,
                    body=pod,
                )
            else:
                raise RuntimeError(f"Failed to create pod: {e}")

        # Wait for pod to be ready
        await self._wait_for_pod_ready()

        # Create required directories
        mkdir_result = await self.exec(
            f"mkdir -p {EnvironmentPaths.agent_dir} {EnvironmentPaths.verifier_dir}"
        )
        if mkdir_result.return_code != 0:
            raise RuntimeError(
                f"Failed to create log directories in pod {self.pod_name}: "
                f"stdout={mkdir_result.stdout}, stderr={mkdir_result.stderr}"
            )

    async def stop(self, delete: bool):
        """Stop/delete the pod."""
        if self._client_manager is None:
            return

        try:
            if delete:
                try:
                    await asyncio.to_thread(
                        self._api.delete_namespaced_pod,
                        name=self.pod_name,
                        namespace=self.namespace,
                        body=k8s_client.V1DeleteOptions(
                            grace_period_seconds=0,
                            propagation_policy="Foreground",
                        ),
                    )
                    # Wait for pod to be deleted
                    for _ in range(60):
                        try:
                            await asyncio.to_thread(
                                self._api.read_namespaced_pod,
                                name=self.pod_name,
                                namespace=self.namespace,
                            )
                            await asyncio.sleep(1)
                        except ApiException as e:
                            if e.status == 404:
                                break
                    else:
                        self.logger.warning(
                            f"Pod {self.pod_name} did not terminate within 60 seconds."
                        )
                except ApiException as e:
                    if e.status != 404:
                        raise
        finally:
            # Release the client reference (actual cleanup happens at program exit)
            if self._client_manager:
                try:
                    await self._client_manager.release_client()
                except Exception as e:
                    self.logger.error(f"Error releasing Kubernetes client: {e}")
                finally:
                    self._client_manager = None
                    self._core_api = None

    async def exec(
        self,
        command: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout_sec: int | None = None,
    ) -> ExecResult:
        """Execute command in pod using kubectl exec equivalent."""
        await self._ensure_client()

        full_command = f"bash -c {shlex.quote(command)}"

        if env:
            for key, value in env.items():
                full_command = f"{key}={shlex.quote(value)} {full_command}"

        if cwd:
            full_command = f"cd {cwd} && {full_command}"

        exec_command = ["sh", "-c", full_command]

        resp = None
        try:
            resp = await asyncio.to_thread(
                stream,
                self._api.connect_get_namespaced_pod_exec,
                self.pod_name,
                self.namespace,
                command=exec_command,
                stderr=True,
                stdin=False,
                stdout=True,
                tty=False,
                _preload_content=False,
            )

            if timeout_sec:
                stdout, stderr = await asyncio.wait_for(
                    asyncio.to_thread(self._read_exec_output, resp),
                    timeout=timeout_sec,
                )
            else:
                stdout, stderr = await asyncio.to_thread(self._read_exec_output, resp)

            resp.run_forever(timeout=0)
            return_code = resp.returncode if resp.returncode is not None else 0

            return ExecResult(
                stdout=stdout,
                stderr=stderr,
                return_code=return_code,
            )

        except asyncio.TimeoutError:
            return ExecResult(
                stdout=None,
                stderr=f"Command timed out after {timeout_sec} seconds",
                return_code=124,
            )
        except ApiException as e:
            if e.status == 404:
                return ExecResult(
                    stdout=None,
                    stderr=f"Pod {self.pod_name} not found (404).",
                    return_code=1,
                )
            elif e.status == 500:
                error_body = str(e.body) if hasattr(e, "body") else str(e)
                if "No agent available" in error_body:
                    return ExecResult(
                        stdout=None,
                        stderr=f"Pod {self.pod_name} unavailable: No agent available.",
                        return_code=1,
                    )
                return ExecResult(
                    stdout=None,
                    stderr=f"Internal server error on pod {self.pod_name}: {e.reason}",
                    return_code=1,
                )
            else:
                return ExecResult(
                    stdout=None,
                    stderr=f"API error ({e.status}) on pod {self.pod_name}: {e.reason}",
                    return_code=1,
                )
        except Exception as e:
            return ExecResult(
                stdout=None,
                stderr=str(e),
                return_code=1,
            )
        finally:
            if resp is not None:
                try:
                    resp.close()
                except Exception:
                    pass

    def _read_exec_output(self, resp):
        """Read output from exec stream."""
        stdout = ""
        stderr = ""

        while resp.is_open():
            resp.update(timeout=1)
            if resp.peek_stdout():
                stdout += resp.read_stdout()
            if resp.peek_stderr():
                stderr += resp.read_stderr()

        return stdout, stderr

    async def _wait_for_container_exec_ready(self, max_attempts: int = 60):
        """Wait for container to be ready for exec operations."""
        for attempt in range(max_attempts):
            try:
                test_command = ["true"]
                resp = await asyncio.to_thread(
                    stream,
                    self._api.connect_get_namespaced_pod_exec,
                    self.pod_name,
                    self.namespace,
                    command=test_command,
                    stderr=False,
                    stdin=False,
                    stdout=True,
                    tty=False,
                    _preload_content=False,
                )
                resp.close()
                return
            except ApiException as e:
                if "container not found" in str(e) or e.status == 500:
                    if attempt % 10 == 0:
                        self.logger.debug(
                            f"Container not ready, attempt {attempt + 1}/{max_attempts}"
                        )
                    await asyncio.sleep(3)
                    continue
                else:
                    raise
            except Exception as e:
                if attempt < max_attempts - 1:
                    if attempt % 10 == 0:
                        self.logger.debug(f"Error checking container readiness: {e}")
                    await asyncio.sleep(3)
                    continue
                else:
                    raise

        raise RuntimeError(
            f"Container not ready for exec after {max_attempts} attempts"
        )

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def upload_file(self, source_path: Path | str, target_path: str):
        """Upload file using kubectl cp equivalent."""
        await self._ensure_client()

        await self._wait_for_container_exec_ready()

        source_path = Path(source_path)

        tar_buffer = io.BytesIO()
        with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
            tar.add(str(source_path), arcname=Path(target_path).name)
        tar_buffer.seek(0)

        target_dir = str(Path(target_path).parent)
        await self.exec(f"mkdir -p {target_dir}")

        exec_command = ["tar", "xf", "-", "-C", target_dir]

        resp = await asyncio.to_thread(
            stream,
            self._api.connect_get_namespaced_pod_exec,
            self.pod_name,
            self.namespace,
            command=exec_command,
            stderr=True,
            stdin=True,
            stdout=True,
            tty=False,
            _preload_content=False,
        )

        resp.write_stdin(tar_buffer.read())
        resp.run_forever(timeout=1)
        resp.close()

    @retry(
        stop=stop_after_attempt(5),
        wait=wait_exponential(multiplier=1, min=2, max=30),
        reraise=True,
    )
    async def upload_dir(self, source_dir: Path | str, target_dir: str):
        """Upload directory using kubectl cp equivalent."""
        await self._ensure_client()

        await self._wait_for_container_exec_ready()

        source_dir = Path(source_dir)

        files_to_upload = []
        for item in source_dir.rglob("*"):
            if item.is_file():
                arcname = str(item.relative_to(source_dir))
                files_to_upload.append(arcname)

        if not files_to_upload:
            self.logger.warning(f"No files to upload from {source_dir}")
            return

        tar_buffer = io.BytesIO()
        with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
            for item in source_dir.rglob("*"):
                if item.is_file():
                    arcname = str(item.relative_to(source_dir))
                    tar.add(str(item), arcname=arcname)
        tar_buffer.seek(0)
        tar_size = len(tar_buffer.getvalue())

        mkdir_result = await self.exec(f"mkdir -p {target_dir}")
        if mkdir_result.return_code != 0:
            raise RuntimeError(
                f"Failed to create target directory {target_dir}: {mkdir_result.stderr}"
            )

        exec_command = ["tar", "xf", "-", "-C", target_dir]

        try:
            resp = await asyncio.to_thread(
                stream,
                self._api.connect_get_namespaced_pod_exec,
                self.pod_name,
                self.namespace,
                command=exec_command,
                stderr=True,
                stdin=True,
                stdout=True,
                tty=False,
                _preload_content=False,
            )
        except ApiException as e:
            if e.status == 500:
                raise RuntimeError(
                    f"Pod {self.pod_name} returned 500 error during upload."
                )
            raise

        try:
            resp.write_stdin(tar_buffer.read())
        except Exception as e:
            raise RuntimeError(f"Failed to write tar data to pod {self.pod_name}: {e}")

        resp.run_forever(timeout=1)
        resp.close()
        self.logger.debug(
            f"Successfully uploaded {len(files_to_upload)} files ({tar_size} bytes) to {target_dir}"
        )

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        reraise=True,
    )
    async def download_file(self, source_path: str, target_path: Path | str):
        """Download file from pod."""
        await self._ensure_client()

        target_path = Path(target_path)
        target_path.parent.mkdir(parents=True, exist_ok=True)

        exec_command = ["tar", "cf", "-", source_path]

        resp = await asyncio.to_thread(
            stream,
            self._api.connect_get_namespaced_pod_exec,
            self.pod_name,
            self.namespace,
            command=exec_command,
            stderr=True,
            stdin=False,
            stdout=True,
            tty=False,
            _preload_content=False,
        )

        tar_data = b""
        while resp.is_open():
            resp.update(timeout=1)
            if resp.peek_stdout():
                data = resp.read_stdout()
                if isinstance(data, str):
                    data = data.encode("utf-8", errors="surrogateescape")
                tar_data += data

        tar_buffer = io.BytesIO(tar_data)
        with tarfile.open(fileobj=tar_buffer, mode="r") as tar:
            for member in tar.getmembers():
                if member.name == source_path or member.name.startswith(
                    source_path.lstrip("/")
                ):
                    member.name = target_path.name
                    tar.extract(member, path=str(target_path.parent))
                    break

    @retry(
        stop=stop_after_attempt(5),
        wait=wait_exponential(multiplier=1, min=2, max=30),
        reraise=True,
    )
    async def download_dir(self, source_dir: str, target_dir: Path | str):
        """Download directory from pod."""
        await self._ensure_client()

        target_dir = Path(target_dir)
        target_dir.mkdir(parents=True, exist_ok=True)

        exec_command = ["sh", "-c", f"cd {source_dir} && tar cf - ."]

        try:
            resp = await asyncio.to_thread(
                stream,
                self._api.connect_get_namespaced_pod_exec,
                self.pod_name,
                self.namespace,
                command=exec_command,
                stderr=True,
                stdin=False,
                stdout=True,
                tty=False,
                _preload_content=False,
            )
        except ApiException as e:
            if e.status == 404:
                raise RuntimeError(f"Pod {self.pod_name} not found (404).")
            elif e.status == 500:
                raise RuntimeError(f"Pod {self.pod_name} is in an error state (500).")
            raise

        tar_data = b""
        stderr_data = ""
        while resp.is_open():
            resp.update(timeout=1)
            if resp.peek_stdout():
                data = resp.read_stdout()
                if isinstance(data, str):
                    data = data.encode("utf-8", errors="surrogateescape")
                tar_data += data
            if resp.peek_stderr():
                stderr_data += resp.read_stderr()

        if stderr_data and (
            "No such file or directory" in stderr_data or "cannot cd" in stderr_data
        ):
            raise RuntimeError(
                f"Failed to access directory {source_dir} in pod {self.pod_name}: {stderr_data.strip()}"
            )

        if not tar_data:
            raise RuntimeError(
                f"No data received when downloading {source_dir} from pod {self.pod_name}."
            )

        tar_buffer = io.BytesIO(tar_data)
        try:
            with tarfile.open(fileobj=tar_buffer, mode="r") as tar:
                tar.extractall(path=str(target_dir))
        except tarfile.TarError as e:
            raise RuntimeError(
                f"Failed to extract directory {source_dir} from pod {self.pod_name}: {e}"
            )

    async def _wait_for_pod_ready(self, timeout_sec: int = 300):
        """Wait for pod to be ready."""
        self.logger.debug(f"Waiting for pod {self.pod_name} to be ready...")

        for attempt in range(timeout_sec):
            try:
                pod = await asyncio.to_thread(
                    self._api.read_namespaced_pod,
                    name=self.pod_name,
                    namespace=self.namespace,
                )

                if pod.status.phase == "Running":
                    if pod.status.container_statuses:
                        if all(c.ready for c in pod.status.container_statuses):
                            self.logger.debug(f"Pod {self.pod_name} is ready!")
                            return

                elif pod.status.phase in ["Failed", "Unknown", "Error"]:
                    error_details = self._get_pod_failure_summary(pod)
                    raise RuntimeError(f"Pod failed to start: {error_details}")

                elif pod.status.phase == "Pending":
                    # Check for image pull errors
                    if pod.status.container_statuses:
                        for c in pod.status.container_statuses:
                            if c.state.waiting:
                                if (
                                    "ImagePullBackOff" in c.state.waiting.reason
                                    or "ErrImagePull" in c.state.waiting.reason
                                ):
                                    raise RuntimeError(
                                        f"Failed to pull image: {c.state.waiting.message or c.state.waiting.reason}"
                                    )

                if attempt % 10 == 0:
                    self.logger.debug(
                        f"Pod status: {pod.status.phase} ({attempt}s elapsed)"
                    )

            except ApiException as e:
                if e.status != 404:
                    raise RuntimeError(f"Kubernetes API error: {e.status} - {e.reason}")

            await asyncio.sleep(1)

        raise RuntimeError(f"Pod not ready after {timeout_sec} seconds")

    def _get_pod_failure_summary(self, pod) -> str:
        """Get a summary of pod failure reasons."""
        reasons = []

        if pod.status.reason:
            reasons.append(f"Reason: {pod.status.reason}")
        if pod.status.message:
            reasons.append(f"Message: {pod.status.message}")

        if pod.status.container_statuses:
            for c in pod.status.container_statuses:
                if c.state.waiting:
                    reasons.append(
                        f"Container {c.name} waiting: {c.state.waiting.reason}"
                    )
                elif c.state.terminated:
                    reasons.append(
                        f"Container {c.name} terminated: {c.state.terminated.reason} "
                        f"(exit code {c.state.terminated.exit_code})"
                    )

        return "; ".join(reasons) if reasons else "Unknown error"
