"""
Sandai Operator SDK - Channel implementations

Channels handle input/output communication for operators.
"""

from abc import ABCMeta, abstractmethod
from queue import Empty, Queue
from typing import Any, Dict, Optional, Tuple
from dataclasses import dataclass
import json
import logging
import socket
import threading
import time
import os
from pathlib import Path
from urllib.parse import urlparse
from .types import EOFMarker
from .env import get_celery_config, get_default_celery_config

try:
    from watchdog.observers import Observer
    from watchdog.events import FileSystemEventHandler
    WATCHDOG_AVAILABLE = True
except ImportError:
    WATCHDOG_AVAILABLE = False
    # Create a dummy class if watchdog is not available
    class FileSystemEventHandler:
        pass

from celery import Celery
import redis


logger = logging.getLogger(__name__)


@dataclass
class CeleryInflightTask:
    sdk_task_id: str
    client_task_id: Optional[str]
    routing_key: str
    priority: str
    message: Any
    delivery_attempt: Optional[int] = None


class ChannelRuntimeError(RuntimeError):
    """Raised when a channel cannot safely process or finalize a task."""


class BaseChannel(metaclass=ABCMeta):
    """Base class for all channels"""

    def __init__(self):
        self.source_queue = Queue()
        self.sink_queue = Queue()
    
    def is_empty(self):
        """Check if source queue is empty"""
        return self.source_queue.empty()
    
    def push(self, output: Dict):
        """Push output to sink queue"""
        return self.sink_queue.put(output)
    
    def pull(self, timeout=None) -> Dict:
        """Pull task from source queue"""
        task = self.source_queue.get(timeout=timeout)
        self.source_queue.task_done()
        return task
    
    def sync(self):
        """Wait for all tasks to complete"""
        self.source_queue.join()
        self.sink_queue.join()
    
    @abstractmethod
    def close(self):
        """Close the channel"""
        raise NotImplementedError
    
    def is_ready(self):
        """Check if channel is ready"""
        return True  # Base class is always ready
    
    def wait_ready(self, timeout=None):
        """Wait for channel to be ready"""
        return True  # Base class is always ready
    
    @abstractmethod
    def graceful_shutdown(self):
        raise NotImplementedError


class FileWatchHandler(FileSystemEventHandler):
    """File system event handler for watching source file changes"""
    
    def __init__(self, file_path: str, callback, debug=False):
        self.file_path = os.path.abspath(file_path)
        self.callback = callback
        self.debug = debug

    def _debug(self, message: str):
        if self.debug:
            logger.debug(message)

    def on_modified(self, event):
        if not event.is_directory:
            event_path = os.path.abspath(event.src_path)
            self._debug(f"Modified event: {event_path}")
            if event_path == self.file_path:
                self._debug("Target file modified, triggering callback")
                self.callback()
    
    def on_moved(self, event):
        # Handle file rename/move operations (common with editors like vim)
        if not event.is_directory:
            dest_path = os.path.abspath(event.dest_path)
            self._debug(f"Moved event: {os.path.abspath(event.src_path)} -> {dest_path}")
            if dest_path == self.file_path:
                self._debug("Target file moved to, triggering callback")
                self.callback()
    
    def on_created(self, event):
        # Handle file creation (when editors create new files)
        if not event.is_directory:
            event_path = os.path.abspath(event.src_path)
            self._debug(f"Created event: {event_path}")
            if event_path == self.file_path:
                self._debug("Target file created, triggering callback")
                self.callback()
    
    def on_deleted(self, event):
        # Handle file deletion (vim might delete and recreate)
        if not event.is_directory:
            event_path = os.path.abspath(event.src_path)
            self._debug(f"Deleted event: {event_path}")
            if event_path == self.file_path:
                self._debug("Target file deleted, will wait for recreation")
                # File was deleted, trigger callback to reopen file handle
                self.callback()


