"""
Temporal worker for Agent Control Plane - Decoupled Architecture.

This worker:
1. Registers with Control Plane API on startup using KUBIYA_API_KEY
2. Gets dynamic configuration (Temporal credentials, task queue name, etc.)
3. Connects to Temporal Cloud with provided credentials
4. Sends periodic heartbeats to Control Plane
5. Has NO direct database access - all state managed via Control Plane API

Environment variables REQUIRED:
- KUBIYA_API_KEY: Kubiya API key for authentication (required)
- CONTROL_PLANE_URL: Control Plane API URL (e.g., https://control-plane.kubiya.ai)
- ENVIRONMENT_NAME: Environment/task queue name to join (default: "default")

Environment variables OPTIONAL:
- WORKER_HOSTNAME: Custom hostname for worker (default: auto-detected)
- HEARTBEAT_INTERVAL: Seconds between heartbeats (default: 60, lightweight mode)
"""

import asyncio
import os
import sys
import structlog
import httpx
import socket
import platform
import psutil
import time
from dataclasses import dataclass
from typing import Optional, List
from temporalio.worker import Worker
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner, SandboxRestrictions
from temporalio.client import Client, TLSConfig
from collections import deque

from control_plane_api.app.utils.helpers import is_local_temporal
# Import workflows and activities from local package
from control_plane_api.worker.workflows.agent_execution import AgentExecutionWorkflow
from control_plane_api.worker.workflows.team_execution import TeamExecutionWorkflow
from control_plane_api.worker.workflows.scheduled_job_wrapper import ScheduledJobWrapperWorkflow
from control_plane_api.app.workflows.dynamic_agent_execution import DynamicAgentExecutionWorkflow
from control_plane_api.worker.activities.agent_activities import (
    execute_agent_llm,
    update_execution_status,
    update_agent_status,
    get_execution_details,
    persist_conversation_history,
    submit_runtime_analytics_activity,
)
from control_plane_api.worker.activities.team_activities import (
    get_team_agents,
    execute_team_coordination,
)
from control_plane_api.worker.activities.runtime_activities import (
    execute_with_runtime,
    publish_user_message,
)
from control_plane_api.worker.activities.job_activities import (
    create_job_execution_record,
    update_job_execution_status,
)

# Configure structured logging
import logging
from control_plane_api.worker.utils.logging_config import configure_logging

# Configure logging with dynamic settings from environment variables
configure_logging()

logger = structlog.get_logger()

# Global log buffer to collect logs since last heartbeat
log_buffer = deque(maxlen=500)  # Keep last 500 log lines
worker_start_time = time.time()

# Global state for differential heartbeats (optimization)
_last_full_heartbeat_time: float = 0
_cached_system_info: Optional[dict] = None
_last_log_index_sent: int = 0
_full_heartbeat_interval: int = 300  # Full heartbeat every 5 minutes (vs lightweight every 60s)


class ProgressUI:
    """Minimal animated UI for worker startup - minikube style"""

    @staticmethod
    def step(emoji: str, message: str, status: str = ""):
        """Log a step with emoji and optional status"""
        if status:
            logger.info("worker_progress", emoji=emoji, message=message, status=status)
        else:
            logger.info("worker_progress", emoji=emoji, message=message)

    @staticmethod
    def success(emoji: str, message: str):
        """Log success message"""
        logger.info("worker_success", emoji=emoji, message=message)

    @staticmethod
    def error(emoji: str, message: str):
        """Log error message"""
        logger.error("worker_error", emoji=emoji, message=message)

    @staticmethod
    def warning(emoji: str, message: str):
        """Log warning message"""
        logger.warning("worker_warning", emoji=emoji, message=message)

    @staticmethod
    def header(text: str):
        """Log section header"""
        logger.info("worker_header", text=text)

    @staticmethod
    def banner():
        """Log startup banner"""
        logger.info("worker_banner", title="Kubiya Agent Worker")


