import grpc
import logging
import time
import json
import urllib.request
import urllib.parse
import urllib.error
import random
import threading
import os
import signal
import queue
import psutil
import importlib
import inspect
import functools
import typing
import socket
import ipaddress
from typing import Any, Callable, Dict, Optional, TypeVar, Iterator, List, Tuple
from types import ModuleType
import hashlib
import msgspec
try:
    import torch
except Exception:  # pragma: no cover - optional dependency
    torch = None
import asyncio

# JWT verification for worker auth (scheduler-issued)
import jwt
_jwt_algorithms: Optional[ModuleType]
try:
    import jwt.algorithms as _jwt_algorithms
except Exception:  # pragma: no cover - optional crypto backend
    _jwt_algorithms = None
RSAAlgorithm: Optional[Any] = getattr(_jwt_algorithms, "RSAAlgorithm", None) if _jwt_algorithms else None
# Use relative imports within the package
from .pb import worker_scheduler_pb2 as _pb
from .pb import worker_scheduler_pb2_grpc as _pb_grpc

pb: Any = _pb
pb_grpc: Any = _pb_grpc

WorkerSchedulerMessage = Any
WorkerEvent = Any
WorkerResources = Any
WorkerRegistration = Any
LoadModelCommand = Any
LoadModelResult = Any
UnloadModelResult = Any
TaskExecutionRequest = Any
TaskExecutionResult = Any
from .decorators import ResourceRequirements # Import ResourceRequirements for type hints if needed
from .errors import RetryableError, FatalError

from .model_interface import ModelManagementInterface
from .downloader import CozyHubDownloader, ModelDownloader
from .types import Asset

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__) # Use __name__ for logger

# Type variables for generic function signatures
I = TypeVar('I')  # Input type
O = TypeVar('O')  # Output type

# Generic type for action functions
ActionFunc = Callable[[Any, I], O]

HEARTBEAT_INTERVAL = 10  # seconds


def _encode_ref_for_url(ref: str) -> str:
    ref = ref.strip().lstrip("/")
    parts = [urllib.parse.quote(p, safe="") for p in ref.split("/") if p]
    return "/".join(parts)


def _infer_mime_type(ref: str, head: bytes) -> str:
    # Prefer magic bytes when available.
    if head.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if head.startswith(b"\xff\xd8\xff"):
        return "image/jpeg"
    if head.startswith(b"GIF87a") or head.startswith(b"GIF89a"):
        return "image/gif"
    if len(head) >= 12 and head[0:4] == b"RIFF" and head[8:12] == b"WEBP":
        return "image/webp"

    # Fall back to extension.
    import mimetypes

    guessed, _ = mimetypes.guess_type(ref)
    return guessed or "application/octet-stream"


def _default_output_prefix(run_id: str) -> str:
    return f"runs/{run_id}/outputs/"


def _require_file_api_base_url() -> str:
    base = os.getenv("FILE_API_BASE_URL", "").strip()
    if not base:
        base = os.getenv("ORCHESTRATOR_HTTP_URL", "").strip()
    if not base:
        base = os.getenv("COZY_HUB_URL", "").strip()
    if not base:
        raise RuntimeError("FILE_API_BASE_URL is required for file operations")
    return base.rstrip("/")


def _require_file_api_token() -> str:
    token = os.getenv("FILE_API_TOKEN", "").strip()
    if not token:
        token = os.getenv("COZY_HUB_TOKEN", "").strip()
    if not token:
        raise RuntimeError("FILE_API_TOKEN is required for file operations")
    return token


def _http_request(
    method: str,
    url: str,
    token: str,
    body: Optional[bytes] = None,
    content_type: Optional[str] = None,
) -> urllib.request.Request:
    req = urllib.request.Request(url, data=body, method=method)
    req.add_header("Authorization", f"Bearer {token}")
    tenant_id = os.getenv("TENANT_ID", "").strip()
    if tenant_id:
        req.add_header("X-Cozy-Tenant-Id", tenant_id)
    if content_type:
        req.add_header("Content-Type", content_type)
    return req


def _is_private_ip_str(ip_str: str) -> bool:
    try:
        ip = ipaddress.ip_address(ip_str)
    except Exception:
        return True
    return bool(
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_multicast
        or ip.is_reserved
        or ip.is_unspecified
    )


def _url_is_blocked(url_str: str) -> bool:
    try:
        u = urllib.parse.urlparse(url_str)
    except Exception:
        return True
    if u.scheme not in ("http", "https"):
        return True
    host = (u.hostname or "").strip()
    if not host:
        return True
    try:
        infos = socket.getaddrinfo(host, None)
    except Exception:
        return True
    for info in infos:
        sockaddr = info[4]
        if not sockaddr:
            continue
        ip_str = str(sockaddr[0])
        if _is_private_ip_str(ip_str):
            return True
    return False

class _JWKSCache:
    def __init__(self, url: str, ttl_seconds: int = 300) -> None:
        self._url = url
        self._ttl_seconds = max(ttl_seconds, 0)
        self._lock = threading.Lock()
        self._fetched_at = 0.0
        self._keys: Dict[str, Any] = {}

    def _fetch(self) -> None:
        if RSAAlgorithm is None:
            raise RuntimeError(
                "PyJWT RSA support is unavailable (missing cryptography). "
                "Install gen-worker with a JWT/RSA-capable build of PyJWT."
            )
        with urllib.request.urlopen(self._url, timeout=5) as resp:
            body = resp.read()
        payload = json.loads(body.decode("utf-8"))
        keys: Dict[str, Any] = {}
        for jwk in payload.get("keys", []):
            kid = jwk.get("kid")
            if not kid:
                continue
            try:
                keys[kid] = RSAAlgorithm.from_jwk(json.dumps(jwk))
            except Exception:
                continue
        self._keys = keys
        self._fetched_at = time.time()

    def _needs_refresh(self) -> bool:
        if not self._keys:
            return True
        if self._ttl_seconds <= 0:
            return False
        return (time.time() - self._fetched_at) > self._ttl_seconds

    def get_key(self, kid: Optional[str]) -> Optional[Any]:
        with self._lock:
            if self._needs_refresh():
                self._fetch()
            if kid and kid in self._keys:
                return self._keys[kid]
            # refresh on miss (rotation)
            self._fetch()
            if kid and kid in self._keys:
                return self._keys[kid]
            return None