class FileChannel(BaseChannel):
    """
    File-based channel for reading tasks from JSONL files
    
    Supports real-time file monitoring and vim-compatible editing.
    """

    def __init__(self, source_jsonl_path: str, sink_jsonl_path: str, max_prefetch_size=100, use_file_watcher=True, debug=False, job_mode=False):
        super().__init__()
        self.max_prefetch_size = max_prefetch_size
        self.current_process_map = {}
        self._sdrun_enabled = os.getenv("SDRUN_MODE", "").lower() == "true"
        self._sdrun_world_size = self._parse_sdrun_int("SDRUN_WORLD_SIZE") or 1
        self._sdrun_rank = self._parse_sdrun_int("SDRUN_RANK") or 0
        self._line_index = 0
        self.source_jsonl_path = os.path.abspath(source_jsonl_path)
        self.sink_jsonl_path = self._resolve_sink_jsonl_path(sink_jsonl_path)
        self.source_file = open(source_jsonl_path, 'r', encoding='utf-8')
        self.sink_file = open(self.sink_jsonl_path, 'a', encoding='utf-8')
        self.use_file_watcher = use_file_watcher and WATCHDOG_AVAILABLE
        self.debug = debug
        self.job_mode = job_mode
        self._stop_event = threading.Event()
        self._file_updated_event = threading.Event()
        self._ready_event = threading.Event()
        self._source_thread = None
        self._sink_thread = None
        self._lock = threading.Lock()
        self._condition = threading.Condition(self._lock)
        self._observer = None
        self._eof_sent = False
        self._closed = False
        self._rewrite_seen_task_ids = None
        # Record file inode and modification time to detect file replacement
        self._file_stat = os.stat(source_jsonl_path)
        self._setup_file_watcher()
        self.start_prepare_source()
        self.start_prepare_sink()

    def _parse_sdrun_int(self, env_name: str) -> Optional[int]:
        value = os.getenv(env_name)
        if value is None:
            return None
        try:
            return int(value)
        except ValueError:
            return None

    def _resolve_sink_jsonl_path(self, sink_jsonl_path: str) -> str:
        if not self._sdrun_enabled:
            return os.path.abspath(sink_jsonl_path)

        path = Path(sink_jsonl_path)
        suffix = path.suffix
        if suffix:
            resolved = path.with_name(f"{path.stem}.{self._sdrun_rank}{suffix}")
        else:
            resolved = path.with_name(f"{path.name}.{self._sdrun_rank}")
        return os.path.abspath(str(resolved))

    def _should_take_line(self, line_index: int) -> bool:
        if not self._sdrun_enabled:
            return True
        if self._sdrun_world_size <= 1:
            return True
        return line_index % self._sdrun_world_size == self._sdrun_rank

    def _debug(self, message: str):
        if self.debug:
            logger.debug(message)

    def _log_exception(self, message: str):
        logger.exception(message)

    def _log_warning(self, message: str):
        logger.warning(message)

    def _safe_close_file(self, file_obj):
        try:
            file_obj.close()
        except OSError:
            pass

    def _mark_rewrite_window(self):
        self._rewrite_seen_task_ids = set(self.current_process_map.keys())

    def _clear_rewrite_window_if_idle(self):
        if self._rewrite_seen_task_ids is not None and not self.current_process_map:
            self._rewrite_seen_task_ids = None

    def _should_skip_rewrite_duplicate(self, task_id: str) -> bool:
        return self._rewrite_seen_task_ids is not None and task_id in self._rewrite_seen_task_ids
    
    def _setup_file_watcher(self):
        """Setup file system watcher if available"""
        if self.use_file_watcher:
            try:
                self._observer = Observer()
                handler = FileWatchHandler(self.source_jsonl_path, self._on_file_updated, self.debug)
                watch_dir = os.path.dirname(self.source_jsonl_path)
                self._observer.schedule(handler, watch_dir, recursive=False)
                self._observer.start()
                self._debug(f"File watcher setup for: {self.source_jsonl_path}")
                self._debug(f"Watching directory: {watch_dir}")
            except Exception as e:
                self._log_warning(f"Failed to setup file watcher, falling back to polling: {e}")
                self.use_file_watcher = False
    
    def _on_file_updated(self):
        """Callback when file is updated"""
        self._debug(f"File update detected: {self.source_jsonl_path}")
        # Reopen file handle to handle file replacement
        self._reopen_source_file()
        self._file_updated_event.set()
    
    def _reopen_source_file(self):
        """Reopen source file handle"""
        try:
            # Check if file still exists and is readable
            if not os.path.exists(self.source_jsonl_path):
                self._debug("Source file does not exist, waiting for recreation...")
                return
            
            # Get new file status
            new_stat = os.stat(self.source_jsonl_path)
            
            # Check if it's the same file (by inode)
            is_same_file = (hasattr(self, '_file_stat') and 
                          new_stat.st_ino == self._file_stat.st_ino and
                          new_stat.st_dev == self._file_stat.st_dev)
            
            # Try to get current position
            try:
                current_pos = self.source_file.tell()
            except OSError:
                current_pos = getattr(self, '_last_known_position', 0)
            
            # Try to reopen file
            new_file = open(self.source_jsonl_path, 'r', encoding='utf-8')
            
            try:
                new_file.seek(0, 2)  # Move to end of file
                new_size = new_file.tell()
                
                if not is_same_file:
                    # For rewritten files, restart scanning from the beginning.
                    # Skip only tasks that were already in-flight during the rewrite scan.
                    new_file.seek(0)
                    self._last_known_position = 0
                    self._line_index = 0
                    self._mark_rewrite_window()
                    self._debug("File was replaced (inode changed), restarting from beginning")
                elif new_size >= current_pos:
                    # Same file and grew, continue from previous position
                    new_file.seek(current_pos)
                    self._debug(f"File size increased, continuing from position {current_pos}")
                else:
                    # Same file but smaller, start from beginning
                    new_file.seek(0)
                    self._last_known_position = 0
                    self._line_index = 0
                    self._debug("Source file was truncated, restarting from beginning")
                
                # Update file status
                self._file_stat = new_stat
                
                # Close old file handle, use new one
                old_file = self.source_file
                self.source_file = new_file
                self._safe_close_file(old_file)

                self._debug(f"Successfully reopened source file at position {new_file.tell()}")

            except Exception as e:
                # If error, close new file handle, continue with old one
                new_file.close()
                self._debug(f"Failed to reopen source file: {e}")
                    
        except Exception as e:
            self._debug(f"Error reopening source file: {e}")
    
    def start_prepare_source(self):
        """Read JSON tasks from source file in a separate thread, controlling process-map size"""
        def _source_worker():
            try:
                first_read_attempt = True
                while not self._stop_event.is_set():
                    # Read available lines from current position using readline()
                    lines_read = False
                    
                    while True:
                        if self._stop_event.is_set():
                            return
                        
                        try:
                            # Record current position
                            current_pos = self.source_file.tell()
                            self._last_known_position = current_pos
                        except OSError:
                            # tell() might be disabled, skip position recording
                            pass
                        
                        line = self.source_file.readline()
                        
                        # If no content read, reached end of file
                        if not line:
                            # After first read attempt, set ready flag
                            if first_read_attempt:
                                first_read_attempt = False
                                self._ready_event.set()
                                self._debug("Channel is now ready - first read attempt completed")
                            break
                        
                        line = line.strip()
                        current_line_index = self._line_index
                        self._line_index += 1
                        if not line:
                            continue
                        if not self._should_take_line(current_line_index):
                            continue
                        
                        lines_read = True
                        
                        try:
                            task = json.loads(line)
                            task_id = task.get('task_id')
                            if not task_id:
                                continue
                            if self._should_skip_rewrite_duplicate(task_id):
                                continue
                            
                            # Wait until process-map size is less than max_prefetch_size
                            with self._condition:
                                while len(self.current_process_map) >= self.max_prefetch_size:
                                    if self._stop_event.is_set():
                                        return
                                    self._condition.wait()  # Wait for pull to trigger notification
                                
                                # Add to process-map and put into queue
                                self.current_process_map[task_id] = task
                            self.source_queue.put(task)
                            
                            # After first successful task read, set ready flag
                            if first_read_attempt:
                                first_read_attempt = False
                                self._ready_event.set()
                                self._debug("Channel is now ready - first task read successfully")
                            
                        except json.JSONDecodeError:
                            continue  # Skip invalid JSON lines
                    
                    # If no lines were read, wait for file updates (unless in job_mode and EOF sent)
                    if not lines_read:
                        # In job_mode, if EOF reached and not sent yet, send EOF marker
                        if self.job_mode and not self._eof_sent:
                            self._eof_sent = True
                            eof_marker = EOFMarker()
                            self.source_queue.put(eof_marker.model_dump())
                            self._debug("Job mode: EOF reached, sent EOF marker")
                            break
                        
                        # In job_mode, if EOF marker already sent, stop waiting for file updates
                        if self.job_mode and self._eof_sent:
                            self._debug("Job mode: EOF sent, stopping source worker")
                            break
                        
                        if self.use_file_watcher:
                            # Use file system events - more efficient
                            # But also check periodically to prevent missing events
                            event_triggered = self._file_updated_event.wait(timeout=1.0)
                            self._file_updated_event.clear()

                            if event_triggered:
                                self._debug("File watcher event triggered, checking for updates...")
                            
                            # Check file updates regardless of events (forced polling)
                            if self._check_file_updated():
                                self._debug("File updated detected, continuing to read...")
                                continue  # Has updates, continue reading
                            else:
                                if event_triggered:
                                    self._debug("File watcher event triggered but no actual update detected")
                                # No actual update, sleep briefly and continue loop
                                time.sleep(0.2)
                        else:
                            # Fallback to polling
                            if self._check_file_updated():
                                self._debug("File updated detected via polling, continuing to read...")
                                continue  # Has updates, continue reading
                            else:
                                time.sleep(0.5)  # No updates, sleep a bit
                        
            except Exception as e:
                self._log_exception(f"Source worker error: {e}")
        
        self._source_thread = threading.Thread(target=_source_worker, daemon=True)
        self._source_thread.start()
    
    def _check_file_updated(self) -> bool:
        """Check if source file has been updated/appended (fallback method)"""
        try:
            # Try to get current position
            try:
                current_pos = self.source_file.tell()
            except OSError:
                # tell() disabled, use recorded position
                current_pos = getattr(self, '_last_known_position', 0)
            
            # First check if file still exists
            if not os.path.exists(self.source_jsonl_path):
                self._debug("Source file does not exist")
                return False
            
            # Get actual file size
            file_stat = os.stat(self.source_jsonl_path)
            file_size = file_stat.st_size
            previous_stat = getattr(self, "_file_stat", None)

            if previous_stat and (
                file_stat.st_ino != previous_stat.st_ino
                or file_stat.st_dev != previous_stat.st_dev
            ):
                self._debug("Source file inode changed, reopening from beginning")
                self._reopen_source_file()
                return True

            self._debug(f"File check: current_pos={current_pos}, file_size={file_size}")
            
            # If file size > current position, there's new content
            if file_size > current_pos:
                self._debug(f"File has new content: {file_size - current_pos} bytes")
                return True
            
            # If file size < current position, file was truncated or rewritten
            if file_size < current_pos:
                self._debug("Source file was truncated, restarting from beginning")
                try:
                    self.source_file.seek(0)
                    self._last_known_position = 0
                    self._line_index = 0
                except OSError:
                    # If seek also fails, reopen file
                    self._reopen_source_file()
                return True
            
            self._debug("No file update detected")
            return False
            
        except OSError as e:
            self._debug(f"Error checking file update: {e}")
            return False
    
    def start_prepare_sink(self):
        """Read from sink queue and write to output file in a separate thread"""
        def _sink_worker():
            try:
                while not self._stop_event.is_set() or not self.sink_queue.empty():
                    try:
                        # Get result from sink queue with timeout to avoid infinite waiting
                        result = self.sink_queue.get(timeout=0.1)
                    except Empty:
                        continue
                    try:
                        
                        # Write to output file
                        json_line = json.dumps(result, ensure_ascii=False)
                        self.sink_file.write(json_line + '\n')
                        self.sink_file.flush()
                        
                        # Remove completed task from process-map
                        task_id = result.get('task_id')
                        if task_id:
                            with self._condition:
                                self.current_process_map.pop(task_id, None)
                                self._clear_rewrite_window_if_idle()
                                self._condition.notify()  # Notify source thread to continue reading
                        
                        self.sink_queue.task_done()
                    except Exception as e:
                        self.sink_queue.task_done()
                        self._log_exception(f"Sink worker failed to write result: {e}")
                        
            except Exception as e:
                self._log_exception(f"Sink worker error: {e}")
        
        self._sink_thread = threading.Thread(target=_sink_worker, daemon=True)
        self._sink_thread.start()
    
    def pull(self, timeout=None) -> Dict:
        """Override pull method to notify source thread when task is pulled"""
        task = super().pull(timeout=timeout)
        
        # In job_mode, immediately remove task from process_map since we don't rely on sink worker
        if self.job_mode:
            task_id = task.get('task_id')
            if task_id:
                with self._condition:
                    self.current_process_map.pop(task_id, None)
                    self._clear_rewrite_window_if_idle()
                    self._condition.notify()  # Notify source thread to continue reading
        else:
            # In normal mode, only notify source thread, cleanup handled by sink worker
            with self._condition:
                self._condition.notify()
        
        return task
    
    def close(self):
        """Close channel and stop all threads"""
        if self._closed:
            return
        self._closed = True
        self._stop_event.set()
        self._file_updated_event.set()  # Wake up waiting source thread
        with self._condition:
            self._condition.notify_all()
        
        # Stop file watcher
        if self._observer:
            self._observer.stop()
            self._observer.join()
        
        # Wait for threads to finish
        if self._source_thread and self._source_thread.is_alive():
            self._source_thread.join(timeout=1.0)
        if self._sink_thread and self._sink_thread.is_alive():
            self._sink_thread.join(timeout=1.0)
        
        self._safe_close_file(self.source_file)
        self._safe_close_file(self.sink_file)

    def graceful_shutdown(self):
        self._stop_event.set()
        self._file_updated_event.set()
        with self._condition:
            self._condition.notify_all()
    
    def is_ready(self):
        """Check if channel is ready (has attempted to read file at least once)"""
        return self._ready_event.is_set()
    
    def wait_ready(self, timeout=None):
        """Wait for channel to be ready"""
        return self._ready_event.wait(timeout)