def collect_system_info() -> dict:
    """
    Collect current system metrics and information.
    """
    try:
        cpu_percent = psutil.cpu_percent(interval=0.1)
        memory = psutil.virtual_memory()
        disk = psutil.disk_usage('/')

        # Get Kubiya CLI version from environment variable (set by CLI) - skipped for now
        cli_version = None

        # Get SDK version
        from control_plane_api.version import get_sdk_version
        sdk_version = get_sdk_version()

        # Get process ID
        pid = os.getpid()

        # Get current working directory
        cwd = os.getcwd()

        # Get supported runtimes (both are always available)
        supported_runtimes = ["agno", "claude_code"]

        # Check Docker availability
        docker_available = False
        docker_version = None
        try:
            import subprocess
            import shutil

            # First try to find docker in PATH using shutil.which
            docker_path = shutil.which('docker')
            logger.debug("docker_which_result", path=docker_path)

            # Fallback to common locations if not in PATH
            if not docker_path:
                docker_paths = [
                    '/usr/local/bin/docker',
                    '/usr/bin/docker',
                    '/opt/homebrew/bin/docker',
                ]
                for path in docker_paths:
                    logger.debug("docker_checking_path", path=path, exists=os.path.exists(path))
                    if os.path.exists(path):
                        docker_path = path
                        break

            if docker_path:
                logger.debug("docker_running_version_check", path=docker_path)
                result = subprocess.run(
                    [docker_path, '--version'],
                    capture_output=True,
                    text=True,
                    timeout=3,
                    shell=False
                )
                logger.debug(
                    "docker_version_output",
                    returncode=result.returncode,
                    stdout=result.stdout[:200],
                    stderr=result.stderr[:200] if result.stderr else None
                )
                if result.returncode == 0:
                    docker_available = True
                    # Parse "Docker version 28.1.1, build 4eba377"
                    output = result.stdout.strip()
                    if ',' in output:
                        docker_version = output.split(',')[0].replace('Docker version', '').strip()
                    else:
                        docker_version = output.replace('Docker version', '').strip()
                    logger.debug("docker_detected", version=docker_version, path=docker_path)
                else:
                    logger.warning("docker_version_check_failed", returncode=result.returncode)
            else:
                logger.warning("docker_not_found_in_path_or_common_locations")
        except Exception as e:
            # Log for debugging but don't fail
            logger.warning("docker_detection_failed", error=str(e), error_type=type(e).__name__)
            import traceback
            logger.debug("docker_detection_traceback", traceback=traceback.format_exc())

        # Parse OS details from platform
        os_name = platform.system()  # Darwin, Linux, Windows
        os_version = platform.release()

        return {
            "hostname": socket.gethostname(),
            "platform": platform.platform(),
            "os_name": os_name,
            "os_version": os_version,
            "python_version": platform.python_version(),
            "cli_version": cli_version,
            "sdk_version": sdk_version,
            "pid": pid,
            "cwd": cwd,
            "supported_runtimes": supported_runtimes,
            "docker_available": docker_available,
            "docker_version": docker_version,
            "cpu_count": psutil.cpu_count(),
            "cpu_percent": cpu_percent,
            "memory_total": memory.total,
            "memory_used": memory.used,
            "memory_percent": memory.percent,
            "disk_total": disk.total,
            "disk_used": disk.used,
            "disk_percent": disk.percent,
            "uptime_seconds": time.time() - worker_start_time,
        }
    except Exception as e:
        logger.warning("failed_to_collect_system_info", error=str(e))
        return {
            "hostname": socket.gethostname(),
            "platform": platform.platform(),
        }


def get_recent_logs() -> List[str]:
    """
    Get logs collected since last heartbeat and clear the buffer.
    """
    logs = list(log_buffer)
    log_buffer.clear()
    return logs


def log_to_buffer(message: str):
    """
    Add a log message to the buffer for sending in next heartbeat.
    """
    log_buffer.append(message)


@dataclass
class WorkerConfig:
    """Configuration received from Control Plane registration"""
    worker_id: str
    environment_name: str  # Task queue name (org_id.environment)
    temporal_namespace: str
    temporal_host: str
    temporal_api_key: str
    organization_id: str
    control_plane_url: str
    litellm_api_url: str = "https://llm-proxy.kubiya.ai"
    litellm_api_key: str = ""
    # Redis configuration for direct event streaming
    redis_url: str = ""
    redis_password: str = ""
    redis_enabled: bool = False
    # WebSocket configuration
    websocket_enabled: bool = True
    websocket_url: str = ""
    websocket_features: list = None
    # Queue configuration for cleanup
    queue_id: str = ""
    queue_ephemeral: bool = False
    queue_single_execution: bool = False