class ActionContext:
    """Context object passed to action functions, allowing cancellation."""
    def __init__(
        self,
        run_id: str,
        emitter: Optional[Callable[[Dict[str, Any]], None]] = None,
        tenant_id: Optional[str] = None,
        user_id: Optional[str] = None,
        timeout_ms: Optional[int] = None,
    ) -> None:
        self._run_id = run_id
        self._tenant_id = tenant_id
        self._user_id = user_id
        self._timeout_ms = timeout_ms
        self._started_at = time.time()
        self._deadline: Optional[float] = None
        if timeout_ms is not None and timeout_ms > 0:
            self._deadline = self._started_at + (timeout_ms / 1000.0)
        self._canceled = False
        self._cancel_event = threading.Event()
        self._emitter = emitter

    @property
    def run_id(self) -> str:
        return self._run_id

    @property
    def tenant_id(self) -> Optional[str]:
        return self._tenant_id

    @property
    def user_id(self) -> Optional[str]:
        return self._user_id

    @property
    def timeout_ms(self) -> Optional[int]:
        return self._timeout_ms

    @property
    def deadline(self) -> Optional[float]:
        return self._deadline

    def time_remaining_s(self) -> Optional[float]:
        if self._deadline is None:
            return None
        return max(0.0, self._deadline - time.time())

    def is_canceled(self) -> bool:
        """Check if the action was canceled."""
        return self._canceled

    def cancel(self) -> None:
        """Mark the action as canceled."""
        if not self._canceled:
            self._canceled = True
            self._cancel_event.set()
            logger.info(f"Action {self.run_id} marked for cancellation.")

    def done(self) -> threading.Event:
        """Returns an event that is set when the action is cancelled."""
        return self._cancel_event

    def emit(self, event_type: str, payload: Optional[Dict[str, Any]] = None) -> None:
        """Emit a progress/event payload (best-effort)."""
        if not self._emitter:
            logger.debug(f"emit({event_type}) dropped: no emitter configured")
            return
        event = {
            "run_id": self._run_id,
            "type": event_type,
            "payload": payload or {},
            "timestamp": time.time(),
        }
        self._emitter(event)

    def progress(self, progress: float, stage: Optional[str] = None) -> None:
        payload: Dict[str, Any] = {"progress": progress}
        if stage is not None:
            payload["stage"] = stage
        self.emit("job.progress", payload)

    def log(self, message: str, level: str = "info") -> None:
        self.emit("job.log", {"message": message, "level": level})

    def save_bytes(self, ref: str, data: bytes) -> Asset:
        if not isinstance(data, (bytes, bytearray)):
            raise TypeError("save_bytes expects bytes")
        data = bytes(data)
        max_bytes = int(os.getenv("WORKER_MAX_OUTPUT_FILE_BYTES", str(200 * 1024 * 1024)))
        if max_bytes > 0 and len(data) > max_bytes:
            raise ValueError("output file too large")
        ref = ref.strip().lstrip("/")
        if not ref.startswith(_default_output_prefix(self.run_id)):
            raise ValueError(f"ref must start with '{_default_output_prefix(self.run_id)}'")

        base = _require_file_api_base_url()
        token = _require_file_api_token()
        url = f"{base}/api/v1/file/{_encode_ref_for_url(ref)}"
        # Default behavior is upsert: PUT to the tenant file store.
        req = _http_request("PUT", url, token, body=data, content_type="application/octet-stream")
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                body = resp.read()
                if resp.status < 200 or resp.status >= 300:
                    raise RuntimeError(f"file save failed ({resp.status})")
                try:
                    meta = json.loads(body.decode("utf-8"))
                except Exception:
                    meta = {}
        except urllib.error.HTTPError as e:
            raise RuntimeError(f"file save failed ({getattr(e, 'code', 'unknown')})") from e

        return Asset(
            ref=ref,
            tenant_id=self.tenant_id,
            local_path=None,
            mime_type=str(meta.get("mime_type") or "") or None,
            size_bytes=int(meta.get("size_bytes") or 0) or len(data),
            sha256=str(meta.get("sha256") or "") or None,
        )

    def save_file(self, ref: str, local_path: str) -> Asset:
        with open(local_path, "rb") as f:
            data = f.read()
        return self.save_bytes(ref, data)

    def save_bytes_create(self, ref: str, data: bytes) -> Asset:
        if not isinstance(data, (bytes, bytearray)):
            raise TypeError("save_bytes_create expects bytes")
        data = bytes(data)
        max_bytes = int(os.getenv("WORKER_MAX_OUTPUT_FILE_BYTES", str(200 * 1024 * 1024)))
        if max_bytes > 0 and len(data) > max_bytes:
            raise ValueError("output file too large")
        ref = ref.strip().lstrip("/")
        if not ref.startswith(_default_output_prefix(self.run_id)):
            raise ValueError(f"ref must start with '{_default_output_prefix(self.run_id)}'")

        base = _require_file_api_base_url()
        token = _require_file_api_token()
        url = f"{base}/api/v1/file/{_encode_ref_for_url(ref)}"
        req = _http_request("POST", url, token, body=data, content_type="application/octet-stream")
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                body = resp.read()
                if resp.status < 200 or resp.status >= 300:
                    raise RuntimeError(f"file save failed ({resp.status})")
                try:
                    meta = json.loads(body.decode("utf-8"))
                except Exception:
                    meta = {}
        except urllib.error.HTTPError as e:
            if getattr(e, "code", None) == 409:
                raise RuntimeError("output path already exists") from e
            raise RuntimeError(f"file save failed ({getattr(e, 'code', 'unknown')})") from e

        return Asset(
            ref=ref,
            tenant_id=self.tenant_id,
            local_path=None,
            mime_type=str(meta.get("mime_type") or "") or None,
            size_bytes=int(meta.get("size_bytes") or 0) or len(data),
            sha256=str(meta.get("sha256") or "") or None,
        )

    def save_file_create(self, ref: str, local_path: str) -> Asset:
        with open(local_path, "rb") as f:
            data = f.read()
        return self.save_bytes_create(ref, data)

    def save_bytes_overwrite(self, ref: str, data: bytes) -> Asset:
        # Back-compat alias: overwrite is the default save_bytes behavior.
        return self.save_bytes(ref, data)

    def save_file_overwrite(self, ref: str, local_path: str) -> Asset:
        return self.save_file(ref, local_path)

# Define the interceptor class correctly
class _AuthInterceptor(grpc.StreamStreamClientInterceptor):
    def __init__(self, token: str) -> None:
        self._token = token

    def intercept_stream_stream(self, continuation: Any, client_call_details: Any, request_iterator: Any) -> Any:
        metadata = list(client_call_details.metadata or [])
        metadata.append(('authorization', f'Bearer {self._token}'))
        new_details = client_call_details._replace(metadata=metadata)
        return continuation(new_details, request_iterator)