class CeleryChannel(BaseChannel):
    """
    Celery-based channel for receiving tasks from Celery and returning results
    
    This channel integrates with Celery to receive tasks from a Celery queue,
    process them, and return results back to Celery.
    """

    def __init__(self, operator_name: str = "unknown",
                 operator_version: str = "1.0.0", priorities: Optional[list] = None,
                 visibility_timeout: int = 300, max_prefetch_size: int = 100, debug: bool = False,
                 job_mode: bool = False, job_mode_timeout: int = 30, result_expires: int = 86400,
                 broker_url: Optional[str] = None, result_backend: Optional[str] = None,
                 task_serializer: str = "json", result_serializer: str = "json",
                 accept_content: Optional[list] = None,
                 max_delivery_attempts: Optional[int] = None):
        
        super().__init__()
        self.operator_name = operator_name
        self.operator_version = operator_version
        
        # 生成算子特定的任务名称，避免多个算子间的配置干扰
        self.task_name = f"sandai_operator.{operator_name}.{operator_version}.process_task"
        
        # 为每个算子创建独立的 Celery 应用实例，避免配置干扰
        app_name = f"sandai_operator_{operator_name}_{operator_version}"
        self.celery_app = Celery(app_name)
        
        # 设置优先级列表（默认包含 high, normal, low, very_low）
        self.priorities = priorities or ["high", "normal", "low", "very_low"]
        
        # 生成多个优先级队列：<priority>.<name>.<version>
        self.queue_names = [f"{priority}.{operator_name}.{operator_version}" for priority in self.priorities]
        
        # 为了兼容性，保留主队列名称（使用第一个优先级）
        self.queue_name = self.queue_names[0] if self.queue_names else f"{operator_name}.{operator_version}"
        
        # 设置 visibility timeout（消息可见性超时）
        self.visibility_timeout = visibility_timeout
        self.result_expires = result_expires
        self.broker_url = broker_url
        self.result_backend = result_backend
        self.task_serializer = task_serializer
        self.result_serializer = result_serializer
        self.accept_content = list(accept_content or ["json", "msgpack"])
        self.max_delivery_attempts = max(0, int(max_delivery_attempts or 0))

        self.max_prefetch_size = max_prefetch_size
        self.debug = debug
        self.job_mode = job_mode
        self.job_mode_timeout = job_mode_timeout  # job mode 超时时间（秒）
        self._stop_event = threading.Event()
        self._ready_event = threading.Event()
        self._source_thread = None
        self._sink_thread = None
        self._lock = threading.Lock()
        self._condition = threading.Condition(self._lock)
        self._current_process_map = {}
        self._last_task_time = None  # 记录最后一次接收到任务的时间
        self._eof_sent = False  # 是否已发送 EOF 标记
        self._closed = False
        self._shutdown_started = False
        self._state_redis_client = None

        # 设置 Celery 消费者
        self._setup_celery_consumer()
        self._setup_state_redis_client()
        self.start_prepare_source()
        self.start_prepare_sink()

    def _debug(self, message: str):
        if self.debug:
            logger.debug(message)

    def _log_exception(self, message: str):
        logger.exception(message)

    def _log_warning(self, message: str):
        logger.warning(message)

    def _redis_connection_kwargs(self, connection_url: Optional[str]) -> Optional[Dict[str, Any]]:
        if not connection_url:
            return None

        parsed = urlparse(connection_url)
        if parsed.scheme not in {"redis", "rediss"}:
            return None

        db = 0
        if parsed.path:
            raw_db = parsed.path.lstrip("/")
            if raw_db:
                try:
                    db = int(raw_db)
                except ValueError:
                    self._log_warning(
                        f"Invalid Redis database in URL {connection_url}, falling back to db=0"
                    )

        connection_kwargs: Dict[str, Any] = {
            "host": parsed.hostname or "localhost",
            "port": parsed.port or 6379,
            "username": parsed.username or None,
            "password": parsed.password or None,
            "db": db,
            "decode_responses": True,
        }
        if parsed.scheme == "rediss":
            connection_kwargs["ssl"] = True
        return connection_kwargs

    def _setup_state_redis_client(self):
        if self.max_delivery_attempts <= 0:
            return

        conf = self.celery_app.conf
        result_backend = (
            conf.get("result_backend") if isinstance(conf, dict) else getattr(conf, "result_backend", None)
        )
        broker_url = conf.get("broker_url") if isinstance(conf, dict) else getattr(conf, "broker_url", None)

        for label, candidate_url in (("result backend", result_backend), ("broker", broker_url)):
            connection_kwargs = self._redis_connection_kwargs(candidate_url)
            if connection_kwargs is None:
                continue

            try:
                client = redis.Redis(**connection_kwargs)
                client.ping()
            except Exception as exc:
                self._log_warning(
                    f"CeleryChannel state Redis from {label} is unavailable: {candidate_url} ({exc})"
                )
                continue

            self._state_redis_client = client
            self._debug(f"CeleryChannel using Redis state store from {label}: {candidate_url}")
            return

        self._log_warning(
            "CeleryChannel delivery-attempt protection disabled because neither result backend nor broker uses Redis"
        )

    def _delivery_tracking_ttl(self) -> int:
        return max(self.result_expires, self.visibility_timeout * (self.max_delivery_attempts + 1), 3600)

    def _delivery_attempt_key(self, client_task_id: str) -> str:
        return (
            f"sandai:operator:{self.operator_name}:{self.operator_version}:"
            f"delivery-attempts:{client_task_id}"
        )

    def _finalized_result_key(self, client_task_id: str) -> str:
        return (
            f"sandai:operator:{self.operator_name}:{self.operator_version}:"
            f"finalized-result:{client_task_id}"
        )

    def _record_delivery_attempt(self, inflight_task: CeleryInflightTask) -> Optional[int]:
        if (
            self.max_delivery_attempts <= 0
            or not inflight_task.client_task_id
            or self._state_redis_client is None
        ):
            return None

        ttl = self._delivery_tracking_ttl()
        attempt_key = self._delivery_attempt_key(inflight_task.client_task_id)
        try:
            attempt = int(self._state_redis_client.incr(attempt_key))
            self._state_redis_client.expire(attempt_key, ttl)
            inflight_task.delivery_attempt = attempt
            return attempt
        except Exception as exc:
            self._log_warning(
                f"Failed to track delivery attempt for task {inflight_task.client_task_id}: {exc}"
            )
            return None

    def _claim_terminal_result(self, client_task_id: Optional[str], status: str) -> bool:
        if not client_task_id or self._state_redis_client is None:
            return True

        try:
            claimed = self._state_redis_client.set(
                self._finalized_result_key(client_task_id),
                status,
                nx=True,
                ex=self._delivery_tracking_ttl(),
            )
        except Exception as exc:
            self._log_warning(
                f"Failed to claim terminal result for task {client_task_id}: {exc}"
            )
            return True

        return bool(claimed)

    def _create_delivery_limit_error_output(
        self, inflight_task: CeleryInflightTask, delivery_attempt: int
    ) -> Dict[str, Any]:
        return {
            "task_id": inflight_task.sdk_task_id,
            "status": "error",
            "error": {
                "error_type": "TaskDeliveryLimitExceededError",
                "error_message": (
                    "Task exceeded the maximum number of delivery attempts "
                    f"({delivery_attempt} > {self.max_delivery_attempts})"
                ),
                "error_details": {
                    "client_task_id": inflight_task.client_task_id,
                    "delivery_attempt": delivery_attempt,
                    "max_delivery_attempts": self.max_delivery_attempts,
                    "routing_key": inflight_task.routing_key,
                },
            },
            "artifacts": [],
            "results": None,
            "__keep_raw_output__": True,
        }

    def _finalize_excessive_redelivery(
        self, inflight_task: CeleryInflightTask, delivery_attempt: int
    ) -> None:
        if self._claim_terminal_result(
            inflight_task.client_task_id, "delivery-limit-exceeded"
        ):
            self._store_celery_result(
                inflight_task,
                self._create_delivery_limit_error_output(inflight_task, delivery_attempt),
            )
        else:
            self._debug(
                "Skipping delivery-limit failure for task %s because another worker already finalized it"
                % (inflight_task.client_task_id or inflight_task.sdk_task_id)
            )

        self._finalize_message(
            inflight_task.message,
            ack=True,
            task_id=inflight_task.sdk_task_id,
        )

    def _prepare_incoming_task(self, body: Any, message) -> Tuple[CeleryInflightTask, Dict[str, Any]]:
        routing_key = message.delivery_info.get('routing_key', 'unknown')
        task_priority = self._extract_priority(routing_key)
        inflight_task = self._build_inflight_task(message, routing_key, task_priority)

        self._debug(f"Generated task ID: {inflight_task.sdk_task_id}")
        self._debug(f"Task priority: {task_priority} (from queue: {routing_key})")
        self._debug(f"Raw body type: {type(body)}")
        self._debug(f"Raw body content: {body}")

        task_data = self._extract_task_data(body)
        if task_data is None:
            raise ChannelRuntimeError(f"Unsupported Celery task payload: {type(body)}")

        task_data = self._normalize_task_payload(task_data, inflight_task)
        if task_data.get("task_type") != 'desc':
            task_data.setdefault('artifacts', [])
            task_data.setdefault('options', {})

        return inflight_task, task_data

    def _register_inflight_task(self, inflight_task: CeleryInflightTask):
        with self._condition:
            while len(self._current_process_map) >= self.max_prefetch_size:
                if self._stop_event.is_set():
                    return False
                self._condition.wait()

            self._current_process_map[inflight_task.sdk_task_id] = inflight_task
            return True

    def _should_emit_job_mode_eof(self, current_time: Optional[float] = None) -> bool:
        if not self.job_mode or self._eof_sent or not self._last_task_time:
            return False
        current_time = current_time or time.time()
        return current_time - self._last_task_time > self.job_mode_timeout

    def _emit_job_mode_eof(self, reason: str):
        self._eof_sent = True
        eof_marker = EOFMarker()
        self.source_queue.put(eof_marker.model_dump())
        self._debug(f"Job mode: Sent EOF marker due to {reason}")

    def _is_timeout_error(self, error: Exception) -> bool:
        if isinstance(error, (TimeoutError, socket.timeout)):
            return True

        error_str = str(error).strip().lower()
        return "timed out" in error_str or error_str == "timeout"

    def _build_consumer_queues(self):
        from kombu import Queue

        queues = []
        for queue_name in self.queue_names:
            queue = Queue(
                queue_name,
                queue_arguments={
                    'x-visibility-timeout': self.visibility_timeout
                }
            )
            queues.append(queue)
        return queues

    def _handle_incoming_message(self, body: Any, message):
        self._debug(f"Received Celery task: {body}")
        self._debug(f"Message properties: {message.properties}")

        inflight_task, task_data = self._prepare_incoming_task(body, message)
        celery_task_id = inflight_task.sdk_task_id
        self._debug(f"Using Celery task ID: {celery_task_id}")

        delivery_attempt = self._record_delivery_attempt(inflight_task)
        if delivery_attempt is not None:
            task_data.setdefault("meta", {}).setdefault("celery", {})[
                "delivery_attempt"
            ] = delivery_attempt
            self._debug(
                f"Tracked delivery attempt {delivery_attempt} for task {inflight_task.client_task_id}"
            )
            if delivery_attempt > self.max_delivery_attempts:
                self._log_warning(
                    "Task %s exceeded max delivery attempts: %s > %s"
                    % (
                        inflight_task.client_task_id or celery_task_id,
                        delivery_attempt,
                        self.max_delivery_attempts,
                    )
                )
                self._finalize_excessive_redelivery(inflight_task, delivery_attempt)
                return

        self._last_task_time = time.time()

        if not self._register_inflight_task(inflight_task):
            self._finalize_message(message, ack=False, task_id=celery_task_id)
            return

        self.source_queue.put(task_data)

        if not self._ready_event.is_set():
            self._ready_event.set()
            self._debug("Celery channel is now ready")

    def _drain_consumer_loop(self, connection):
        while not self._stop_event.is_set():
            try:
                if self._should_emit_job_mode_eof():
                    self._emit_job_mode_eof("timeout")
                    break

                timeout = 1.0 if self.job_mode else 5.0
                connection.drain_events(timeout=timeout)

            except Exception as e:
                if not self._is_timeout_error(e):
                    self._log_warning(f"Celery drain_events error: {e}")
                if self._should_emit_job_mode_eof():
                    self._emit_job_mode_eof("timeout during drain_events")
                    break

                continue

    def _extract_priority(self, routing_key: str) -> str:
        for priority in self.priorities:
            if routing_key.startswith(f"{priority}."):
                return priority
        return "normal"

    def _build_celery_task_id(self, message, routing_key: str) -> str:
        correlation_id = message.properties.get("correlation_id")
        if correlation_id:
            return f"{routing_key}:{correlation_id}:{id(message)}"
        return f"{routing_key}:{id(message)}"

    def _build_inflight_task(self, message, routing_key: str, priority: str) -> CeleryInflightTask:
        return CeleryInflightTask(
            sdk_task_id=self._build_celery_task_id(message, routing_key),
            client_task_id=message.properties.get("correlation_id"),
            routing_key=routing_key,
            priority=priority,
            message=message,
        )

    def _normalize_task_payload(self, task_data: Dict[str, Any], inflight_task: CeleryInflightTask) -> Dict[str, Any]:
        normalized_task = task_data.copy()
        normalized_task["task_id"] = inflight_task.sdk_task_id
        normalized_task["priority"] = inflight_task.priority

        meta = dict(normalized_task.get("meta") or {})
        celery_meta = dict(meta.get("celery") or {})
        celery_meta.update(
            {
                "client_task_id": inflight_task.client_task_id,
                "routing_key": inflight_task.routing_key,
                "sdk_task_id": inflight_task.sdk_task_id,
            }
        )
        meta["celery"] = celery_meta
        normalized_task["meta"] = meta
        return normalized_task

    def _extract_task_data(self, body: Any) -> Optional[Dict[str, Any]]:
        if isinstance(body, dict):
            return body.copy()

        if isinstance(body, (list, tuple)) and len(body) >= 2:
            args, kwargs = body[0], body[1]
            if isinstance(kwargs, dict):
                return kwargs.copy()
            if isinstance(args, list) and args and isinstance(args[0], dict):
                return args[0].copy()

        return None

    def _store_celery_result(self, inflight_task: CeleryInflightTask, result: Dict[str, Any]):
        if not inflight_task.client_task_id:
            return

        payload = dict(result)
        keep_raw_output = payload.pop("__keep_raw_output__", None)

        if payload.get("status") in {"success", "error"}:
            celery_result = payload
            if not keep_raw_output:
                celery_result["artifacts"] = payload.get("artifacts", [])
                celery_result["results"] = payload.get("results", {})
                celery_result["error"] = payload.get("error", None)
            celery_status = "SUCCESS"
        else:
            error_info = payload.get("error", {})
            if isinstance(error_info, dict):
                error_message = error_info.get("error_message", "Unknown error")
                error_type = error_info.get("error_type", "ProcessingError")
            else:
                error_message = str(error_info) if error_info else "Unknown error"
                error_type = "ProcessingError"

            celery_result = Exception(f"{error_type}: {error_message}")
            celery_status = "FAILURE"

        self.celery_app.backend.store_result(
            inflight_task.client_task_id,
            celery_result,
            celery_status,
        )

    def _complete_inflight_task(self, process_info: CeleryInflightTask, result: Dict[str, Any]):
        celery_task_id = result.get('task_id', process_info.sdk_task_id)
        original_task_id = process_info.client_task_id
        message = process_info.message

        if not self._claim_terminal_result(original_task_id, result.get("status", "unknown")):
            self._debug(
                f"Skipping result store for task {original_task_id or celery_task_id} because another worker already finalized it"
            )
            self._finalize_message(message, ack=True, task_id=celery_task_id)
            return

        try:
            self._debug(f"Processed Celery task {celery_task_id}")
            self._debug(f"Original task ID: {original_task_id}")
            self._debug(f"Result status: {result.get('status', 'unknown')}")

            self._store_celery_result(process_info, result)
            if original_task_id:
                self._debug(f"Stored result for task {original_task_id}")
            self._finalize_message(message, ack=True, task_id=celery_task_id)
        except Exception as e:
            self._log_exception(f"Error sending result for task {celery_task_id}: {e}")
            self._finalize_message(message, ack=False, task_id=celery_task_id, error=e)

    def _finalize_message(self, message, *, ack: bool, task_id: str, error: Optional[Exception] = None):
        action = message.ack if ack else message.reject
        try:
            action()
        except Exception as exc:
            suffix = f": {error}" if error is not None else ""
            self._log_exception(
                f"Failed to {'ack' if ack else 'reject'} Celery task {task_id}{suffix}: {exc}"
            )

    def _pop_process_info(self, celery_task_id: str) -> Optional[CeleryInflightTask]:
        with self._condition:
            process_info = self._current_process_map.pop(celery_task_id, None)
            self._condition.notify_all()
            return process_info
    
    def _setup_celery_consumer(self):
        """设置 Celery 消费者配置"""
        # 从显式配置获取 Celery 配置，未提供时再回退到环境/默认配置
        env_celery_config = None
        if self.broker_url and self.result_backend:
            env_celery_config = type("ResolvedCeleryConfig", (), {
                "broker_url": self.broker_url,
                "result_backend": self.result_backend,
                "task_serializer": self.task_serializer,
                "result_serializer": self.result_serializer,
                "accept_content": self.accept_content,
            })()
        else:
            env_celery_config = get_celery_config()
            if not env_celery_config:
                # 如果环境未配置，使用默认配置
                env_celery_config = get_default_celery_config()
        
        self.celery_app.conf.update({
            'broker_url': env_celery_config.broker_url,
            'result_backend': env_celery_config.result_backend,
            'task_serializer': env_celery_config.task_serializer,
            'result_serializer': env_celery_config.result_serializer,
            'accept_content': env_celery_config.accept_content,

            # 设置 Redis broker 的 visibility timeout
            'broker_transport_options': {
                'visibility_timeout': self.visibility_timeout,
                'fanout_prefix': True,
                'fanout_patterns': True
            },
            # 设置 result backend 过期时间，避免结果键长期堆积
            'result_expires': self.result_expires,
        })
        
        self._debug(f"Celery channel setup for operator: {self.operator_name} v{self.operator_version}")
        self._debug(f"Task name: {self.task_name}")
        self._debug(f"Celery app: {self.celery_app.main}")
        self._debug(f"Broker: {env_celery_config.broker_url}")
        self._debug(f"Backend: {env_celery_config.result_backend}")
        self._debug(f"Queues: {', '.join(self.queue_names)}")
        self._debug(f"Priority order: {' > '.join(self.priorities)}")
        self._debug(f"Visibility timeout: {self.visibility_timeout} seconds")
        self._debug(f"Result expires: {self.result_expires} seconds")
    
    def start_prepare_source(self):
        """从 Celery 队列接收任务的线程"""
        def _source_worker():
            try:
                # 设置 Celery 消费者
                with self.celery_app.connection() as connection:
                    queues = self._build_consumer_queues()
                    
                    self._debug(
                        f"Created {len(queues)} priority queues with visibility_timeout={self.visibility_timeout}s:"
                    )
                    for i, (priority, queue_name) in enumerate(zip(self.priorities, self.queue_names)):
                        self._debug(f"  {i+1}. {priority}: {queue_name}")

                    def process_task(body, message):
                        """处理接收到的 Celery 任务"""
                        try:
                            self._handle_incoming_message(body, message)
                        except Exception as e:
                            self._log_exception(f"Error processing Celery task: {e}")
                            self._finalize_message(message, ack=False, task_id=f"incoming:{id(message)}", error=e)
                    
                    # 创建消费者并设置多个优先级队列
                    # 队列按优先级顺序排列，Celery 会优先处理前面的队列
                    consumer = connection.Consumer(queues=queues, callbacks=[process_task], no_ack=False)
                    
                    # 启动消费
                    consumer.consume()
                    self._ready_event.set()
                    
                    # 初始化最后任务时间
                    self._last_task_time = time.time()

                    self._drain_consumer_loop(connection)
                    
            except Exception as e:
                self._log_exception(f"Celery source worker error: {e}")
        
        self._source_thread = threading.Thread(target=_source_worker, daemon=True)
        self._source_thread.start()
    
    def start_prepare_sink(self):
        """处理结果并返回给 Celery 的线程"""
        def _sink_worker():
            try:
                while not self._stop_event.is_set() or not self.sink_queue.empty():
                    try:
                        # 从 sink 队列获取结果
                        result = self.sink_queue.get(timeout=0.1)
                    except Empty:
                        continue
                    try:
                        
                        celery_task_id = result.get('task_id')
                        if not celery_task_id:
                            self.sink_queue.task_done()
                            continue
                        
                        # 获取对应的 Celery 消息
                        process_info = self._pop_process_info(celery_task_id)

                        if not process_info:
                            self._log_warning(
                                f"Missing Celery process info for task {celery_task_id}, cannot ack/reject"
                            )
                            self.sink_queue.task_done()
                            continue

                        self._complete_inflight_task(process_info, result)

                        self.sink_queue.task_done()
                    except Exception as e:
                        self.sink_queue.task_done()
                        self._log_exception(f"Celery sink worker error: {e}")
                        
            except Exception as e:
                self._log_exception(f"Celery sink worker failed: {e}")
        
        self._sink_thread = threading.Thread(target=_sink_worker, daemon=True)
        self._sink_thread.start()
    
    def pull(self, timeout=None) -> Dict:
        """从源队列拉取任务"""
        task = super().pull(timeout=timeout)

        # 由 sink worker 统一负责清理 _current_process_map 并 ack/reject
        with self._condition:
            self._condition.notify()

        return task
    
    def graceful_shutdown(self):
        """优雅关闭，释放所有未完成的任务"""
        with self._condition:
            if self._shutdown_started:
                return
            self._shutdown_started = True
            self._stop_event.set()
            self._condition.notify_all()
            pending_processes = list(self._current_process_map.items())
            self._current_process_map.clear()

        self._debug("Starting graceful shutdown of CeleryChannel...")

        rejected_count = 0
        for task_id, process_info in pending_processes:
            try:
                process_info.message.reject()
                rejected_count += 1
                self._debug(f"Rejected task {task_id} for reprocessing")
            except Exception as e:
                self._log_warning(f"Error rejecting task {task_id}: {e}")

        self._debug(f"Graceful shutdown completed: rejected {rejected_count} tasks")
    
    def close(self):
        """关闭 channel 并停止所有线程"""
        if self._closed:
            return
        self._closed = True

        self.graceful_shutdown()
        
        # 等待线程完成
        if self._source_thread and self._source_thread.is_alive():
            self._source_thread.join(timeout=2.0)
        if self._sink_thread and self._sink_thread.is_alive():
            self._sink_thread.join(timeout=2.0)
    
    def is_ready(self):
        """检查 channel 是否就绪"""
        return self._ready_event.is_set()
    
    def wait_ready(self, timeout=None):
        """等待 channel 就绪"""
        return self._ready_event.wait(timeout)