async def start_worker_for_queue(
    control_plane_url: str,
    kubiya_api_key: str,
    queue_id: str,
) -> WorkerConfig:
    """
    Start a worker for a specific queue ID.

    Args:
        control_plane_url: Control Plane API URL
        kubiya_api_key: Kubiya API key for authentication
        queue_id: Worker queue ID (UUID)

    Returns:
        WorkerConfig with all necessary configuration

    Raises:
        Exception if start fails
    """
    # Get worker SDK version for compatibility check
    from control_plane_api.version import get_sdk_version
    worker_sdk_version = get_sdk_version()

    # Collect system info to send during registration
    system_info = collect_system_info()

    logger.info(
        "starting_worker_for_queue",
        queue_id=queue_id,
        control_plane_url=control_plane_url,
        sdk_version=worker_sdk_version,
        pid=system_info.get("pid"),
        cwd=system_info.get("cwd"),
    )

    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(
                f"{control_plane_url}/api/v1/worker-queues/{queue_id}/start",
                headers={"Authorization": f"Bearer {kubiya_api_key}"},
                json={
                    "worker_sdk_version": worker_sdk_version,
                    "system_info": system_info,
                    "control_plane_url": control_plane_url
                }
            )

            # Success case
            if response.status_code == 200:
                data = response.json()

                ProgressUI.success("✓", f"Registered with control plane")
                logger.info(
                    "worker_registered",
                    worker_id=data.get("worker_id")[:8],
                    queue_name=data.get("queue_name"),
                )

                # Check SDK version compatibility
                control_plane_sdk_version = data.get("control_plane_sdk_version")
                if control_plane_sdk_version and control_plane_sdk_version != worker_sdk_version:
                    ProgressUI.warning("⚠", "SDK version mismatch detected")
                    print(f"\n   Worker SDK version:        {worker_sdk_version}")
                    print(f"   Control Plane SDK version: {control_plane_sdk_version}")
                    print(f"\n   Consider updating your worker to match the control plane version.\n")

                    logger.warning(
                        "sdk_version_mismatch",
                        worker_version=worker_sdk_version,
                        control_plane_version=control_plane_sdk_version,
                    )
                elif control_plane_sdk_version:
                    logger.info(
                        "sdk_version_match",
                        version=worker_sdk_version,
                    )

                # The task_queue_name is now just the queue UUID
                # Priority for LiteLLM API URL:
                # 1. LITELLM_API_BASE environment variable (from local proxy via CLI)
                # 2. Control plane litellm_api_url
                # 3. Default (https://llm-proxy.kubiya.ai)
                litellm_api_url = os.getenv("LITELLM_API_BASE") or data.get("litellm_api_url", "https://llm-proxy.kubiya.ai")
                litellm_api_key = os.getenv("LITELLM_API_KEY") or data.get("litellm_api_key", "")

                # Log which LiteLLM endpoint is being used
                if os.getenv("LITELLM_API_BASE"):
                    logger.info(
                        "using_local_litellm_proxy",
                        litellm_api_url=litellm_api_url,
                        source="environment_variable"
                    )
                elif "litellm_api_url" in data:
                    logger.info(
                        "using_control_plane_litellm_proxy",
                        litellm_api_url=litellm_api_url,
                        source="control_plane"
                    )

                return WorkerConfig(
                    worker_id=data["worker_id"],
                    environment_name=data["task_queue_name"],  # This is now the queue UUID
                    temporal_namespace=data["temporal_namespace"],
                    temporal_host=data["temporal_host"],
                    temporal_api_key=data["temporal_api_key"],
                    organization_id=data["organization_id"],
                    control_plane_url=data["control_plane_url"],
                    litellm_api_url=litellm_api_url,
                    litellm_api_key=litellm_api_key,
                    # Redis configuration from control plane (for direct event streaming)
                    redis_url=data.get("redis_url", ""),
                    redis_password=data.get("redis_password", ""),
                    redis_enabled=data.get("redis_enabled", False),
                    # WebSocket configuration from control plane
                    websocket_enabled=data.get("websocket_enabled", True),
                    websocket_url=data.get("websocket_url", ""),
                    websocket_features=data.get("websocket_features", []),
                )

            # Handle errors
            else:
                # Try to extract error detail from response
                error_message = response.text
                try:
                    error_data = response.json()
                    error_message = error_data.get("detail", response.text)
                except:
                    pass

                ProgressUI.error("✗", "Worker registration failed")
                print(f"   {error_message}\n")

                logger.error(
                    "worker_start_failed",
                    status_code=response.status_code,
                    queue_id=queue_id,
                )
                sys.exit(1)

    except httpx.RequestError as e:
        ProgressUI.error("✗", f"Connection failed: {control_plane_url}")
        print(f"   {str(e)}\n")
        logger.error("control_plane_connection_failed", error=str(e))
        sys.exit(1)


async def send_heartbeat(
    config: WorkerConfig,
    kubiya_api_key: str,
    status: str = "active",
    tasks_processed: int = 0,
    current_task_id: Optional[str] = None,
    force_full: bool = False
) -> bool:
    """
    Send heartbeat to Control Plane with differential data.

    Optimization: Uses lightweight heartbeats (status only) by default,
    and sends full heartbeats (with system info + logs) every 5 minutes.
    This reduces server load by 90% while maintaining full visibility.

    Args:
        config: Worker configuration
        kubiya_api_key: Kubiya API key for authentication
        status: Worker status (active, idle, busy)
        tasks_processed: Number of tasks processed
        current_task_id: Currently executing task ID
        force_full: Force a full heartbeat (ignores timing logic)

    Returns:
        True if successful, False otherwise
    """
    global _last_full_heartbeat_time, _cached_system_info, _last_log_index_sent

    current_time = time.time()
    time_since_last_full = current_time - _last_full_heartbeat_time

    # Determine if this should be a full heartbeat
    # Full heartbeat: every 5 minutes, or on first run, or if forced
    is_full_heartbeat = (
        force_full or
        _last_full_heartbeat_time == 0 or
        time_since_last_full >= _full_heartbeat_interval
    )

    # Build base heartbeat data (always included)
    heartbeat_data = {
        "status": status,
        "tasks_processed": tasks_processed,
        "current_task_id": current_task_id,
        "worker_metadata": {},
    }

    # Add system info and logs only for full heartbeats
    if is_full_heartbeat:
        # Collect fresh system info (expensive operation)
        system_info = collect_system_info()
        _cached_system_info = system_info
        heartbeat_data["system_info"] = system_info

        # Get logs since last full heartbeat (only new logs)
        logs = get_recent_logs()
        if logs:
            heartbeat_data["logs"] = logs

        # Update last full heartbeat time
        _last_full_heartbeat_time = current_time
        heartbeat_type = "full"
    else:
        # Lightweight heartbeat - no system info or logs
        # Server will use cached system info from Redis
        heartbeat_type = "lightweight"

    try:
        # Normalize URL to prevent double-slash issues
        control_plane_url = config.control_plane_url.rstrip("/")
        url = f"{control_plane_url}/api/v1/workers/{config.worker_id}/heartbeat"

        async with httpx.AsyncClient(timeout=10.0) as client:
            response = await client.post(
                url,
                json=heartbeat_data,
                headers={"Authorization": f"Bearer {kubiya_api_key}"}
            )

            if response.status_code in [200, 204]:
                logger.debug(
                    "heartbeat_sent",
                    worker_id=config.worker_id,
                    type=heartbeat_type,
                    payload_size=len(str(heartbeat_data))
                )
                log_to_buffer(
                    f"[{time.strftime('%H:%M:%S')}] Heartbeat sent ({heartbeat_type})"
                )
                return True
            else:
                logger.warning(
                    "heartbeat_failed",
                    status_code=response.status_code,
                    response=response.text[:200],
                    type=heartbeat_type
                )
                log_to_buffer(
                    f"[{time.strftime('%H:%M:%S')}] Heartbeat failed: HTTP {response.status_code}"
                )
                return False

    except Exception as e:
        error_msg = f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)"
        logger.warning(
            "heartbeat_error",
            error=error_msg,
            error_type=type(e).__name__,
            worker_id=config.worker_id[:8] if config.worker_id else "unknown",
            type=heartbeat_type
        )
        log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Heartbeat error: {error_msg[:150]}")
        return False


