import asyncio
import json
import logging
import random
import socket
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

import numpy as np
import ray
import requests
import torch
from flask import abort, Flask, request, Response
from tensordict import TensorDict

from verl import DataProto

from .lightningrl import LightningRL
from .types import Triplet, Rollout

# Configure logging
from .logging import configure_logger


logger = configure_logger(level=logging.INFO, name=__name__)


def _find_available_port() -> int:
    """Find an available port for the proxy server."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", 0))
        return s.getsockname()[1]


class TaskStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    RETRYING = "retrying"


@dataclass
class AgentTask:
    task_id: str
    agent_type: str
    task_data: Dict[str, Any]
    endpoint: str
    status: TaskStatus
    future: Any = None
    result: Any = None
    error: Optional[str] = None
    created_at: float = None
    completed_at: Optional[float] = None
    retry_count: int = 0
    max_retries: int = 3
    timeout_seconds: float = 60.0


@dataclass
class RetryConfig:
    """
    Configuration for task retry behavior.
    
    Note: The retry system only retries tasks that fail due to timeout.
    Other types of failures (like agent logic errors, invalid input, etc.)
    are not retried to avoid wasting resources on tasks that are likely to fail again.
    """

    max_retries: int = 3
    timeout_seconds: float = 60.0
    retry_delay_seconds: float = 1.0
    exponential_backoff: bool = True
    backoff_multiplier: float = 2.0


class TaskRetryManager:
    """Manages timeout and retry logic for agent tasks."""

    def __init__(self, retry_config: RetryConfig = None):
        self.config = retry_config or RetryConfig()
        self.logger = logging.getLogger(f"{__name__}.TaskRetryManager")

    def should_retry(self, task: AgentTask) -> bool:
        """
        Check if a task should be retried.
        
        Only retries tasks that failed due to timeout, not other types of failures.
        This ensures that only network/response timeouts are retried, while
        other errors (like invalid input, agent logic errors, etc.) are not.
        """
        # Only retry if we haven't exceeded max retries
        if task.retry_count >= task.max_retries:
            return False
        
        # Only retry if the failure was due to timeout
        if task.error and "timed out" in task.error.lower():
            return True
        
        return False

    def is_task_timed_out(self, task: AgentTask, current_time: float = None) -> bool:
        """Check if a task has exceeded its timeout."""
        if current_time is None:
            current_time = time.time()

        if task.status != TaskStatus.RUNNING or task.created_at is None:
            return False

        elapsed_time = current_time - task.created_at
        return elapsed_time > task.timeout_seconds

    def get_retry_delay(self, retry_count: int) -> float:
        """Calculate delay before retry based on configuration."""
        if not self.config.exponential_backoff:
            return self.config.retry_delay_seconds

        delay = self.config.retry_delay_seconds * (
            self.config.backoff_multiplier**retry_count
        )
        return min(delay, 30.0)  # Cap at 30 seconds

    def prepare_retry(self, task: AgentTask) -> AgentTask:
        """Prepare a task for retry by updating its state."""
        task.retry_count += 1
        task.status = TaskStatus.RETRYING
        task.future = None
        task.error = None

        self.logger.info(
            f"Preparing retry {task.retry_count}/{task.max_retries} for task {task.task_id}"
        )
        return task

    def mark_task_failed(self, task: AgentTask, error_message: str) -> AgentTask:
        """Mark a task as permanently failed after exhausting retries."""
        task.status = TaskStatus.FAILED
        task.error = f"Failed after {task.retry_count} retries: {error_message}"
        task.completed_at = time.time()

        self.logger.error(f"Task {task.task_id} permanently failed: {task.error}")
        return task


class AgentOrchestrator:
    def __init__(
        self,
        server_addresses=None,
        ray_config=None,
        tokenizer=None,
        pad_token_id=None,
        reward_fillna_value=0.0,
        proxy_port=39795,
        llm_timeout_seconds=600.0,
        collector_timeout_ms=100,

        retry_config=None,
        task_timeout_seconds=90.0,
        max_task_retries=3,
    ):
        """
        Initialize the Agent Orchestrator.

        Args:
            server_addresses: Optional list of backend LLM server addresses. If provided,
                            the proxy server will be started during initialization.
            ray_config: Optional Ray configuration dictionary for cluster initialization
            tokenizer: Optional tokenizer for processing text
            pad_token_id: Token ID to use for padding (defaults to 0)
            reward_fillna_value: Value to use when filling missing rewards (defaults to 0.0)
            proxy_port: Port for the proxy server (defaults to 39795)
            llm_timeout_seconds: Timeout for LLM requests in seconds (defaults to 600.0)
            collector_timeout_ms: Timeout for async collector in milliseconds (defaults to 100)

            retry_config: Optional RetryConfig object for customizing retry behavior
            task_timeout_seconds: Default timeout for agent tasks in seconds (defaults to 60.0)
            max_task_retries: Default maximum number of retries per task (defaults to 3)
        """
        self.agent_types: Dict[str, Callable] = {}
        self.tasks: Dict[str, AgentTask] = {}
        self.futures: Dict[str, Any] = {}
        self.ray_initialized = False

        self.tokenizer = tokenizer
        self.pad_token_id = pad_token_id if pad_token_id is not None else 0
        self.reward_fillna_value = reward_fillna_value
        self._completed_rollouts: Dict[str, Any] = {}
        self._task_id_to_original_sample: Dict[str, Dict] = {}
        self._total_tasks_queued = 0
        self.is_train = True

        # Initialize retry configuration
        if retry_config is None:
            retry_config = RetryConfig(
                max_retries=max_task_retries, timeout_seconds=task_timeout_seconds
            )
        self.retry_manager = TaskRetryManager(retry_config)
        self.default_task_timeout = task_timeout_seconds
        self.default_max_retries = max_task_retries
        self.rollout_traj_from_proxy: Dict[str, List[Tuple[List[int], List[int]]]] = {}

        # Async collector configuration
        self.collector_timeout_ms = collector_timeout_ms
        self._completion_queue: Optional[asyncio.Queue] = None
        self._collector_task: Optional[asyncio.Task] = None
        self._collector_running = False
        self._event_loop: Optional[asyncio.AbstractEventLoop] = None
        self._executor = ThreadPoolExecutor(
            max_workers=2, thread_name_prefix="orchestrator"
        )
        self._shutdown_event = threading.Event()

        # Initialize LightningRL component for training batch preparation
        self.lightning_rl = LightningRL(
            pad_token_id=self.pad_token_id, reward_fillna_value=self.reward_fillna_value
        )

        # Proxy server configuration
        self.proxy_port = proxy_port
        self.llm_timeout_seconds = llm_timeout_seconds
        self.backend_llm_server_addresses: List[str] = server_addresses or []
        self._proxy_thread: Optional[threading.Thread] = None
        self._proxy_server_running = False

        # Initialize Ray cluster
        self.init_ray(ray_config)

        # Start proxy server if server addresses are provided
        if self.backend_llm_server_addresses:
            logger.info("Starting proxy server during initialization...")
            self._start_proxy_server()

    def register_agent_type(
        self,
        agent_name: str,
        rollout_function: Callable,
        endpoint: str = "https://api.openai.com/v1",
        model: str = "gpt-4o-mini",
    ):
        """
        Register a new agent rollout function.

        Args:
            agent_name: Name identifier for the agent type
            rollout_function: Ray remote function for agent execution
            endpoint: Optional LLM endpoint URL
            model: Model name to use for the agent
        """
        self.agent_types[agent_name] = {
            "function": rollout_function,
            "endpoint": endpoint,
            "model": model,
        }
        logger.info(
            f"Registered agent type: {agent_name} with endpoint: {endpoint} and model: {model}"
        )

    def init_ray(self, ray_config: Optional[Dict[str, Any]] = None):
        """
        Initialize Ray cluster.

        Args:
            ray_config: Optional Ray configuration dictionary
        """
        if self.ray_initialized:
            logger.info("Ray is already initialized")
            return

        try:
            if not ray.is_initialized():
                if ray_config:
                    ray.init(**ray_config)
                else:
                    ray.init()
            self.ray_initialized = True
            logger.info("Ray initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize Ray: {e}")
            raise

    def init_server(
        self,
        endpoint: str = "https://api.openai.com/v1",
        agent_types: Optional[List[tuple]] = None,
    ):
        """
        Initialize the orchestrator server by registering agent types.

        Args:
            endpoint: Default LLM endpoint for agents
            agent_types: Optional list of (agent_name, rollout_function) tuples to register
        """
        # Register provided agent types
        if agent_types:
            for agent_name, rollout_function in agent_types:
                self.register_agent_type(agent_name, rollout_function, endpoint)

        logger.info(f"Orchestrator server initialized with endpoint: {endpoint}")

    def agent_rollout(
        self,
        agent_type: str,
        task: Dict[str, Any],
        endpoint: Optional[str] = None,
        timeout_seconds: Optional[float] = None,
        max_retries: Optional[int] = None,
    ) -> str:
        """
        Launch an agent task on Ray cluster with timeout and retry support.

        Args:
            agent_type: Type of agent to use
            task: Task data to process
            endpoint: Optional endpoint override
            timeout_seconds: Optional timeout override for this task
            max_retries: Optional max retries override for this task

        Returns:
            Task ID for tracking the task
        """
        if agent_type not in self.agent_types:
            raise ValueError(
                f"Agent type '{agent_type}' not registered. Available types: {list(self.agent_types.keys())}"
            )

        # Generate unique task ID
        task_id = str(uuid.uuid4())

        # Use provided endpoint or fall back to registered endpoint
        use_endpoint = endpoint or self.agent_types[agent_type]["endpoint"]
        use_model = self.agent_types[agent_type]["model"]

        # Create task record with retry configuration
        agent_task = AgentTask(
            task_id=task_id,
            agent_type=agent_type,
            task_data=task,
            endpoint=use_endpoint,
            status=TaskStatus.PENDING,
            created_at=time.time(),
            timeout_seconds=timeout_seconds or self.default_task_timeout,
            max_retries=max_retries or self.default_max_retries,
        )

        return self._launch_agent_task(agent_task)

    def _launch_agent_task(self, agent_task: AgentTask) -> str:
        """
        Internal method to launch or relaunch an agent task.

        Args:
            agent_task: AgentTask object to launch

        Returns:
            Task ID for tracking the task
        """
        try:
            # Launch the Ray task
            rollout_function = self.agent_types[agent_task.agent_type]["function"]

            future = rollout_function.remote(
                agent_task.task_data,
                agent_task.endpoint,
                self.agent_types[agent_task.agent_type]["model"],
                agent_task.task_id,
            )

            # Update task record
            agent_task.future = future
            agent_task.status = TaskStatus.RUNNING
            agent_task.created_at = time.time()  # Reset creation time for retries

            # Store task and future
            self.tasks[agent_task.task_id] = agent_task
            self.futures[agent_task.task_id] = future

            logger.info(
                f"Launched task {agent_task.task_id} (attempt {agent_task.retry_count + 1})"
            )
            return agent_task.task_id

        except Exception as e:
            error_msg = str(e)
            # Check if this is a Ray timeout exception
            if "timeout" in error_msg.lower() or "timed out" in error_msg.lower():
                error_msg = f"Task timed out: {error_msg}"
            
            agent_task.status = TaskStatus.FAILED
            agent_task.error = error_msg
            self.tasks[agent_task.task_id] = agent_task

            logger.error(f"Failed to launch task {agent_task.task_id}: {error_msg}")
            raise

    def _retry_agent_task(self, task_id: str, error_message: str) -> bool:
        """
        Attempt to retry a failed or timed-out task.

        Args:
            task_id: Task ID to retry
            error_message: Error message from the failure

        Returns:
            True if retry was initiated, False if task permanently failed
        """
        if task_id not in self.tasks:
            logger.error(f"Cannot retry unknown task {task_id}")
            return False

        task = self.tasks[task_id]

        # Check if we should retry this task
        if not self.retry_manager.should_retry(task):
            # Mark as permanently failed
            self.retry_manager.mark_task_failed(task, error_message)
            return False

        # Prepare task for retry
        self.retry_manager.prepare_retry(task)

        # Calculate retry delay
        retry_delay = self.retry_manager.get_retry_delay(
            task.retry_count - 1
        )  # retry_count was already incremented

        # Schedule retry after delay
        def delayed_retry():
            time.sleep(retry_delay)
            try:
                self._launch_agent_task(task)
                logger.info(f"Retrying task {task_id} after {retry_delay:.1f}s delay")
            except Exception as e:
                error_msg = str(e)
                # Check if this is a Ray timeout exception
                if "timeout" in error_msg.lower() or "timed out" in error_msg.lower():
                    error_msg = f"Task timed out: {error_msg}"
                logger.error(f"Failed to retry task {task_id}: {error_msg}")
                self._retry_agent_task(task_id, error_msg)

        # Execute retry in background thread to avoid blocking
        retry_thread = threading.Thread(target=delayed_retry, daemon=True)
        retry_thread.start()

        return True

    def _handle_task_timeout(self, task_id: str):
        """Handle a task that has timed out."""
        if task_id not in self.tasks:
            return

        task = self.tasks[task_id]

        # Cancel the Ray task if it's still running
        if task.future:
            try:
                ray.cancel(task.future)
                logger.info(f"Cancelled timed-out task {task_id}")
            except Exception as e:
                logger.warning(f"Failed to cancel task {task_id}: {e}")

        # Clean up future reference
        if task_id in self.futures:
            del self.futures[task_id]

        # Attempt to retry the task
        timeout_error = f"Task timed out after {task.timeout_seconds}s"
        retry_success = self._retry_agent_task(task_id, timeout_error)

        if not retry_success:
            logger.error(f"Task {task_id} permanently failed due to timeout")

    def get_completed_tasks(self) -> List[str]:
        """Get list of completed task IDs."""
        return [
            task_id
            for task_id, task in self.tasks.items()
            if task.status == TaskStatus.COMPLETED
        ]


    def wait_for_any_completion(self, timeout: Optional[float] = None) -> Optional[str]:
        """
        Wait for any running task to complete.

        Args:
            timeout: Optional timeout in seconds

        Returns:
            Task ID of completed task or None if timeout
        """
        running_futures = [
            (task_id, task.future)
            for task_id, task in self.tasks.items()
            if task.status == TaskStatus.RUNNING and task.future
        ]

        if not running_futures:
            return None

        futures = [future for _, future in running_futures]
        task_id_map = {id(future): task_id for task_id, future in running_futures}

        try:
            ready, not_ready = ray.wait(futures, num_returns=1, timeout=timeout)
            if ready:
                completed_future = ready[0]
                task_id = task_id_map[id(completed_future)]
                result = ray.get(completed_future)
                self._update_task_completion(task_id, result)
                return task_id
        except Exception as e:
            error_msg = str(e)
            # Check if this is a Ray timeout exception
            if "timeout" in error_msg.lower() or "timed out" in error_msg.lower():
                error_msg = f"Task timed out: {error_msg}"
            logger.error(f"Error waiting for task completion: {error_msg}")

        return None
    
    def shutdown(self):
        """Shutdown the orchestrator and Ray."""
        logger.info("Shutting down orchestrator...")

        # Signal shutdown to async collector
        self._shutdown_event.set()

        # Stop async collector
        self._stop_collector()

        # Stop proxy server
        if self._proxy_server_running:
            self._stop_proxy_server()

        # Shutdown executor
        self._executor.shutdown(wait=True)

        if self.ray_initialized:
            ray.shutdown()
            self.ray_initialized = False
        logger.info("Orchestrator shutdown complete")

    def _create_proxy_app(self) -> Flask:
        """Create and configure the Flask proxy application."""
        app = Flask(__name__)
        
        @app.route(
            "/v1/<path:path>",
            methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
        )
        def proxy_handler(path):
            return self._handle_proxy_request(path)

        return app

    def _handle_proxy_request(self, path: str) -> Response:
        """Handle individual proxy requests with proper error handling."""
        # Check if backends are available
        if not self.backend_llm_server_addresses:
            abort(503, description="No backend LLM servers available")

        # Select target server using round-robin or random selection
        target_server = random.choice(self.backend_llm_server_addresses)
        target_url = f"http://{target_server}/v1/{path}"

        try:
            # Forward request to backend
            response, task_id = self._forward_request(target_url)
            return self._process_response(response, task_id)

        except requests.exceptions.RequestException as e:
            logger.error(f"Proxy request failed: {e}")
            abort(500, description=f"Error proxying request: {e}")

    def _forward_request(self, target_url: str) -> requests.Response:
        """Forward the incoming request to the target backend server."""
        # Prepare headers, excluding hop-by-hop headers
        headers = {
            key: value
            for key, value in request.headers
            if key.lower() not in {"host", "connection"}
        } 
        # FIXME: consider using other cleaner way to get task_id
        task_id = headers['Authorization'].split(' ')[1]
        # The request object is automatically provided by Flask and is available in the context of Flask route handlers. 
        return requests.request(
            method=request.method,
            url=target_url,
            headers=headers,
            params=request.args,
            data=request.get_data(), 
            cookies=request.cookies,
            allow_redirects=False,
            timeout=self.llm_timeout_seconds,
        ), task_id

    def _process_response(self, response: requests.Response, task_id: str) -> Response:
        """Process the backend response and return it to the client."""
        # Filter out hop-by-hop headers
        hop_by_hop_headers = {
            "content-encoding",
            "content-length",
            "transfer-encoding",
            "connection",
            "keep-alive",
            "proxy-authenticate",
            "proxy-authorization",
            "te",
            "trailers",
            "upgrade",
        }

        response_headers = [
            (name, value)
            for name, value in response.raw.headers.items()
            if name.lower() not in hop_by_hop_headers
        ]

        # Handle successful responses with potential JSON processing
        if response.status_code == 200:
            try:
                # Process JSON response if needed (placeholder for future token processing)
                response_json = json.loads(response.content.decode("utf-8"))
                processed_content = json.dumps(response_json).encode("utf-8")

                if task_id not in self.rollout_traj_from_proxy:
                    self.rollout_traj_from_proxy[task_id] = []
                # FIXME: consider when there are multiple response_token_ids 
                self.rollout_traj_from_proxy[task_id].append((response_json['response_token_ids'][0], response_json["prompt_token_ids"]))
                 
                return Response(
                    processed_content,
                    status=response.status_code,
                    headers=response_headers,
                )
            except (json.JSONDecodeError, UnicodeDecodeError):
                # Return raw content if JSON processing fails
                pass

        return Response(response.content, response.status_code, response_headers)

    def _start_proxy_server(self):
        """Start the Flask proxy server in a separate thread."""
        if self._proxy_server_running:
            logger.warning("Proxy server is already running")
            return

        app = self._create_proxy_app()

        def run_server():
            app.run(host="0.0.0.0", port=self.proxy_port, threaded=True, debug=False)

        self._proxy_thread = threading.Thread(target=run_server, daemon=True)
        self._proxy_thread.start()
        self._proxy_server_running = True

        logger.info(f"Proxy server running on http://localhost:{self.proxy_port}/v1")

    def _stop_proxy_server(self):
        """Stop the proxy server."""
        if self._proxy_server_running:
            self._proxy_server_running = False
            logger.info("Proxy server stopped")

    def start_proxy_server(
        self, backend_llm_server_addresses: Optional[List[str]] = None
    ):
        """
        Start the proxy server with backend LLM server addresses.

        Args:
            backend_llm_server_addresses: Optional list of backend server addresses to proxy to.
                                        If not provided, uses the instance's stored addresses.
        """
        if backend_llm_server_addresses is not None:
            self.backend_llm_server_addresses = backend_llm_server_addresses

        if not self.backend_llm_server_addresses:
            logger.warning(
                "No backend LLM server addresses provided. Proxy server will not be started."
            )
            return

        self._start_proxy_server()
        logger.info(
            f"Proxy server started with {len(self.backend_llm_server_addresses)} backends"
        )

    def get_proxy_endpoint(self) -> str:
        """Get the proxy server endpoint URL."""
        return f"http://localhost:{self.proxy_port}/v1"

    def _update_task_completion(self, task_id: str, result: Any):
        """Internal method to update task as completed."""
        # result is a reward in this case
        task = self.tasks[task_id]
        task.status = TaskStatus.COMPLETED
        
        # Convert agent result to Rollout object if it's not already one
        if not isinstance(result, Rollout):
            result = self._convert_agent_result_to_rollout(task_id, result)
        
        task.result = result
        task.completed_at = time.time()

        # Clean up future reference
        if task_id in self.futures:
            del self.futures[task_id]
        
        return result

    def _update_task_failure(self, task_id: str, error: str):
        """Internal method to update task as failed."""
        task = self.tasks[task_id]
        task.status = TaskStatus.FAILED
        task.error = error
        task.completed_at = time.time()

        # Clean up future reference
        if task_id in self.futures:
            del self.futures[task_id]
        logger.error(f"Task {task_id} failed: {error}")

    def _convert_agent_result_to_rollout(self, task_id: str, agent_result: Any) -> Rollout:
        """
        Convert agent result to Rollout object.
        
        Args:
            task_id: Task ID for the rollout
            agent_result: Result from agent (expected to be a dict with reward and other data)
            
        Returns:
            Rollout: Rollout object with triplets
        """
        try:
            # Extract data from agent result
            if isinstance(agent_result, dict):
                reward = agent_result.get('reward', 0.0)
                prompt = agent_result.get('prompt', '')
                response_text = agent_result.get('response_text', '')
                question = agent_result.get('question', '')
                ground_truth = agent_result.get('ground_truth', '')
                answer = agent_result.get('answer', '')
                model = agent_result.get('model', '')
                endpoint = agent_result.get('endpoint', '')
                metadata = agent_result.get('metadata', {})
            elif isinstance(agent_result, float):

                reward = agent_result

            else:
                reward = getattr(agent_result, 'reward', 0.0)
                prompt = getattr(agent_result, 'prompt', '')
                response_text = getattr(agent_result, 'response_text', '')
                question = getattr(agent_result, 'question', '')
                ground_truth = getattr(agent_result, 'ground_truth', '')
                answer = getattr(agent_result, 'answer', '')
                model = getattr(agent_result, 'model', '')
                endpoint = getattr(agent_result, 'endpoint', '')
                metadata = getattr(agent_result, 'metadata', {})

            # Create simple token IDs (for now, use character-based encoding as placeholder)
            # TODO: Replace with proper tokenizer when available
            triplets = []
            if task_id in self.rollout_traj_from_proxy:
                for response_token_ids, prompt_token_ids in self.rollout_traj_from_proxy[task_id]:
                    triplet = Triplet(
                        prompt={"token_ids": prompt_token_ids},
                        response={"token_ids": response_token_ids},
                        reward=reward,
                        metadata={
                            "task_id": task_id,
                            "agent_name": self.tasks[task_id].agent_type,
                        }
                    )
                    triplets.append(triplet)
            else:
                raise ValueError(f"Task {task_id} not found in rollout_traj_from_proxy")
             
            
            # Create and return Rollout object
            rollout = Rollout(
                rollout_id=task_id,
                final_reward=reward,
                triplets=triplets,
                metadata={
                    "question": question,
                    "ground_truth": ground_truth,
                    "answer": answer,
                    "model": model,
                    "endpoint": endpoint,
                    **metadata
                }
            )
            
            return rollout
            
        except Exception as e:
            logger.error(f"Failed to convert agent result to rollout for task {task_id}: {e}")
            # Return a minimal failure rollout with empty triplets
            return Rollout(
                rollout_id=task_id,
                final_reward=0.0,
                triplets=[],
                metadata={"error": str(e)}
            )
        finally:
            # Always clean up rollout_traj_from_proxy to prevent memory leaks
            if task_id in self.rollout_traj_from_proxy:
                del self.rollout_traj_from_proxy[task_id]

    def initialize_batch_processing(self, data, is_train=True):
        """
        Initialize batch processing with data.

        Args:
            data: Dictionary containing batch data
            is_train: Whether this is for training or validation
        """
        logger.info(f"Setting up data, is_train={is_train}")

        self.reset_orchestrator_state()
        self.is_train = is_train

        # Queue tasks for agents to process
        if not data:
            logger.warning("No data provided to initialize_batch_processing")
            return

        keys = list(data.keys())
        num_samples = len(data[keys[0]]) if keys else 0
        rollouts_per_sample = 1  # Can be configured based on training needs

        logger.info(
            f"Processing {num_samples} samples with {rollouts_per_sample} rollouts each"
        )

        # TODO: add a scheduling policy
        for i in range(num_samples):
            data_id = str(uuid.uuid4())
            original_sample = {key: data[key][i] for key in keys}
            original_sample["data_id"] = data_id

            # For training, each sample can be rolled out multiple times
            for j in range(rollouts_per_sample):
                # Assume we have a default agent type registered
                if not self.agent_types:
                    logger.error(
                        "No agent types registered. Call register_agent_type() first."
                    )
                    raise RuntimeError(
                        "No agent types registered. Call register_agent_type() first."
                    )

                default_agent_type = list(self.agent_types.keys())[0]

                # Use proxy endpoint for agent tasks if proxy server is running
                proxy_endpoint = (
                    self.get_proxy_endpoint() if self._proxy_server_running else None
                )

                task_id = self.agent_rollout(
                    default_agent_type, original_sample, endpoint=proxy_endpoint
                )

                # Store original sample data to reconstruct batch information later
                self._task_id_to_original_sample[task_id] = original_sample
                self._total_tasks_queued += 1

        logger.info(
            f"Successfully queued {self._total_tasks_queued} tasks for {num_samples} samples"
        )

    def wait_for_batch_completion(self, verbose=True):
        """Wait for all queued tasks to be completed using async collector."""
        if self._total_tasks_queued == 0:
            logger.warning("No tasks were queued.")
            return

        logger.info(
            f"Starting to wait for {self._total_tasks_queued} tasks to complete"
        )
        start_time = time.time()

        # Start the async collector if not already running
        self._ensure_collector_running()

        while len(self._completed_rollouts) < self._total_tasks_queued:
            # Wait for completion events from the async collector
            try:
                # Use a short timeout to allow periodic progress logging
                future = asyncio.run_coroutine_threadsafe(
                    asyncio.wait_for(self._completion_queue.get(), timeout=1.0),
                    self._event_loop,
                )

                try:
                    completion_event = future.result(
                        timeout=1.1
                    )  # Slightly longer than asyncio timeout
                    task_id, result, error = completion_event

                    if task_id not in self._completed_rollouts:
                        if error:
                            logger.error(f"Task {task_id} failed: {error}")
                            # Still count as completed for batch completion tracking
                        else:
                            rollout = self._update_task_completion(task_id, result)
                            self._completed_rollouts[task_id] = rollout

                except asyncio.TimeoutError:
                    # Normal timeout for progress logging - check for actual task timeouts
                    logger.debug("Timeout waiting for completion events - checking for timed out tasks")
                    
                    # Check all running tasks for actual timeouts and handle them
                    timed_out_tasks = []
                    for tid, task in self.tasks.items():
                        if task.status == TaskStatus.RUNNING and self.retry_manager.is_task_timed_out(task):
                            timed_out_tasks.append(tid)
                    
                    # Cancel and retry timed out tasks
                    for tid in timed_out_tasks:
                        logger.warning(f"Task {tid} timed out - cancelling and retrying")
                        self._handle_task_timeout(tid)

            except Exception as e:
                logger.warning(f"Waiting for completion events: {e}")
                # Continue with next iteration - async collector will handle completions

            if verbose:
                elapsed_time = time.time() - start_time
                logger.info(
                    f"Completed {len(self._completed_rollouts)}/{self._total_tasks_queued} tasks... (elapsed: {elapsed_time:.1f}s)"
                )

        elapsed_time = time.time() - start_time
        logger.info(
            f"All {self._total_tasks_queued} tasks finished in {elapsed_time:.1f}s."
        )

    def collect_validation_metrics(self):
        """Calculate and return metrics for a validation run."""
        assert not self.is_train, "This method should only be called during validation."
        assert len(self._completed_rollouts) == self._total_tasks_queued

        sample_stats = []
        for task_id, result in self._completed_rollouts.items():
            # Extract metrics from result - handle both old and new result formats
            reward = self.lightning_rl._fillna_reward(result)

            # For Rollout objects, calculate response length from triplets
            response_length = 0
            turn_count = 1

            if hasattr(result, "triplets") and result.triplets:
                # New Rollout object format
                turn_count = len(result.triplets)
                for triplet in result.triplets:
                    if hasattr(triplet, "response") and isinstance(
                        triplet.response, dict
                    ):
                        response_text = triplet.response.get("text", "")
                        response_length += (
                            len(response_text.split()) if response_text else 0
                        )
            else:
                # Old CalcAgentResult format
                response_length = getattr(result, "response_length", 0)
                turn_count = getattr(result, "turn_count", 1)

            sample_stats.append(
                {
                    "reward": reward,
                    "response_length": response_length,
                    "turn_count": turn_count,
                }
            )

        if not sample_stats:
            return {
                "val/reward": 0.0,
                "val/mean_response_length": 0.0,
                "val/sum_response_length": 0.0,
                "val/turn_count": 0.0,
            }

        return {
            "val/reward": np.mean([stat["reward"] for stat in sample_stats]),
            "val/mean_response_length": np.mean(
                [stat["response_length"] for stat in sample_stats]
            ),
            "val/sum_response_length": np.mean(
                [stat["response_length"] for stat in sample_stats]
            ),
            "val/turn_count": np.mean([stat["turn_count"] for stat in sample_stats]),
        }

    def generate_training_batch(self, max_prompt_length, max_response_length, device):
        """
        Processes completed rollouts to generate a training data batch using LightningRL.

        This method delegates to the LightningRL component which implements the core
        agentic RL algorithm for decomposing trajectories into transitions and grouping
        them by task for advantage estimation.
        """
        assert self.is_train, "This method should only be called during training."
        assert len(self._completed_rollouts) == self._total_tasks_queued

        return self.lightning_rl.generate_training_batch(
            completed_rollouts=self._completed_rollouts,
            task_id_to_original_sample=self._task_id_to_original_sample,
            max_prompt_length=max_prompt_length,
            max_response_length=max_response_length,
            device=device,
        )

    def reset_orchestrator_state(self):
        """Reset the internal state of the orchestrator for the next run."""
        logger.info(f"Clearing orchestrator state:")
        logger.info(f"  - Completed rollouts: {len(self._completed_rollouts)}")
        logger.info(
            f"  - Task-to-sample mappings: {len(self._task_id_to_original_sample)}"
        )
        logger.info(f"  - Total tasks queued: {self._total_tasks_queued}")
        logger.info(f"  - Active tasks: {len(self.tasks)}")
        logger.info(f"  - Active futures: {len(self.futures)}")

        self._completed_rollouts.clear()
        self._task_id_to_original_sample.clear()
        self._total_tasks_queued = 0
        # Clear Ray task management state to prevent accumulation across runs
        self.tasks.clear()
        self.futures.clear()

        logger.info("Orchestrator state cleared")

    def _ensure_collector_running(self):
        """Ensure the async collector is running."""
        if not self._collector_running:
            self._start_collector()

    def _start_collector(self):
        """Start the async task completion collector."""
        if self._collector_running:
            logger.warning("Collector is already running")
            return

        try:
            # Create new event loop in executor thread
            def run_collector_loop():
                self._event_loop = asyncio.new_event_loop()
                asyncio.set_event_loop(self._event_loop)

                self._completion_queue = asyncio.Queue(maxsize=1000)
                self._collector_task = self._event_loop.create_task(
                    self._collector_loop()
                )
                self._collector_running = True

                logger.info("Started async task completion collector")
                self._event_loop.run_forever()

            self._executor.submit(run_collector_loop)

            # Wait a moment for the loop to start
            time.sleep(0.1)

        except Exception as e:
            logger.error(f"Failed to start collector: {e}")
            raise

    def _stop_collector(self):
        """Stop the async task completion collector."""
        if not self._collector_running:
            return

        try:
            if self._event_loop and self._collector_task:
                # Cancel the collector task
                future = asyncio.run_coroutine_threadsafe(
                    self._cancel_collector_task(), self._event_loop
                )
                future.result(timeout=2.0)

                # Stop the event loop
                self._event_loop.call_soon_threadsafe(self._event_loop.stop)

            self._collector_running = False
            self._event_loop = None
            self._collector_task = None
            self._completion_queue = None

            logger.info("Stopped async task completion collector")

        except Exception as e:
            logger.warning(f"Error stopping collector: {e}")

    async def _cancel_collector_task(self):
        """Cancel the collector task."""
        if self._collector_task:
            self._collector_task.cancel()
            try:
                await self._collector_task
            except asyncio.CancelledError:
                pass

    async def _collector_loop(self):
        """Main async loop for collecting task completions and monitoring timeouts."""
        logger.info("Collector loop started")

        while not self._shutdown_event.is_set():
            try:
                current_time = time.time()

                # Check for timed-out tasks first
                timed_out_tasks = []
                for task_id, task in self.tasks.items():
                    if (
                        task.status == TaskStatus.RUNNING
                        and self.retry_manager.is_task_timed_out(task, current_time)
                    ):
                        timed_out_tasks.append(task_id)

                # Handle timed-out tasks
                for task_id in timed_out_tasks:
                    logger.warning(
                        f"Detected timeout for task {task_id} in collector loop"
                    )
                    self._handle_task_timeout(task_id)

                # Get all running futures
                running_futures = [
                    (task_id, task.future)
                    for task_id, task in self.tasks.items()
                    if task.status == TaskStatus.RUNNING and task.future
                ]

                if not running_futures:
                    # No running tasks, sleep briefly
                    await asyncio.sleep(self.collector_timeout_ms / 1000.0)
                    continue

                futures = [future for _, future in running_futures]
                task_id_map = {
                    id(future): task_id for task_id, future in running_futures
                }

                # Wait for any completions with timeout
                try:
                    ready, not_ready = ray.wait(
                        futures,
                        num_returns=len(futures),  # Get all ready futures
                        timeout=self.collector_timeout_ms / 1000.0,
                    )

                    # Process all ready futures
                    for completed_future in ready:
                        task_id = task_id_map[id(completed_future)]

                        try:
                            result = ray.get(completed_future)
                            
                            rollout = self._update_task_completion(task_id, result)
                            self._completed_rollouts[task_id] = rollout

                            # Publish completion event
                            try:
                                await asyncio.wait_for(
                                    self._completion_queue.put((task_id, result, None)),
                                    timeout=0.1,
                                )
                            except asyncio.TimeoutError:
                                logger.warning(
                                    f"Completion queue full, dropping event for task {task_id}"
                                )

                        except Exception as e:
                            error_msg = str(e)
                            # Check if this is a Ray timeout exception
                            if "timeout" in error_msg.lower() or "timed out" in error_msg.lower():
                                error_msg = f"Task timed out: {error_msg}"
                            
                            # Check if this is a timeout failure that should trigger retry
                            retry_success = self._retry_agent_task(task_id, error_msg)

                            if not retry_success:
                                # Task permanently failed, update normally
                                self._update_task_failure(task_id, error_msg)

                                # Publish failure event
                                try:
                                    await asyncio.wait_for(
                                        self._completion_queue.put(
                                            (task_id, None, error_msg)
                                        ),
                                        timeout=0.1,
                                    )
                                except asyncio.TimeoutError:
                                    logger.warning(
                                        f"Completion queue full, dropping failure event for task {task_id}"
                                    )

                except Exception as e:
                    logger.warning(f"Error in collector loop ray.wait: {e}")
                    await asyncio.sleep(self.collector_timeout_ms / 1000.0)

            except asyncio.CancelledError:
                logger.info("Collector loop cancelled")
                break
            except Exception as e:
                logger.error(f"Unexpected error in collector loop: {e}")
                await asyncio.sleep(self.collector_timeout_ms / 1000.0)

        logger.info("Collector loop stopped")