class Worker:
    """Worker implementation that connects to the scheduler via gRPC."""

    def __init__(
        self,
        scheduler_addr: str = "localhost:8080",
        scheduler_addrs: Optional[List[str]] = None,
        user_module_names: List[str] = ["functions"], # Add new parameter for user modules
        worker_id: Optional[str] = None,
        auth_token: Optional[str] = None,
        use_tls: bool = False,
        reconnect_delay: int = 5,
        max_reconnect_attempts: int = 0,  # 0 means infinite retries
        model_manager: Optional[ModelManagementInterface] = None, # Optional model manager
        downloader: Optional[ModelDownloader] = None,  # Optional model downloader
    ) -> None:
        """Initialize a new worker.

        Args:
            scheduler_addr: Address of the scheduler service.
            scheduler_addrs: Optional list of seed scheduler addresses.
            user_module_names: List of Python module names containing user-defined @worker_function functions.
            worker_id: Unique ID for this worker (generated if not provided).
            auth_token: Optional authentication token.
            use_tls: Whether to use TLS for the connection.
            reconnect_delay: Seconds to wait between reconnection attempts.
            max_reconnect_attempts: Max reconnect attempts (0 = infinite).
            model_manager: Optional model manager.
            downloader: Optional model downloader.
        """
        self.scheduler_addr = scheduler_addr
        self.scheduler_addrs = self._normalize_scheduler_addrs(scheduler_addr, scheduler_addrs)
        self.user_module_names = user_module_names # Store module names
        self.worker_id = worker_id or f"py-worker-{os.getpid()}"
        self.auth_token = auth_token
        self.use_tls = use_tls
        self.reconnect_delay = reconnect_delay
        self.max_reconnect_attempts = max_reconnect_attempts
        self.max_input_bytes = int(os.getenv("WORKER_MAX_INPUT_BYTES", "0"))
        self.max_output_bytes = int(os.getenv("WORKER_MAX_OUTPUT_BYTES", "0"))

        self._jwks_url = os.getenv("SCHEDULER_JWKS_URL", "").strip()
        self._jwks_ttl_seconds = int(os.getenv("SCHEDULER_JWKS_TTL_SECONDS", "300"))
        self._jwt_issuer = os.getenv("SCHEDULER_JWT_ISSUER", "").strip()
        self._jwt_audience = os.getenv("SCHEDULER_JWT_AUDIENCE", "").strip()
        self._jwks_cache: Optional[_JWKSCache] = _JWKSCache(self._jwks_url, self._jwks_ttl_seconds) if self._jwks_url else None

        self.deployment_id = os.getenv("DEPLOYMENT_ID", "") # Read DEPLOYMENT_ID env var
        if not self.deployment_id:
            logger.warning("DEPLOYMENT_ID environment variable not set for this worker!")

        self.tenant_id = os.getenv("TENANT_ID", "")
        self.runpod_pod_id = os.getenv("RUNPOD_POD_ID", "") # Read injected pod ID
        if not self.runpod_pod_id:
            logger.warning("RUNPOD_POD_ID environment variable not set for this worker!")

        logger.info(f"RUNPOD_POD_ID: {self.runpod_pod_id}")

        self._actions: Dict[str, Callable[[ActionContext, Optional[Any], bytes], bytes]] = {}
        self._active_tasks: Dict[str, ActionContext] = {}
        self._active_tasks_lock = threading.Lock()
        self._active_function_counts: Dict[str, int] = {}
        self.max_concurrency = int(os.getenv("WORKER_MAX_CONCURRENCY", "0"))
        self._drain_timeout_seconds = int(os.getenv("WORKER_DRAIN_TIMEOUT_SECONDS", "0"))
        self._draining = False
        self._discovered_resources: Dict[str, ResourceRequirements] = {} # Store resources per function
        self._function_schemas: Dict[str, Tuple[bytes, bytes]] = {}  # func_name -> (input_schema_json, output_schema_json)

        self._gpu_busy_lock = threading.Lock()
        self._is_gpu_busy = False

        self._channel: Optional[Any] = None
        self._stub: Optional[Any] = None
        self._stream: Optional[Any] = None
        self._running = False
        self._stop_event = threading.Event()
        self._reconnect_count = 0
        self._outgoing_queue: queue.Queue[Any] = queue.Queue()
        self._leader_hint: Optional[str] = None

        self._receive_thread: Optional[threading.Thread] = None
        self._heartbeat_thread: Optional[threading.Thread] = None

        self._reconnect_delay_base = max(0, reconnect_delay)
        self._reconnect_delay_max = int(os.getenv("RECONNECT_MAX_DELAY", "60"))
        self._reconnect_jitter = float(os.getenv("RECONNECT_JITTER_SECONDS", "1.0"))

        resolved_model_manager = model_manager
        if resolved_model_manager is None:
            model_manager_path = os.getenv("MODEL_MANAGER_CLASS", "").strip()
            if model_manager_path:
                try:
                    module_path, _, class_name = model_manager_path.partition(":")
                    if not module_path or not class_name:
                        raise ValueError("MODEL_MANAGER_CLASS must be in module:Class format")
                    module = importlib.import_module(module_path)
                    manager_cls = getattr(module, class_name)
                    resolved_model_manager = manager_cls()
                    logger.info(f"Loaded ModelManager from MODEL_MANAGER_CLASS={model_manager_path}")
                except Exception as e:
                    logger.exception(f"Failed to load MODEL_MANAGER_CLASS '{model_manager_path}': {e}")
        self._model_manager = resolved_model_manager
        self._downloader = downloader
        if self._downloader is None:
            base_url = os.getenv("COZY_HUB_URL", "").strip()
            token = os.getenv("COZY_HUB_TOKEN", "").strip() or None
            if base_url:
                self._downloader = CozyHubDownloader(base_url, token=token)
        self._supported_model_ids_from_scheduler: Optional[List[str]] = None # To store IDs from scheduler
        self._model_init_done_event = threading.Event() # To signal model init is complete

        if self._model_manager:
            logger.info(f"ModelManager of type '{type(self._model_manager).__name__}' provided.")
        else:
            logger.info("No ModelManager provided. Worker operating in simple mode regarding models.")
            self._model_init_done_event.set() # No model init to wait for if no manager
        if self._downloader:
            logger.info(f"ModelDownloader of type '{type(self._downloader).__name__}' configured.")

        logger.info(f"Created worker: ID={self.worker_id}, DeploymentID={self.deployment_id or 'N/A'}, Scheduler={scheduler_addr}")

        # Discover functions before setting signals? Maybe after. Let's do it here.
        self._discover_and_register_functions()

        self._verify_auth_token()

        signal.signal(signal.SIGINT, self._handle_interrupt)
        signal.signal(signal.SIGTERM, self._handle_interrupt)

    @staticmethod
    def _normalize_scheduler_addrs(primary: str, addrs: Optional[List[str]]) -> List[str]:
        unique: List[str] = []
        for addr in [primary] + (addrs or []):
            addr = (addr or "").strip()
            if addr and addr not in unique:
                unique.append(addr)
        return unique

    @staticmethod
    def _extract_leader_addr(details: Optional[str]) -> Optional[str]:
        if not details:
            return None
        if details.startswith("not_leader:"):
            leader = details.split("not_leader:", 1)[1].strip()
            return leader or None
        return None

    def _set_scheduler_addr(self, addr: str) -> None:
        addr = addr.strip()
        if not addr:
            return
        self.scheduler_addr = addr
        if addr not in self.scheduler_addrs:
            self.scheduler_addrs.insert(0, addr)

    def _iter_scheduler_addrs(self) -> Iterator[str]:
        seen = set()
        for addr in self.scheduler_addrs:
            addr = addr.strip()
            if not addr or addr in seen:
                continue
            seen.add(addr)
            yield addr

    def _verify_auth_token(self) -> None:
        if not self.auth_token or not self._jwks_cache:
            return
        try:
            header = jwt.get_unverified_header(self.auth_token)
            kid = header.get("kid")
            key = self._jwks_cache.get_key(kid)
            if not key:
                raise ValueError("JWKS key not found for token")
            options = {"verify_aud": bool(self._jwt_audience)}
            jwt.decode(
                self.auth_token,
                key=key,
                algorithms=["RS256"],
                audience=self._jwt_audience or None,
                issuer=self._jwt_issuer or None,
                options=options,
            )
            logger.info("Worker auth token verified against scheduler JWKS.")
        except Exception as e:
            logger.error(f"Worker auth token verification failed: {e}")
            raise

    def _format_error(self, message: str, retryable: bool) -> str:
        return json.dumps({
            "message": message,
            "retryable": retryable,
        })

    def _emit_progress_event(self, event: Dict[str, Any]) -> None:
        try:
            run_id = event.get("run_id") or ""
            event_type = event.get("type") or ""
            payload = event.get("payload") or {}
            if "timestamp" not in payload:
                payload = dict(payload)
                payload["timestamp"] = event.get("timestamp", time.time())
            payload_json = json.dumps(payload).encode("utf-8")
            msg = pb.WorkerSchedulerMessage(
                worker_event=pb.WorkerEvent(
                    run_id=run_id,
                    event_type=event_type,
                    payload_json=payload_json,
                )
            )
            self._send_message(msg)
        except Exception:
            logger.exception("Failed to emit progress event")


    def _set_gpu_busy_status(self, busy: bool, func_name_for_log: str = "") -> None:
        with self._gpu_busy_lock:
            if self._is_gpu_busy == busy:
                return
            self._is_gpu_busy = busy
        if func_name_for_log:
            logger.info(f"GPU status changed to {busy} due to function '{func_name_for_log}'.")
        else:
            logger.info(f"GPU status changed to {busy}.")


    def _get_gpu_busy_status(self) -> bool:
        with self._gpu_busy_lock:
            return self._is_gpu_busy


    def _discover_and_register_functions(self) -> None:
        """Discover and register functions marked with @worker_function."""
        logger.info(f"Discovering worker functions in modules: {self.user_module_names}...")
        discovered_count = 0
        for module_name in self.user_module_names:
            try:
                module = importlib.import_module(module_name)
                logger.debug(f"Inspecting module: {module_name}")
                for name, obj in inspect.getmembers(module):
                    if inspect.isfunction(obj) and hasattr(obj, '_is_worker_function'):
                        if getattr(obj, '_is_worker_function') is True:
                            # Found a decorated function
                            original_func = obj # Keep reference to the actual decorated function
                            func_name = original_func.__name__ # Use the real function name

                            if func_name in self._actions:
                                logger.warning(f"Function '{func_name}' from module '{module_name}' conflicts with an already registered function. Skipping.")
                                continue

                            resources: ResourceRequirements = getattr(original_func, '_worker_resources', ResourceRequirements())
                            self._discovered_resources[func_name] = resources

                            expects_pipeline_flag = resources.expects_pipeline_arg
                            payload_type = self._infer_payload_type(original_func, expects_pipeline_flag)
                            return_type = self._infer_return_type(original_func)
                            if payload_type is None or return_type is None:
                                logger.error(
                                    "Skipping function '%s' due to invalid or missing payload type annotation.",
                                    func_name,
                                )
                                continue

                            try:
                                input_schema = msgspec.json.schema(payload_type)
                                output_schema = msgspec.json.schema(return_type)
                                self._function_schemas[func_name] = (
                                    json.dumps(input_schema, separators=(",", ":"), sort_keys=True).encode("utf-8"),
                                    json.dumps(output_schema, separators=(",", ":"), sort_keys=True).encode("utf-8"),
                                )
                            except Exception as exc:
                                logger.error("Failed to generate msgspec JSON schema for '%s': %s", func_name, exc)
                                continue

                            # Create the wrapper for gRPC/msgpack interaction
                            def create_wrapper(
                                captured_func: Callable[..., Any],
                                captured_name: str,
                                captured_payload_type: type[msgspec.Struct],
                                captured_return_type: type[msgspec.Struct],
                                func_expects_pipeline: bool = False,
                            ) -> Callable[[ActionContext, Optional[Any], bytes], bytes]:
                                @functools.wraps(captured_func) # Preserve metadata of original user func
                                def wrapper(ctx: ActionContext, pipeline_instance: Optional[Any], input_bytes: bytes) -> bytes:
                                    try:
                                        input_obj = msgspec.msgpack.decode(input_bytes, type=captured_payload_type)
                                        self._materialize_assets(ctx.run_id, input_obj)
                                        # Pass the context and deserialized input to the *original* user function
                                        if func_expects_pipeline: # Only pass pipeline if function expects it
                                            if pipeline_instance is None:
                                                err_msg = f"Function '{captured_name}' expected a pipeline argument, but None was provided by the Worker core."
                                                logger.error(err_msg)
                                                raise ValueError(err_msg)
                                            result = captured_func(ctx, pipeline_instance, input_obj)
                                        else:
                                            result = captured_func(ctx, input_obj) # For functions not needing a model

                                        if ctx.is_canceled():
                                             raise InterruptedError("Task was canceled during execution")
                                        if not isinstance(result, captured_return_type):
                                            raise TypeError(
                                                f"Function {captured_name} returned {type(result)!r}, "
                                                f"expected {captured_return_type!r}"
                                            )
                                        # Ensure result is bytes after msgspec msgpack serialization
                                        packed_result = msgspec.msgpack.encode(result)
                                        if not isinstance(packed_result, bytes):
                                            raise TypeError(
                                                f"Function {captured_name} did not return msgspec-serializable data resulting in bytes"
                                            )
                                        return packed_result
                                    except InterruptedError as ie: # Catch cancellation specifically
                                        logger.warning(f"Function {captured_name} run {ctx.run_id} was interrupted.")
                                        raise # Re-raise to be handled in _execute_function
                                    except Exception as e:
                                        logger.exception(f"Error during execution of function {captured_name} (run_id: {ctx.run_id})")
                                        raise # Re-raise to be caught by _execute_function
                                return wrapper

                            self._actions[func_name] = create_wrapper(
                                original_func,
                                func_name,
                                payload_type,
                                return_type,
                                func_expects_pipeline=expects_pipeline_flag,
                            )
                            logger.info(f"Registered function: '{func_name}' from module '{module_name}' with resources: {resources}")
                            discovered_count += 1

            except ImportError:
                logger.error(f"Could not import user module: {module_name}")
            except Exception as e:
                logger.exception(f"Error during discovery in module {module_name}: {e}")

        if discovered_count == 0:
             logger.warning(f"No functions decorated with @worker_function found in specified modules: {self.user_module_names}")
        else:
             logger.info(f"Discovery complete. Found {discovered_count} worker functions.")

    def _infer_payload_type(
        self,
        func: Callable[..., Any],
        expects_pipeline: bool,
    ) -> Optional[type[msgspec.Struct]]:
        signature = inspect.signature(func)
        params = list(signature.parameters.values())
        expected_params = 3 if expects_pipeline else 2
        if len(params) != expected_params:
            logger.error(
                "Function '%s' has %d parameters but expected %d.",
                func.__name__,
                len(params),
                expected_params,
            )
            return None

        payload_param = params[-1]
        try:
            type_hints = typing.get_type_hints(func, globalns=func.__globals__)
        except Exception as exc:
            logger.error("Failed to resolve type hints for '%s': %s", func.__name__, exc)
            return None

        payload_type = type_hints.get(payload_param.name)
        if payload_type is None:
            logger.error("Function '%s' is missing a payload type annotation.", func.__name__)
            return None

        if not isinstance(payload_type, type) or not issubclass(payload_type, msgspec.Struct):
            logger.error(
                "Function '%s' payload type must be a msgspec.Struct, got %r.",
                func.__name__,
                payload_type,
            )
            return None

        return payload_type

    def _infer_return_type(self, func: Callable[..., Any]) -> Optional[type[msgspec.Struct]]:
        try:
            type_hints = typing.get_type_hints(func, globalns=func.__globals__)
        except Exception as exc:
            logger.error("Failed to resolve return type hints for '%s': %s", func.__name__, exc)
            return None

        return_type = type_hints.get("return")
        if return_type is None:
            logger.error("Function '%s' is missing a return type annotation.", func.__name__)
            return None

        if not isinstance(return_type, type) or not issubclass(return_type, msgspec.Struct):
            logger.error(
                "Function '%s' return type must be a msgspec.Struct, got %r.",
                func.__name__,
                return_type,
            )
            return None

        return return_type

    def _send_message(self, message: WorkerSchedulerMessage) -> None:
        """Add a message to the outgoing queue."""
        if self._running and not self._stop_event.is_set():
            try:
                self._outgoing_queue.put_nowait(message)
            except queue.Full:
                 logger.error("Outgoing message queue is full. Message dropped!")
        else:
            logger.warning("Attempted to send message while worker is stopping or stopped.")

    def _materialize_assets(self, run_id: str, obj: Any) -> None:
        if isinstance(obj, Asset):
            self._materialize_asset(run_id, obj)
            return
        if isinstance(obj, list):
            for it in obj:
                self._materialize_assets(run_id, it)
            return
        if isinstance(obj, dict):
            for it in obj.values():
                self._materialize_assets(run_id, it)
            return
        fields = getattr(obj, "__struct_fields__", None)
        if fields and isinstance(fields, (tuple, list)):
            for name in fields:
                try:
                    self._materialize_assets(run_id, getattr(obj, name))
                except Exception:
                    continue

    def _materialize_asset(self, run_id: str, asset: Asset) -> None:
        if asset.local_path:
            return
        ref = (asset.ref or "").strip()
        if not ref:
            return

        base_dir = os.getenv("WORKER_RUN_DIR", "/tmp/cozy").rstrip("/")
        local_inputs_dir = os.path.join(base_dir, run_id, "inputs")
        os.makedirs(local_inputs_dir, exist_ok=True)
        cache_dir = os.getenv("WORKER_CACHE_DIR", os.path.join(base_dir, "cache")).rstrip("/")
        os.makedirs(cache_dir, exist_ok=True)

        max_bytes = int(os.getenv("WORKER_MAX_INPUT_FILE_BYTES", str(200 * 1024 * 1024)))

        # External URL inputs (download directly into the run folder).
        if ref.startswith("http://") or ref.startswith("https://"):
            if _url_is_blocked(ref):
                raise RuntimeError("input url blocked")
            ext = os.path.splitext(urllib.parse.urlparse(ref).path)[1] or os.path.splitext(ref)[1]
            name_hash = hashlib.sha256(ref.encode("utf-8")).hexdigest()[:32]
            local_path = os.path.join(local_inputs_dir, f"{name_hash}{ext}")
            size, sha256_hex, mime = self._download_url_to_file(ref, local_path, max_bytes)
            asset.local_path = local_path
            if not asset.tenant_id:
                asset.tenant_id = self.tenant_id
            asset.mime_type = mime
            asset.size_bytes = size
            asset.sha256 = sha256_hex
            return

        # Cozy Hub file ref (tenant scoped) - use orchestrator file API with HEAD+cache.
        base = _require_file_api_base_url()
        token = _require_file_api_token()
        url = f"{base}/api/v1/file/{_encode_ref_for_url(ref)}"

        head_req = _http_request("HEAD", url, token)
        with urllib.request.urlopen(head_req, timeout=10) as resp:
            if resp.status < 200 or resp.status >= 300:
                raise RuntimeError(f"failed to stat asset ({resp.status})")
            sha256_hex = (resp.headers.get("X-Cozy-SHA256") or "").strip()
            size_hdr = (resp.headers.get("X-Cozy-Size-Bytes") or "").strip()
            mime = (resp.headers.get("X-Cozy-Mime-Type") or "").strip()
        size = int(size_hdr) if size_hdr.isdigit() else 0
        if max_bytes > 0 and size > max_bytes:
            raise RuntimeError("input file too large")

        ext = os.path.splitext(ref)[1]
        if not ext and mime:
            guessed = {
                "image/png": ".png",
                "image/jpeg": ".jpg",
                "image/webp": ".webp",
                "image/gif": ".gif",
            }.get(mime)
            ext = guessed or ""

        if not sha256_hex:
            sha256_hex = hashlib.sha256(ref.encode("utf-8")).hexdigest()
        cache_name = f"{sha256_hex[:32]}{ext}"
        cache_path = os.path.join(cache_dir, cache_name)

        if not os.path.exists(cache_path):
            get_req = _http_request("GET", url, token)
            with urllib.request.urlopen(get_req, timeout=30) as resp:
                if resp.status < 200 or resp.status >= 300:
                    raise RuntimeError(f"failed to download asset ({resp.status})")
                _size, _sha = self._stream_to_file(resp, cache_path, max_bytes)
                if not size:
                    size = _size
                if not sha256_hex:
                    sha256_hex = _sha

        local_path = os.path.join(local_inputs_dir, cache_name)
        if not os.path.exists(local_path):
            try:
                os.link(cache_path, local_path)
            except Exception:
                try:
                    import shutil

                    shutil.copyfile(cache_path, local_path)
                except Exception:
                    local_path = cache_path

        if not mime:
            with open(local_path, "rb") as f:
                head = f.read(512)
            mime = _infer_mime_type(ref, head)

        asset.local_path = local_path
        if not asset.tenant_id:
            asset.tenant_id = self.tenant_id
        asset.mime_type = mime or None
        asset.size_bytes = size or None
        asset.sha256 = sha256_hex or None

    def _download_url_to_file(self, src: str, dst: str, max_bytes: int) -> Tuple[int, str, Optional[str]]:
        attempts = int(os.getenv("WORKER_DOWNLOAD_RETRIES", "3"))
        attempt = 0
        last_err: Optional[Exception] = None
        while attempt < max(1, attempts):
            attempt += 1
            try:
                client = urllib.request.build_opener()
                req = urllib.request.Request(src, method="GET")
                with client.open(req, timeout=30) as resp:
                    size, sha = self._stream_to_file(resp, dst, max_bytes)
                with open(dst, "rb") as f:
                    head = f.read(512)
                mime = _infer_mime_type(src, head)
                return size, sha, mime
            except Exception as e:
                last_err = e
                if attempt >= max(1, attempts):
                    break
                sleep_s = min(10.0, 0.5 * (2 ** (attempt - 1))) + random.random() * 0.2
                time.sleep(sleep_s)
        raise RuntimeError(f"failed to download url: {last_err}")

    def _stream_to_file(self, src: Any, dst: str, max_bytes: int) -> Tuple[int, str]:
        tmp = f"{dst}.tmp-{os.getpid()}-{threading.get_ident()}-{random.randint(0, 1_000_000)}"
        total = 0
        h = hashlib.sha256()
        try:
            with open(tmp, "wb") as out:
                while True:
                    chunk = src.read(1024 * 1024)
                    if not chunk:
                        break
                    total += len(chunk)
                    if total > max_bytes:
                        raise RuntimeError("input file too large")
                    h.update(chunk)
                    out.write(chunk)
            os.replace(tmp, dst)
        finally:
            try:
                if os.path.exists(tmp):
                    os.remove(tmp)
            except Exception:
                pass
        return total, h.hexdigest()

    def connect(self) -> bool:
        """Connect to the scheduler.

        Returns:
            bool: True if connection was successful, False otherwise.
        """
        attempted: set[str] = set()
        while True:
            addr = None
            if self._leader_hint and self._leader_hint not in attempted:
                addr = self._leader_hint
                self._leader_hint = None
            else:
                for candidate in self._iter_scheduler_addrs():
                    if candidate not in attempted:
                        addr = candidate
                        break
            if not addr:
                break
            attempted.add(addr)
            self._set_scheduler_addr(addr)
            if self._connect_once():
                return True
        return False

    def _connect_once(self) -> bool:
        try:
            if self.use_tls:
                # TODO: Add proper credential loading if needed
                creds = grpc.ssl_channel_credentials()
                self._channel = grpc.secure_channel(self.scheduler_addr, creds)
            else:
                self._channel = grpc.insecure_channel(self.scheduler_addr)

            interceptors = []
            if self.auth_token:
                interceptors.append(_AuthInterceptor(self.auth_token))

            if interceptors:
                self._channel = grpc.intercept_channel(self._channel, *interceptors)

            self._stub = pb_grpc.SchedulerWorkerServiceStub(self._channel)

            # Start the bidirectional stream
            request_iterator = self._outgoing_message_iterator()
            self._stream = self._stub.ConnectWorker(request_iterator)

            logger.info(f"Attempting to connect to scheduler at {self.scheduler_addr}...")

            # Send initial registration immediately
            self._register_worker(is_heartbeat=False)

            # Start the receive loop in a separate thread *after* stream is initiated
            self._receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
            self._receive_thread.start()

            # Start heartbeat thread
            self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
            self._heartbeat_thread.start()

            logger.info(f"Successfully connected to scheduler at {self.scheduler_addr}")
            self._reconnect_count = 0
            return True

        except grpc.RpcError as e:
            # Access code() and details() methods for RpcError
            code = e.code() if hasattr(e, 'code') and callable(e.code) else grpc.StatusCode.UNKNOWN
            details = e.details() if hasattr(e, 'details') and callable(e.details) else str(e)
            leader = self._extract_leader_addr(details)
            if code == grpc.StatusCode.FAILED_PRECONDITION and leader:
                logger.warning(f"Scheduler returned not_leader for {self.scheduler_addr}; redirecting to {leader}")
                self._leader_hint = leader
                self._set_scheduler_addr(leader)
            else:
                logger.error(f"Failed to connect to scheduler: {code} - {details}")
            self._close_connection()
            return False
        except Exception as e:
            logger.exception(f"Unexpected error connecting to scheduler: {e}")
            self._close_connection()
            return False

    def _outgoing_message_iterator(self) -> Iterator[WorkerSchedulerMessage]:
        """Yields messages from the outgoing queue to send to the scheduler."""
        while not self._stop_event.is_set():
            try:
                # Block for a short time to allow stopping gracefully
                message = self._outgoing_queue.get(timeout=0.1)
                yield message
                # self._outgoing_queue.task_done() # Not needed if not joining queue
            except queue.Empty:
                continue
            except Exception as e:
                 if not self._stop_event.is_set():
                     logger.exception(f"Error in outgoing message iterator: {e}")
                     self._handle_connection_error()
                     break # Exit iterator on error

    def _heartbeat_loop(self) -> None:
        """Periodically sends heartbeat messages."""
        while not self._stop_event.wait(HEARTBEAT_INTERVAL):
            try:
                self._register_worker(is_heartbeat=True)
                logger.debug("Sent heartbeat to scheduler")
            except Exception as e:
                if not self._stop_event.is_set():
                    logger.error(f"Error sending heartbeat: {e}")
                    self._handle_connection_error()
                    break # Stop heartbeating on error

    def _register_worker(self, is_heartbeat: bool = False) -> None:
        """Create and send a registration/heartbeat message."""
        try:
            mem = psutil.virtual_memory()
            cpu_cores = os.cpu_count() or 0

            gpu_count = 0
            gpu_total_mem = 0
            vram_models = []
            gpu_used_mem = 0
            gpu_free_mem = 0
            gpu_name = ""
            gpu_driver = ""

            if torch and torch.cuda.is_available():
                gpu_count = torch.cuda.device_count()
                if gpu_count > 0:
                    try:
                        props = torch.cuda.get_device_properties(0)
                        gpu_total_mem = props.total_memory
                        gpu_used_mem = torch.cuda.memory_allocated(0)
                        gpu_name = props.name
                        gpu_driver = torch.version.cuda or ""
                        try:
                            free_mem, total_mem = torch.cuda.mem_get_info(0)
                            gpu_total_mem = total_mem
                            gpu_used_mem = total_mem - free_mem
                            gpu_free_mem = free_mem
                        except Exception:
                            pass
                        logger.debug(f"GPU: {props.name}, VRAM total={gpu_total_mem}, used={gpu_used_mem}, cuda={torch.version.cuda}")
                    except Exception as gpu_err:
                         logger.warning(f"Could not get GPU properties: {gpu_err}")

            fake_gpu_count = os.getenv("WORKER_FAKE_GPU_COUNT")
            if fake_gpu_count:
                try:
                    gpu_count = int(fake_gpu_count)
                    if gpu_count > 0:
                        fake_mem = int(os.getenv("WORKER_FAKE_GPU_MEMORY_BYTES", str(24 * 1024 * 1024 * 1024)))
                        gpu_total_mem = fake_mem
                        gpu_used_mem = 0
                        gpu_free_mem = fake_mem
                        gpu_name = os.getenv("WORKER_FAKE_GPU_NAME", "FakeGPU")
                        gpu_driver = os.getenv("WORKER_FAKE_GPU_DRIVER", "fake")
                except ValueError:
                    logger.warning("Invalid WORKER_FAKE_GPU_COUNT; ignoring fake GPU override.")

            supports_model_loading_flag = False
            # current_models = []
            if self._model_manager:
                vram_models = self._model_manager.get_vram_loaded_models()
                supports_model_loading_flag = True 

            function_concurrency = {}
            for func_name, req in self._discovered_resources.items():
                if req and req.max_concurrency:
                    function_concurrency[func_name] = int(req.max_concurrency)

            cuda_version = os.getenv("WORKER_CUDA_VERSION", "").strip()
            torch_version = os.getenv("WORKER_TORCH_VERSION", "").strip()
            if torch is not None:
                if not torch_version:
                    torch_version = getattr(torch, "__version__", "") or ""
                if not cuda_version:
                    cuda_version = getattr(torch.version, "cuda", "") or ""
            if not cuda_version:
                cuda_version = os.getenv("CUDA_VERSION", "").strip() or os.getenv("NVIDIA_CUDA_VERSION", "").strip()

            function_schemas = []
            for fname, (in_schema, out_schema) in self._function_schemas.items():
                try:
                    function_schemas.append(
                        pb.FunctionSchema(
                            name=fname,
                            input_schema_json=in_schema,
                            output_schema_json=out_schema,
                        )
                    )
                except Exception:
                    continue

            resources = pb.WorkerResources(
                worker_id=self.worker_id,
                deployment_id=self.deployment_id,
                # tenant_id=self.tenant_id,
                runpod_pod_id=self.runpod_pod_id,
                gpu_is_busy=self._get_gpu_busy_status(),
                cpu_cores=cpu_cores,
                memory_bytes=mem.total,
                gpu_count=gpu_count,
                gpu_memory_bytes=gpu_total_mem,
                gpu_memory_used_bytes=gpu_used_mem,
                gpu_memory_free_bytes=gpu_free_mem,
                gpu_name=gpu_name,
                gpu_driver=gpu_driver,
                max_concurrency=self.max_concurrency,
                function_concurrency=function_concurrency,
                cuda_version=cuda_version,
                torch_version=torch_version,
                available_functions=list(self._actions.keys()),
                available_models=vram_models,
                supports_model_loading=supports_model_loading_flag,
                function_schemas=function_schemas,
            )
            registration = pb.WorkerRegistration(
                resources=resources,
                is_heartbeat=is_heartbeat
            )
            message = pb.WorkerSchedulerMessage(worker_registration=registration)
            # logger.info(f"DEBUG: Preparing to send registration. Resource object: {resources}")
            # logger.info(f"DEBUG: Value being sent for runpod_pod_id: '{resources.runpod_pod_id}'")
            self._send_message(message)
        except Exception as e:
            logger.error(f"Failed to create or send registration/heartbeat: {e}")

    def run(self) -> None:
        """Run the worker, connecting to the scheduler and processing tasks."""
        if self._running:
            logger.warning("Worker is already running")
            return

        self._running = True
        self._stop_event.clear()
        self._reconnect_count = 0 # Reset reconnect count on new run
        self._draining = False

        while self._running and not self._stop_event.is_set():
            self._reconnect_count += 1
            logger.info(f"Connection attempt {self._reconnect_count}...")
            if self.connect():
                # Successfully connected, wait for stop signal or disconnection
                logger.info("Connection successful. Worker running.")
                self._stop_event.wait() # Wait here until stopped or disconnected
                logger.info("Worker run loop received stop/disconnect signal.")
                # If stopped normally (self.stop() called), _running will be False
                # If disconnected, connect() failed, threads stopped, _handle_connection_error called _stop_event.set()
            else:
                # Connection failed
                if self.max_reconnect_attempts > 0 and self._reconnect_count >= self.max_reconnect_attempts:
                    logger.error("Failed to connect after maximum attempts. Stopping worker.")
                    self._running = False # Ensure loop terminates
                    break

                if self._running and not self._stop_event.is_set():
                    backoff = self._reconnect_delay_base * (2 ** max(self._reconnect_count - 1, 0))
                    if self._reconnect_delay_max > 0:
                        backoff = min(backoff, self._reconnect_delay_max)
                    jitter = random.uniform(0, self._reconnect_jitter) if self._reconnect_jitter > 0 else 0
                    delay = backoff + jitter
                    logger.info(f"Connection attempt {self._reconnect_count} failed. Retrying in {delay:.2f} seconds...")
                    # Wait for delay, but break if stop event is set during wait
                    if self._stop_event.wait(delay):
                        logger.info("Stop requested during reconnect delay.")
                        break # Exit if stopped while waiting
            # After a failed attempt or disconnect, clear stop event for next retry
            if self._running:
                 self._stop_event.clear()

        # Cleanup after loop exits (either max attempts reached or manual stop)
        self.stop()

    def _handle_interrupt(self, sig: int, frame: Optional[Any]) -> None:
        """Handle interrupt signal (Ctrl+C)."""
        logger.info(f"Received signal {sig}, shutting down gracefully.")
        self.stop()

    def stop(self) -> None:
        """Stop the worker and clean up resources."""
        if not self._running and not self._stop_event.is_set(): # Check if already stopped or stopping
            # Avoid multiple stop calls piling up
            # logger.debug("Stop called but worker already stopped or stopping.")
            return

        logger.info("Stopping worker...")
        self._draining = True
        self._running = False # Signal loops to stop
        self._stop_event.set() # Wake up any waiting threads

        # Cancel any active tasks
        active_task_ids = []
        if self._drain_timeout_seconds > 0:
            deadline = time.time() + self._drain_timeout_seconds
            while time.time() < deadline:
                with self._active_tasks_lock:
                    remaining = len(self._active_tasks)
                if remaining == 0:
                    break
                time.sleep(0.2)

        with self._active_tasks_lock:
            active_task_ids = list(self._active_tasks.keys())
            for run_id in active_task_ids:
                ctx = self._active_tasks.get(run_id)
                if ctx:
                    logger.debug(f"Cancelling active task {run_id} during stop.")
                    ctx.cancel()
            # Don't clear here, allow _execute_function to finish and remove

        # Wait for threads (give them a chance to finish)
        # Stop heartbeat first
        if self._heartbeat_thread and self._heartbeat_thread.is_alive():
             logger.debug("Joining heartbeat thread...")
             self._heartbeat_thread.join(timeout=1.0)

        # The outgoing iterator might be blocked on queue.get, stop_event wakes it

        # Close the gRPC connection (this might interrupt the receive loop)
        self._close_connection()

        # Wait for receive thread
        if self._receive_thread and self._receive_thread.is_alive():
            logger.debug("Joining receive thread...")
            self._receive_thread.join(timeout=2.0)

        # Clear outgoing queue after threads are stopped
        logger.debug("Clearing outgoing message queue...")
        while not self._outgoing_queue.empty():
            try:
                self._outgoing_queue.get_nowait()
            except queue.Empty:
                break

        logger.info("Worker stopped.")
        # Reset stop event in case run() is called again
        self._stop_event.clear()

    def _close_connection(self) -> None:
        """Close the gRPC channel and reset state."""
        if self._stream:
             try:
                  # Attempt to cancel the stream from the client side
                  # This might help the server side release resources quicker
                  # Note: Behavior might vary depending on server implementation
                  if hasattr(self._stream, 'cancel') and callable(self._stream.cancel):
                     self._stream.cancel()
                     logger.debug("gRPC stream cancelled.")
             except Exception as e:
                  logger.warning(f"Error cancelling gRPC stream: {e}")
        self._stream = None

        if self._channel:
            try:
                self._channel.close()
                logger.debug("gRPC channel closed.")
            except Exception as e:
                 logger.error(f"Error closing gRPC channel: {e}")
        self._channel = None
        self._stub = None


    def _receive_loop(self) -> None:
        """Loop to receive messages from the scheduler via the stream."""
        logger.info("Receive loop started.")
        try:
            if not self._stream:
                 logger.error("Receive loop started without a valid stream.")
                 # Don't call _handle_connection_error here, connect should have failed
                 return

            for message in self._stream:
                # Check stop event *before* processing
                if self._stop_event.is_set():
                    logger.debug("Stop event set during iteration, exiting receive loop.")
                    break
                try:
                    self._process_message(message)
                except Exception as e:
                    # Log errors processing individual messages but continue loop
                    logger.exception(f"Error processing message: {e}")

        except grpc.RpcError as e:
            # RpcError indicates a problem with the gRPC connection itself
            code = e.code() if hasattr(e, 'code') and callable(e.code) else grpc.StatusCode.UNKNOWN
            details = e.details() if hasattr(e, 'details') and callable(e.details) else str(e)

            if self._stop_event.is_set():
                 # If stopping, cancellation is expected
                 if code == grpc.StatusCode.CANCELLED:
                     logger.info("gRPC stream cancelled gracefully during shutdown.")
                 else:
                     logger.warning(f"gRPC error during shutdown: {code} - {details}")
            elif code == grpc.StatusCode.FAILED_PRECONDITION:
                leader = self._extract_leader_addr(details)
                if leader:
                    logger.warning(f"Scheduler redirect received; reconnecting to leader at {leader}")
                    self._leader_hint = leader
                    self._set_scheduler_addr(leader)
                self._handle_connection_error()
            elif code == grpc.StatusCode.CANCELLED:
                logger.warning("gRPC stream unexpectedly cancelled by server or network.")
                self._handle_connection_error()
            elif code in (grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.DEADLINE_EXCEEDED, grpc.StatusCode.INTERNAL):
                 logger.warning(f"gRPC connection lost ({code}). Attempting reconnect.")
                 self._handle_connection_error()
            else:
                 logger.error(f"Unhandled gRPC error in receive loop: {code} - {details}")
                 self._handle_connection_error() # Attempt reconnect on unknown errors too
        except Exception as e:
            # Catch-all for non-gRPC errors in the loop
            if not self._stop_event.is_set():
                logger.exception(f"Unexpected error in receive loop: {e}")
                self._handle_connection_error() # Attempt reconnect
        finally:
             logger.info("Receive loop finished.")

    def _handle_connection_error(self) -> None:
         """Handles actions needed when a connection error occurs during run."""
         if self._running and not self._stop_event.is_set():
             logger.warning("Connection error detected. Signaling main loop to reconnect...")
             self._close_connection() # Ensure resources are closed before reconnect attempt
             self._stop_event.set() # Signal run loop to attempt reconnection
         # else: # Already stopping or stopped
             # logger.debug("Connection error detected but worker is already stopping.")


    def _process_message(self, message: WorkerSchedulerMessage) -> None:
        """Process a single message received from the scheduler."""
        msg_type = message.WhichOneof('msg')
        # logger.debug(f"Received message of type: {msg_type}")

        if msg_type == 'run_request':
            self._handle_run_request(message.run_request)
        elif msg_type == 'load_model_cmd':
            # TODO: Implement model loading logic
            # model_id = message.load_model_cmd.model_id
            # logger.warning(f"Received load_model_cmd for {model_id}, but not yet implemented.")
            # # Send result back (failure for now)
            # result = pb.LoadModelResult(model_id=model_id, success=False, error_message="Model loading not implemented")
            # self._send_message(pb.WorkerSchedulerMessage(load_model_result=result))
            self._handle_load_model_cmd(message.load_model_cmd)
        elif msg_type == 'unload_model_cmd':
            # TODO: Implement model unloading logic
            model_id = message.unload_model_cmd.model_id
            logger.warning(f"Received unload_model_cmd for {model_id}, but not yet implemented.")
            result = pb.UnloadModelResult(model_id=model_id, success=False, error_message="Model unloading not implemented")
            self._send_message(pb.WorkerSchedulerMessage(unload_model_result=result))
        elif msg_type == 'interrupt_run_cmd':
            run_id = message.interrupt_run_cmd.run_id
            self._handle_interrupt_request(run_id)
        # Add handling for other message types if needed (e.g., config updates)
        elif msg_type == 'deployment_model_config':
            if self._model_manager:
                logger.info(f"Received DeploymentModelConfig: {message.deployment_model_config.supported_model_ids}")
                self._supported_model_ids_from_scheduler = list(message.deployment_model_config.supported_model_ids)
                self._model_init_done_event.clear() # Clear before starting new init
                model_init_thread = threading.Thread(target=self._process_deployment_config_async_wrapper, daemon=True)
                model_init_thread.start()
            else:
                logger.info("Received DeploymentModelConfig, but no model manager configured. Ignoring.")
                self._model_init_done_event.set() # Signal completion as there's nothing to do
        elif msg_type is None:
             logger.warning("Received empty message from scheduler.")
        else:
            logger.warning(f"Received unhandled message type: {msg_type}")

    def _process_deployment_config_async_wrapper(self) -> None:
        if not self._model_manager or self._supported_model_ids_from_scheduler is None:
            self._model_init_done_event.set()
            return
        
        loop = None
        try:
            # Get or create an event loop for this thread
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)

            loop.run_until_complete(
                self._model_manager.process_supported_models_config(
                    self._supported_model_ids_from_scheduler,
                    self._downloader
                )
            )
            logger.info("Model configuration and downloads (if any) processed.")
        except Exception as e:
            logger.exception(f"Error during model_manager.process_supported_models_config: {e}")
        finally:
            if loop and not loop.is_running() and not loop.is_closed(): # Clean up loop if we created it
                loop.close()
            self._model_init_done_event.set() # Signal completion or failure

    def _handle_load_model_cmd(self, cmd: LoadModelCommand) -> None:
        model_id = cmd.model_id
        logger.info(f"Received LoadModelCommand for: {model_id}")
        success = False; error_msg = ""
        if not self._model_manager:
            error_msg = "LoadModelCommand: No model manager configured on worker."
            logger.error(error_msg)
        else:
            try:
                # Wait for initial model downloads if they haven't finished
                if not self._model_init_done_event.is_set():
                    logger.info(f"LoadModelCmd ({model_id}): Waiting for initial model setup...")
                    # Timeout for this wait, can be adjusted
                    if not self._model_init_done_event.wait(timeout=300.0): # 5 minutes
                         raise TimeoutError("Timeout waiting for model initialization before VRAM load.")
                
                logger.info(f"Model Memory Manager attempting to load '{model_id}' into VRAM...")
                # load_model_into_vram is async
                success = asyncio.run(self._model_manager.load_model_into_vram(model_id))
                if success: logger.info(f"Model '{model_id}' loaded to VRAM by Model Memory Manager.")
                else: error_msg = f"MMM.load_model_into_vram failed for '{model_id}'."; logger.error(error_msg)
            except Exception as e: error_msg = f"Exception in mmm.load_model_into_vram for '{model_id}': {e}"; logger.exception(error_msg)
        
        result = pb.LoadModelResult(model_id=model_id, success=success, error_message=error_msg)
        self._send_message(pb.WorkerSchedulerMessage(load_model_result=result))


    def _handle_run_request(self, request: TaskExecutionRequest) -> None:
        """Handle a task execution request from the scheduler."""
        run_id = request.run_id
        function_name = request.function_name
        input_payload = request.input_payload
        required_model_id_for_exec = ""
        timeout_ms = int(getattr(request, "timeout_ms", 0) or 0)
        tenant_id = str(getattr(request, "tenant_id", "") or "") or (self.tenant_id or "")
        user_id = str(getattr(request, "user_id", "") or "")

        if request.required_models and len(request.required_models) > 0:
            required_model_id_for_exec = request.required_models[0]

        logger.info(f"Received Task request: run_id={run_id}, function={function_name}, model='{required_model_id_for_exec or 'None'}'")

        func_wrapper = self._actions.get(function_name)
        if not func_wrapper:
            error_msg = f"Unknown function requested: {function_name}"
            logger.error(error_msg)
            self._send_task_result(run_id, False, None, error_msg)
            return
        if self.max_input_bytes > 0 and len(input_payload) > self.max_input_bytes:
            error_msg = f"Input payload too large: {len(input_payload)} bytes (max {self.max_input_bytes})"
            logger.error(error_msg)
            self._send_task_result(run_id, False, None, error_msg)
            return
        if self._draining:
            error_msg = "Worker is draining; refusing new tasks"
            logger.warning(error_msg)
            self._send_task_result(run_id, False, None, error_msg)
            return

        ctx = ActionContext(
            run_id,
            emitter=self._emit_progress_event,
            tenant_id=tenant_id or None,
            user_id=user_id or None,
            timeout_ms=timeout_ms if timeout_ms > 0 else None,
        )
        # Add to active tasks *before* starting thread
        with self._active_tasks_lock:
             # Double-check if task is already active (race condition mitigation)
             if run_id in self._active_tasks:
                  error_msg = f"Task with run_id {run_id} is already active (race condition?)."
                  logger.error(error_msg)
                  return # Avoid starting duplicate thread
             if self.max_concurrency > 0 and len(self._active_tasks) >= self.max_concurrency:
                  error_msg = f"Worker concurrency limit reached ({self.max_concurrency})."
                  logger.error(error_msg)
                  self._send_task_result(run_id, False, None, error_msg)
                  return
             resource_req = self._discovered_resources.get(function_name)
             func_limit = resource_req.max_concurrency if resource_req and resource_req.max_concurrency else 0
             if func_limit > 0 and self._active_function_counts.get(function_name, 0) >= func_limit:
                  error_msg = f"Function concurrency limit reached for {function_name} ({func_limit})."
                  logger.error(error_msg)
                  self._send_task_result(run_id, False, None, error_msg)
                  return
             self._active_tasks[run_id] = ctx
             if func_limit > 0:
                  self._active_function_counts[function_name] = self._active_function_counts.get(function_name, 0) + 1

        # Execute function in a separate thread to avoid blocking the receive loop
        thread = threading.Thread(
            target=self._execute_function,
            args=(ctx, function_name, func_wrapper, input_payload, required_model_id_for_exec),
            daemon=True,
        )
        thread.start()

    def _handle_interrupt_request(self, run_id: str) -> None:
        """Handle a request to interrupt/cancel a running task."""
        logger.info(f"Received interrupt request for run_id={run_id}")
        with self._active_tasks_lock:
            ctx = self._active_tasks.get(run_id)
            if ctx:
                ctx.cancel() # Set internal flag and event
            else:
                logger.warning(f"Could not interrupt task {run_id}: Not found in active tasks.")

    def _execute_function(
        self,
        ctx: ActionContext,
        function_name: str,
        func_to_execute: Callable[[ActionContext, Optional[Any], bytes], bytes],
        input_payload: bytes,
        required_model_id: str,
    ) -> None:
        """Execute the registered function and send the result/error back."""
        run_id = ctx.run_id
        output_payload: Optional[bytes] = None
        error_message: str = ""
        success = False

        # Determine if this function requires GPU and manage worker's GPU state
        func_requires_gpu = False
        resource_req = self._discovered_resources.get(function_name)
        if resource_req:
            func_requires_gpu = resource_req.requires_gpu
            func_expects_pipeline = resource_req.expects_pipeline_arg

        # Variable to track if this specific thread execution set the GPU busy
        this_thread_set_gpu_busy = False
        if func_requires_gpu:
            with self._gpu_busy_lock: # Lock to check and set self._is_gpu_busy atomically
                if not self._is_gpu_busy:
                    self._is_gpu_busy = True
                    this_thread_set_gpu_busy = True
                    logger.info(f"Worker GPU marked as BUSY by task {run_id} ({function_name}).")
                else:
                    logger.warning(f"Task {run_id} ({function_name}) requires GPU, but worker GPU was already marked busy. Proceeding...")

        active_pipeline_instance = None # To hold the pipeline for the user function
        try:
            if ctx.is_canceled(): 
                raise InterruptedError("Task cancelled before execution")
            
            if func_expects_pipeline:
                if not required_model_id and resource_req and resource_req.model_name:
                    required_model_id = str(resource_req.model_name)
                if not required_model_id:
                    raise ValueError(f"Function '{function_name}' expects a pipeline argument, but no model ID was provided.")
                
                if not self._model_manager:
                    raise RuntimeError(f"Function '{function_name}' expects a pipeline argument, but no model manager configured on worker.")
                
                if not self._model_init_done_event.is_set():
                    logger.info(f"Task {run_id} ({function_name}) waiting for initial model setup...")
                    if not self._model_init_done_event.wait(timeout=300.0): # 5 min timeout
                        raise TimeoutError(f"Timeout waiting for model initialization for task {run_id}")
                    logger.info(f"Initial model setup complete. Proceeding for task {run_id}.")
                
                logger.info(f"Task {run_id} ({function_name}) getting active pipeline for model '{required_model_id}'...")
                # get_active_pipeline is async
                active_pipeline_instance = asyncio.run(self._model_manager.get_active_pipeline(required_model_id))
                if not active_pipeline_instance:
                    raise RuntimeError(f"ModelManager failed to provide active pipeline for '{required_model_id}' for task {run_id}.")
                
                logger.info(f"Task {run_id} ({function_name}) obtained pipeline for model '{required_model_id}'.")

            # Execute the function wrapper (which handles deserialization/serialization)
            output_payload = func_to_execute(ctx, active_pipeline_instance, input_payload)
            # Check for cancellation *during* execution (func should check ctx.is_canceled)
            if ctx.is_canceled():
                raise InterruptedError("Task was cancelled during execution")

            if output_payload is not None and self.max_output_bytes > 0:
                if len(output_payload) > self.max_output_bytes:
                    raise ValueError(f"Output payload too large: {len(output_payload)} bytes (max {self.max_output_bytes})")

            success = True
            logger.info(f"Task {run_id} completed successfully.")

        except InterruptedError as e:
             error_message = self._format_error(str(e) or "Task was canceled", retryable=False)
             logger.warning(f"Task {run_id} was canceled: {error_message}")
             success = False # Explicitly set success to False on cancellation
        except RetryableError as e:
            error_message = self._format_error(f"{type(e).__name__}: {str(e)}", retryable=True)
            logger.error(f"Task {run_id} ({function_name}) retryable failure: {error_message}")
            success = False
        except FatalError as e:
            error_message = self._format_error(f"{type(e).__name__}: {str(e)}", retryable=False)
            logger.error(f"Task {run_id} ({function_name}) fatal failure: {error_message}")
            success = False
        except (ValueError, RuntimeError, TimeoutError) as ve_rte_to: # Catch specific errors we raise
            retryable = isinstance(ve_rte_to, TimeoutError)
            error_message = self._format_error(f"{type(ve_rte_to).__name__}: {str(ve_rte_to)}", retryable=retryable)
            logger.error(f"Task {run_id} ({function_name}) failed pre-execution or during model acquisition: {error_message}")
            success = False
        except Exception as e:
            error_message = self._format_error(f"{type(e).__name__}: {str(e)}", retryable=False)
            logger.exception(f"Error executing function for run_id={run_id}: {error_message}")
            success = False
        finally:
            # Release the GPU if this thread set it busy
            if this_thread_set_gpu_busy:
                with self._gpu_busy_lock: # Lock to set self._is_gpu_busy
                    self._is_gpu_busy = False
                logger.info(f"Worker GPU marked as NOT BUSY by task {run_id} ({function_name}).")

            # Always send a result back, regardless of success, failure, or cancellation
            self._send_task_result(run_id, success, output_payload, error_message)
            # Remove from active tasks *after* sending result
            with self._active_tasks_lock:
                if run_id in self._active_tasks:
                    del self._active_tasks[run_id]
                resource_req = self._discovered_resources.get(function_name)
                func_limit = resource_req.max_concurrency if resource_req and resource_req.max_concurrency else 0
                if func_limit > 0:
                    current = self._active_function_counts.get(function_name, 0) - 1
                    if current <= 0:
                        self._active_function_counts.pop(function_name, None)
                    else:
                        self._active_function_counts[function_name] = current
                # else: # Might have been removed by stop() already
                     # logger.warning(f"Task {run_id} not found in active tasks during cleanup.")


    def _send_task_result(self, run_id: str, success: bool, output_payload: Optional[bytes], error_message: str) -> None:
        """Send a task execution result back to the scheduler via the queue."""
        try:
            result = pb.TaskExecutionResult(
                run_id=run_id,
                success=success,
                output_payload=(output_payload or b'') if success else b'', # Default to b'' if None
                error_message=error_message if not success else ""
            )
            msg = pb.WorkerSchedulerMessage(run_result=result)
            self._send_message(msg)
            logger.debug(f"Queued task result for run_id={run_id}, success={success}")
        except Exception as e:
             # This shouldn't generally fail unless message creation has issues
             logger.error(f"Failed to create or queue task result for run_id={run_id}: {e}")