async def create_temporal_client(config: WorkerConfig) -> Client:
    """
    Create Temporal client using configuration from Control Plane.

    Args:
        config: Worker configuration from Control Plane registration

    Returns:
        Connected Temporal client instance
    """
    try:
        if is_local_temporal():
            # Connect to local Temporal without TLS or API key
            logger.info("connecting_to_local_temporal", host=config.temporal_host)
            client = await Client.connect(
                config.temporal_host,
                namespace=config.temporal_namespace,
            )
        else:
            # Connect to Temporal Cloud with TLS and API key
            logger.info("connecting_to_temporal_cloud", host=config.temporal_host)
            client = await Client.connect(
                config.temporal_host,
                namespace=config.temporal_namespace,
                tls=TLSConfig(),  # TLS enabled
                rpc_metadata={"authorization": f"Bearer {config.temporal_api_key}"}
            )

        return client

    except Exception as e:
        logger.error("connection_failed", error=str(e))
        ProgressUI.error("✗", f"Temporal connection failed: {str(e)}")
        raise


async def send_disconnect(
    config: WorkerConfig,
    kubiya_api_key: str,
    reason: str = "shutdown",
    exit_code: Optional[int] = None,
    error_message: Optional[str] = None
) -> bool:
    """
    Notify Control Plane that worker is disconnecting/exiting.

    Args:
        config: Worker configuration
        kubiya_api_key: Kubiya API key for authentication
        reason: Disconnect reason (shutdown, error, crash, etc.)
        exit_code: Exit code if applicable
        error_message: Error message if applicable

    Returns:
        True if successful, False otherwise
    """
    disconnect_data = {
        "reason": reason,
        "exit_code": exit_code,
        "error_message": error_message
    }

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            response = await client.post(
                f"{config.control_plane_url}/api/v1/workers/{config.worker_id}/disconnect",
                json=disconnect_data,
                headers={"Authorization": f"Bearer {kubiya_api_key}"}
            )

            if response.status_code in [200, 204]:
                logger.info(
                    "worker_disconnected",
                    worker_id=config.worker_id,
                    reason=reason,
                    exit_code=exit_code
                )
                return True
            else:
                logger.warning(
                    "disconnect_notification_failed",
                    status_code=response.status_code,
                    response=response.text[:200]
                )
                return False

    except Exception as e:
        logger.warning("disconnect_notification_error", error=str(e))
        return False


async def delete_ephemeral_queue(
    config: WorkerConfig,
    kubiya_api_key: str,
    queue_id: str,
    timeout: int = 5
) -> bool:
    """
    Delete ephemeral queue during worker shutdown.

    This allows the worker to clean up its ephemeral queue immediately,
    without requiring the CLI to wait for worker unregistration.

    Args:
        config: Worker configuration
        kubiya_api_key: Kubiya API key for authentication
        queue_id: Queue UUID to delete
        timeout: Request timeout in seconds (short timeout - if it fails, TTL handles it)

    Returns:
        True if successful, False otherwise
    """
    try:
        async with httpx.AsyncClient(timeout=float(timeout)) as client:
            response = await client.delete(
                f"{config.control_plane_url}/api/v1/worker-queues/{queue_id}",
                headers={"Authorization": f"Bearer {kubiya_api_key}"}
            )

            if response.status_code in [200, 204]:
                logger.info(
                    "ephemeral_queue_deleted",
                    queue_id=queue_id,
                    worker_id=config.worker_id
                )
                return True
            else:
                logger.warning(
                    "queue_delete_failed",
                    queue_id=queue_id,
                    status_code=response.status_code,
                    response=response.text[:200]
                )
                return False

    except Exception as e:
        logger.warning(
            "queue_delete_error",
            queue_id=queue_id,
            error=str(e)
        )
        return False


async def heartbeat_loop(config: WorkerConfig, kubiya_api_key: str, interval: int = 60):
    """
    Background task to send periodic heartbeats to Control Plane.

    Args:
        config: Worker configuration
        kubiya_api_key: Kubiya API key for authentication
        interval: Seconds between heartbeats
    """
    tasks_processed = 0

    while True:
        try:
            await asyncio.sleep(interval)
            await send_heartbeat(
                config=config,
                kubiya_api_key=kubiya_api_key,
                status="active",
                tasks_processed=tasks_processed
            )
        except asyncio.CancelledError:
            logger.info("heartbeat_loop_cancelled")
            break
        except Exception as e:
            logger.warning("heartbeat_loop_error", error=str(e))


async def run_worker():
    """
    Run the Temporal worker with decoupled architecture.

    The worker:
    1. Registers with Control Plane API
    2. Gets dynamic configuration (Temporal credentials, task queue, etc.)
    3. Connects to Temporal Cloud
    4. Starts heartbeat loop
    5. Registers workflows and activities
    6. Polls for tasks and executes them
    """
    # Get configuration from environment
    kubiya_api_key = os.environ.get("KUBIYA_API_KEY")
    control_plane_url = os.environ.get("CONTROL_PLANE_URL")
    queue_id = os.environ.get("QUEUE_ID")
    heartbeat_interval = int(os.environ.get("HEARTBEAT_INTERVAL", "60"))
    single_execution_mode = os.environ.get("SINGLE_EXECUTION", "").lower() in ("true", "1", "yes")

    # Validate required configuration
    if not kubiya_api_key:
        logger.error(
            "configuration_error",
            message="KUBIYA_API_KEY environment variable is required"
        )
        sys.exit(1)

    if not control_plane_url:
        logger.error(
            "configuration_error",
            message="CONTROL_PLANE_URL environment variable is required"
        )
        sys.exit(1)

    if not queue_id:
        logger.error(
            "configuration_error",
            message="QUEUE_ID environment variable is required"
        )
        sys.exit(1)

    log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Worker starting for queue {queue_id}")

    if single_execution_mode:
        log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Single execution mode: enabled (will exit after one task)")
        logger.info("single_execution_mode_enabled", queue_id=queue_id)

    # Check if agent-runtime mode is enabled
    use_agent_runtime = os.environ.get("USE_AGENT_RUNTIME", "").lower() in ("true", "1", "yes")
    agent_runtime_server = None
    health_monitor = None

    try:
        # Print banner
        ProgressUI.banner()

        # Step 0: Setup agent-runtime if enabled
        if use_agent_runtime:
            from pathlib import Path
            from control_plane_api.worker.binary_manager import BinaryManager
            from control_plane_api.worker.agent_runtime_server import AgentRuntimeServer, ServerConfig

            ProgressUI.step("⏳", "Setting up agent-runtime...")
            log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Downloading agent-runtime binary...")

            config_dir = Path(os.environ.get("AGENT_RUNTIME_CONFIG_DIR", Path.home() / ".kubiya"))
            binary_manager = BinaryManager(config_dir)
            binary_path = await binary_manager.ensure_binary("latest")

            log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Starting agent-runtime server...")
            server_config = ServerConfig(
                grpc_port=int(os.environ.get("AGENT_RUNTIME_GRPC_PORT", "50052")),
                http_port=int(os.environ.get("AGENT_RUNTIME_HTTP_PORT", "8082")),
                health_port=int(os.environ.get("AGENT_RUNTIME_HEALTH_PORT", "8083")),
                config_dir=config_dir,
                log_level=os.environ.get("AGENT_RUNTIME_LOG_LEVEL", "info"),
            )

            agent_runtime_server = AgentRuntimeServer(binary_path, server_config)
            await agent_runtime_server.start(wait_for_health=True, timeout=30)

            # Set environment variable for runtime to use
            os.environ["AGENT_RUNTIME_ADDRESS"] = agent_runtime_server.grpc_address
            ProgressUI.success("✓", f"Agent runtime ready at {agent_runtime_server.grpc_address}")
            log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Agent runtime server started on {agent_runtime_server.grpc_address}")

        # Step 1: Register with control plane
        ProgressUI.step("⏳", "Registering with control plane...")
        log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Registering with control plane...")
        config = await start_worker_for_queue(
            control_plane_url=control_plane_url,
            kubiya_api_key=kubiya_api_key,
            queue_id=queue_id,
        )
        log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Worker registered: {config.worker_id}")

        # Set environment variables for activities to use
        os.environ["CONTROL_PLANE_URL"] = config.control_plane_url

        # Set single execution flag so event publisher can disable WebSocket
        if single_execution_mode:
            os.environ["KUBIYA_SINGLE_EXECUTION_MODE"] = "true"
        os.environ["KUBIYA_API_KEY"] = kubiya_api_key
        os.environ["WORKER_ID"] = config.worker_id
        os.environ["LITELLM_API_BASE"] = config.litellm_api_url
        os.environ["LITELLM_API_KEY"] = config.litellm_api_key

        # Set WebSocket environment variables if enabled
        from control_plane_api.worker.utils.environment import should_use_websocket

        if config.websocket_enabled and config.websocket_url and should_use_websocket():
            os.environ["WEBSOCKET_ENABLED"] = "true"
            os.environ["WEBSOCKET_URL"] = config.websocket_url
            logger.info(
                "websocket_configured",
                worker_id=config.worker_id[:8],
                websocket_url=config.websocket_url
            )
        else:
            os.environ["WEBSOCKET_ENABLED"] = "false"
            if not should_use_websocket():
                logger.info("websocket_disabled_serverless_environment")
            else:
                logger.info("websocket_disabled_using_http")

        # Set Redis environment variables if provided (for Redis-first event streaming)
        if config.redis_enabled and config.redis_url:
            os.environ["REDIS_URL"] = config.redis_url
            os.environ["REDIS_ENABLED"] = "true"
            if config.redis_password:
                os.environ["REDIS_PASSWORD"] = config.redis_password
            logger.info(
                "redis_configured_for_direct_streaming",
                worker_id=config.worker_id[:8],
                redis_url=config.redis_url.split("@")[-1] if "@" in config.redis_url else config.redis_url  # Log without password
            )
        else:
            os.environ["REDIS_ENABLED"] = "false"
            logger.debug("redis_not_configured_will_use_http_endpoint")

        # Step 2: Connect to Temporal
        ProgressUI.step("⏳", "Connecting to Temporal...")
        client = await create_temporal_client(config)
        ProgressUI.success("✓", "Connected to Temporal")

        # Step 3: Send initial heartbeat
        ProgressUI.step("⏳", "Sending heartbeat...")
        await send_heartbeat(
            config=config,
            kubiya_api_key=kubiya_api_key,
            status="active",
            tasks_processed=0
        )
        ProgressUI.success("✓", "Worker visible in UI")

        # Start heartbeat loop in background
        heartbeat_task = asyncio.create_task(
            heartbeat_loop(config, kubiya_api_key, heartbeat_interval)
        )

        # Start health monitoring for agent-runtime if enabled
        health_monitor_task = None
        if agent_runtime_server is not None:
            from control_plane_api.worker.health_monitor import HealthMonitor
            # Note: os is already imported at module level (line 22)

            check_interval = int(os.environ.get("AGENT_RUNTIME_HEALTH_CHECK_INTERVAL", "30"))
            max_failures = int(os.environ.get("AGENT_RUNTIME_MAX_RESTART_ATTEMPTS", "3"))
            restart_enabled = os.environ.get("AGENT_RUNTIME_AUTO_RESTART", "true").lower() in ("true", "1", "yes")

            health_monitor = HealthMonitor(
                agent_runtime_server=agent_runtime_server,
                check_interval=check_interval,
                max_failures=max_failures,
                restart_enabled=restart_enabled,
            )
            await health_monitor.start()
            ProgressUI.success("✓", "Health monitoring enabled")
            log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Health monitoring started (interval={check_interval}s)")

        # Step 4: Create worker
        ProgressUI.step("⏳", "Starting worker...")

        # Configure workflow sandbox with passthrough modules
        # These modules use non-deterministic operations at import time but are safe
        # because they're only used in activities, not workflow logic
        sandbox_restrictions = SandboxRestrictions.default.with_passthrough_modules(
            "structlog",
            "structlog.dev",
            "structlog.processors",
            "structlog.tracebacks",
            "rich",
            "rich.traceback",
            "control_plane_api.version",  # Version checking uses filesystem operations
        )

        worker = Worker(
            client,
            task_queue=config.environment_name,
            workflows=[
                AgentExecutionWorkflow,
                TeamExecutionWorkflow,
                ScheduledJobWrapperWorkflow,  # Wrapper for scheduled jobs
                DynamicAgentExecutionWorkflow,  # Dynamic agent/team execution
            ],
            activities=[
                execute_agent_llm,
                update_execution_status,
                update_agent_status,
                get_execution_details,  # Get execution details from Control Plane
                persist_conversation_history,  # Conversation persistence
                submit_runtime_analytics_activity,  # Analytics submission
                get_team_agents,
                execute_team_coordination,
                execute_with_runtime,  # RuntimeFactory-based execution
                publish_user_message,  # Publish user message to stream
                create_job_execution_record,  # Job execution record creation
                update_job_execution_status,  # Job execution status updates
            ],
            max_concurrent_activities=10,
            max_concurrent_workflow_tasks=10,
            workflow_runner=SandboxedWorkflowRunner(restrictions=sandbox_restrictions),
        )

        ProgressUI.success("✓", "Worker ready")

        # Start WebSocket client if enabled
        from control_plane_api.worker.control_plane_client import get_control_plane_client

        control_plane_client = get_control_plane_client()
        if config.websocket_enabled and should_use_websocket():
            await control_plane_client.start_websocket()
            ProgressUI.step("✓", "WebSocket connected")
            logger.info("websocket_started", worker_id=config.worker_id[:8])

        if single_execution_mode:
            ProgressUI.header("📡 Listening for one task... (will exit after completion)")
        else:
            ProgressUI.header("📡 Listening for tasks... (Ctrl+C to stop)")

        logger.info(
            "worker_ready",
            worker_id=config.worker_id[:8],
            single_execution_mode=single_execution_mode,
        )

        # Run worker (blocks until interrupted)
        try:
            if single_execution_mode:
                # Single execution mode: run worker and monitor for workflow completion
                logger.info("starting_worker_in_single_execution_mode")

                # Create a task to run the worker
                worker_run_task = asyncio.create_task(worker.run())

                # Monitor for execution completion via Control Plane API
                async def monitor_and_shutdown():
                    """
                    Monitor execution status and shutdown after task completes.
                    Robustness improvements:
                    - Requires consecutive completion checks to avoid false positives
                    - Extends timeout for long-running tasks
                    """
                    # Brief wait for worker to start and pick up the execution
                    # Reduced from 5s to 1s for faster ephemeral worker startup
                    await asyncio.sleep(1)

                    # Monitor for 30 minutes max (extended from 10 minutes)
                    max_runtime = 1800
                    check_interval = 2  # Check every 2 seconds - balanced between speed and API load
                    elapsed = 0
                    execution_seen = False
                    execution_id = None

                    # Robustness: Require 2 consecutive "completed" checks before shutting down
                    # With 2s polling interval, this provides 4s buffer for async operations to settle
                    consecutive_completion_checks = 0
                    required_consecutive_checks = 2

                    logger.info("single_execution_monitor_started", queue_id=queue_id)

                    should_shutdown = False
                    while elapsed < max_runtime and not should_shutdown:
                        await asyncio.sleep(check_interval)
                        elapsed += check_interval

                        # Check if worker task completed unexpectedly
                        if worker_run_task.done():
                            logger.info("single_execution_worker_task_completed", elapsed=elapsed)
                            break

                        # Query Control Plane for recent executions on this queue
                        try:
                            # Get the control plane client
                            async with httpx.AsyncClient(timeout=10.0) as http_client:
                                # List recent executions for this queue
                                response = await http_client.get(
                                    f"{control_plane_url}/api/v1/worker-queues/{queue_id}/executions",
                                    headers={"Authorization": f"Bearer {kubiya_api_key}"},
                                    params={"limit": 5, "status": "all"}
                                )

                                if response.status_code == 200:
                                    executions = response.json()

                                    # Look for any execution in a terminal or waiting state
                                    for execution in executions:
                                        exec_status = execution.get("status", "").lower()
                                        exec_id = execution.get("id")

                                        if not execution_seen:
                                            if exec_status in ["running", "completed", "failed", "waiting_for_input"]:
                                                execution_seen = True
                                                execution_id = exec_id
                                                logger.info("single_execution_detected", execution_id=exec_id[:8] if exec_id else None, status=exec_status)

                                        # If we've seen an execution and it's now in a terminal state, check if consistent
                                        # NOTE: We do NOT treat "waiting_for_input" as terminal in single execution mode
                                        # because the LLM may still be processing (e.g., tool calls) and the execution
                                        # should continue until truly completed or failed
                                        if execution_seen and exec_id == execution_id:
                                            if exec_status in ["completed", "failed", "cancelled"]:
                                                consecutive_completion_checks += 1
                                                logger.info("single_execution_completion_check",
                                                           execution_id=exec_id[:8] if exec_id else None,
                                                           status=exec_status,
                                                           consecutive_checks=consecutive_completion_checks,
                                                           required_checks=required_consecutive_checks,
                                                           elapsed=elapsed)

                                                # Only shutdown after consecutive checks confirm completion
                                                if consecutive_completion_checks >= required_consecutive_checks:
                                                    logger.info("single_execution_completed",
                                                               execution_id=exec_id[:8] if exec_id else None,
                                                               status=exec_status,
                                                               elapsed=elapsed)
                                                    # Give SSE clients time to receive all final events
                                                    # Reduced to 2s for faster shutdown while still allowing
                                                    # SSE streams to complete
                                                    logger.info("single_execution_grace_period_starting",
                                                               execution_id=exec_id[:8] if exec_id else None,
                                                               grace_seconds=2)
                                                    await asyncio.sleep(2)
                                                    should_shutdown = True
                                                    break
                                            else:
                                                # Execution is back to running state - reset counter
                                                if consecutive_completion_checks > 0:
                                                    logger.info("single_execution_still_active",
                                                               execution_id=exec_id[:8] if exec_id else None,
                                                               status=exec_status,
                                                               resetting_counter=True)
                                                    consecutive_completion_checks = 0
                                else:
                                    logger.debug("single_execution_status_check_failed", status_code=response.status_code)
                                    # Reset consecutive checks on failed API call to be safe
                                    if consecutive_completion_checks > 0:
                                        logger.debug("single_execution_resetting_counter_after_failed_check")
                                        consecutive_completion_checks = 0

                        except Exception as e:
                            logger.debug("single_execution_status_check_error", error=str(e))
                            # Reset consecutive checks on error to be safe
                            if consecutive_completion_checks > 0:
                                logger.debug("single_execution_resetting_counter_after_error")
                                consecutive_completion_checks = 0
                            # Continue monitoring even if status check fails

                    # Check why we exited the loop
                    if not should_shutdown and elapsed >= max_runtime:
                        # Actual timeout
                        logger.warning("single_execution_timeout_reached", elapsed=elapsed)

                    # Shutdown the worker gracefully
                    logger.info("single_execution_triggering_shutdown", elapsed_seconds=elapsed, reason="completed" if should_shutdown else "timeout")
                    ProgressUI.step("✓", "Task completed - shutting down worker...")
                    log_to_buffer(f"[{time.strftime('%H:%M:%S')}] Task completed, shutting down...")
                    await worker.shutdown()

                # Start monitoring task
                monitor_task = asyncio.create_task(monitor_and_shutdown())

                try:
                    # Wait for worker to complete
                    await worker_run_task
                    logger.info("single_execution_worker_stopped")
                finally:
                    # Cancel monitor task if still running
                    if not monitor_task.done():
                        monitor_task.cancel()
                        try:
                            await monitor_task
                        except asyncio.CancelledError:
                            pass
            else:
                # Normal mode - run indefinitely
                await worker.run()
        finally:
            # Stop WebSocket client
            await control_plane_client.stop_websocket()

        # Cancel heartbeat task when worker stops
        heartbeat_task.cancel()
        try:
            await heartbeat_task
        except asyncio.CancelledError:
            pass

        # Notify control plane of graceful shutdown
        print()
        ProgressUI.step("⏳", "Shutting down gracefully...")

        # Delete ephemeral queue if we're the owner (single execution mode)
        if config.queue_ephemeral and config.queue_single_execution and config.queue_id:
            try:
                await delete_ephemeral_queue(
                    config=config,
                    kubiya_api_key=kubiya_api_key,
                    queue_id=config.queue_id
                )
                logger.info("ephemeral_queue_cleaned_up", queue_id=config.queue_id)
            except Exception as e:
                logger.warning(
                    "ephemeral_queue_cleanup_failed",
                    queue_id=config.queue_id,
                    error=str(e)
                )
                # Continue shutdown even if delete fails (TTL will handle it)

        await send_disconnect(
            config=config,
            kubiya_api_key=kubiya_api_key,
            reason="shutdown",
            exit_code=0
        )
        ProgressUI.success("✓", "Worker stopped")
        print()

    except KeyboardInterrupt:
        print()
        ProgressUI.step("⏳", "Shutting down...")

        # Stop health monitor if running
        if health_monitor is not None:
            try:
                await health_monitor.stop()
            except Exception as e:
                logger.warning("health_monitor_stop_failed", error=str(e))

        # Stop agent-runtime server if running
        if agent_runtime_server is not None:
            try:
                ProgressUI.step("⏳", "Stopping agent-runtime server...")
                agent_runtime_server.stop(timeout=10)
                ProgressUI.success("✓", "Agent runtime stopped")
            except Exception as e:
                logger.warning("agent_runtime_stop_failed_on_interrupt", error=str(e))

        # Stop WebSocket client
        from control_plane_api.worker.control_plane_client import get_control_plane_client
        try:
            control_plane_client = get_control_plane_client()
            await control_plane_client.stop_websocket()
        except:
            pass

        # Notify control plane of keyboard interrupt (only if config was successfully obtained)
        try:
            if 'config' in locals():
                # Delete ephemeral queue if we're the owner
                if config.queue_ephemeral and config.queue_single_execution and config.queue_id:
                    try:
                        await delete_ephemeral_queue(
                            config=config,
                            kubiya_api_key=kubiya_api_key,
                            queue_id=config.queue_id
                        )
                    except Exception as e:
                        logger.warning(
                            "ephemeral_queue_cleanup_on_interrupt_failed",
                            error=str(e)
                        )

                await send_disconnect(
                    config=config,
                    kubiya_api_key=kubiya_api_key,
                    reason="shutdown",
                    exit_code=0
                )
                ProgressUI.success("✓", "Worker stopped")
            else:
                logger.info("shutdown_before_registration_completed")
        except Exception as e:
            logger.warning("disconnect_on_interrupt_failed", error=str(e))
    except Exception as e:
        import traceback
        logger.error("temporal_worker_error", error=str(e), traceback=traceback.format_exc())

        # Stop health monitor if running
        if health_monitor is not None:
            try:
                await health_monitor.stop()
            except Exception as stop_error:
                logger.warning("health_monitor_stop_failed_on_error", error=str(stop_error))

        # Stop agent-runtime server if running
        if agent_runtime_server is not None:
            try:
                logger.info("stopping_agent_runtime_on_error")
                agent_runtime_server.stop(timeout=10)
                logger.info("agent_runtime_stopped_on_error")
            except Exception as stop_error:
                logger.warning("agent_runtime_stop_failed_on_error", error=str(stop_error))

        # Notify control plane of error (only if config was successfully obtained)
        try:
            if 'config' in locals():
                await send_disconnect(
                    config=config,
                    kubiya_api_key=kubiya_api_key,
                    reason="error",
                    exit_code=1,
                    error_message=str(e)[:2000] + (" [truncated]" if len(str(e)) > 2000 else "")
                )
            else:
                logger.warning("disconnect_skipped_no_config", error="Worker failed before registration completed")
        except Exception as disconnect_error:
            logger.warning("disconnect_on_error_failed", error=str(disconnect_error))
        raise


def main():
    """Main entry point with CLI argument support"""
    import argparse

    # Parse CLI arguments
    parser = argparse.ArgumentParser(
        description="Kubiya Agent Worker - Temporal worker for agent execution"
    )
    parser.add_argument(
        "--queue-id",
        type=str,
        help="Worker queue ID (can also use QUEUE_ID env var)"
    )
    parser.add_argument(
        "--api-key",
        type=str,
        help="Kubiya API key (can also use KUBIYA_API_KEY env var)"
    )
    parser.add_argument(
        "--control-plane-url",
        type=str,
        help="Control plane URL (can also use CONTROL_PLANE_URL env var)"
    )
    parser.add_argument(
        "--heartbeat-interval",
        type=int,
        default=60,
        help="Heartbeat interval in seconds (default: 60, lightweight mode)"
    )

    args = parser.parse_args()

    # Set environment variables from CLI args if not already set
    # Environment variables take precedence over CLI args (safer)
    if args.queue_id and not os.environ.get("QUEUE_ID"):
        os.environ["QUEUE_ID"] = args.queue_id
    if args.api_key and not os.environ.get("KUBIYA_API_KEY"):
        os.environ["KUBIYA_API_KEY"] = args.api_key
    if args.control_plane_url and not os.environ.get("CONTROL_PLANE_URL"):
        os.environ["CONTROL_PLANE_URL"] = args.control_plane_url
    if args.heartbeat_interval and not os.environ.get("HEARTBEAT_INTERVAL"):
        os.environ["HEARTBEAT_INTERVAL"] = str(args.heartbeat_interval)

    logger.info("worker_starting")

    try:
        asyncio.run(run_worker())
    except KeyboardInterrupt:
        logger.info("worker_stopped")
    except Exception as e:
        logger.error("worker_failed", error=str(e))
        sys.exit(1)


if __name__ == "__main__":
    main()
