"""
Sandai Operator SDK - Channel implementations

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

from abc import ABCMeta, abstractmethod
from typing import Dict, Optional
import json
import threading
import time
import os
from queue import Queue
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
from celery.result import AsyncResult


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()
    
    def close(self):
        """Close the channel"""
        ...
    
    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
    
    def graceful_shutdown(self):
        ...


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 on_modified(self, event):
        if not event.is_directory:
            event_path = os.path.abspath(event.src_path)
            if self.debug:
                print(f"🔧 Modified event: {event_path}")
            if event_path == self.file_path:
                if self.debug:
                    print(f"✅ 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)
            if self.debug:
                print(f"🔄 Moved event: {os.path.abspath(event.src_path)} → {dest_path}")
            if dest_path == self.file_path:
                if self.debug:
                    print(f"✅ 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)
            if self.debug:
                print(f"🆕 Created event: {event_path}")
            if event_path == self.file_path:
                if self.debug:
                    print(f"✅ 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)
            if self.debug:
                print(f"🗑️ Deleted event: {event_path}")
            if event_path == self.file_path:
                if self.debug:
                    print(f"✅ 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.source_jsonl_path = os.path.abspath(source_jsonl_path)
        self.source_file = open(source_jsonl_path, 'r', encoding='utf-8')
        self.sink_file = open(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
        # 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 _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()
                if self.debug:
                    print(f"File watcher setup for: {self.source_jsonl_path}")
                    print(f"Watching directory: {watch_dir}")
            except Exception as e:
                print(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"""
        if self.debug:
            print(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):
                if self.debug:
                    print(f"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:
                    # File was replaced, handle intelligently
                    # If new file size >= current position, try to continue from current position
                    # Otherwise start from beginning
                    if new_size >= current_pos:
                        new_file.seek(current_pos)
                        if self.debug:
                            print(f"File was replaced (inode changed), continuing from position {current_pos}")
                    else:
                        new_file.seek(0)
                        self._last_known_position = 0
                        if self.debug:
                            print(f"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)
                    if self.debug:
                        print(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
                    if self.debug:
                        print(f"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
                try:
                    old_file.close()
                except:
                    pass  # Old file might be invalid
                
                if self.debug:
                    print(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()
                if self.debug:
                    print(f"Failed to reopen source file: {e}")
                    
        except Exception as e:
            if self.debug:
                print(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()
                                if self.debug:
                                    print("Channel is now ready - first read attempt completed")
                            break
                        
                        line = line.strip()
                        if not line:
                            continue
                        
                        lines_read = True
                        
                        try:
                            task = json.loads(line)
                            task_id = task.get('task_id')
                            if not 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()
                                if self.debug:
                                    print("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())
                            if self.debug:
                                print("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:
                            if self.debug:
                                print("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:
                                if self.debug:
                                    print("File watcher event triggered, checking for updates...")
                            
                            # Check file updates regardless of events (forced polling)
                            if self._check_file_updated():
                                if self.debug:
                                    print("File updated detected, continuing to read...")
                                continue  # Has updates, continue reading
                            else:
                                if self.debug and event_triggered:
                                    print("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():
                                if self.debug:
                                    print("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:
                import traceback
                traceback.print_exc()
                print(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):
                if self.debug:
                    print("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
            
            if self.debug:
                print(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:
                if self.debug:
                    print(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:
                if self.debug:
                    print(f"Source file was truncated, restarting from beginning")
                try:
                    self.source_file.seek(0)
                    self._last_known_position = 0
                except OSError:
                    # If seek also fails, reopen file
                    self._reopen_source_file()
                return True
            
            if self.debug:
                print("No file update detected")
            return False
            
        except OSError as e:
            if self.debug:
                print(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():
                    try:
                        # Get result from sink queue with timeout to avoid infinite waiting
                        result = self.sink_queue.get(timeout=0.1)
                        
                        # 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._condition.notify()  # Notify source thread to continue reading
                        
                        self.sink_queue.task_done()
                        
                    except:
                        # Timeout or other exceptions, continue loop
                        continue
                        
            except Exception as e:
                print(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._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"""
        self._stop_event.set()
        self._file_updated_event.set()  # Wake up waiting source thread
        
        # 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)
        
        super().close()
        self.source_file.close()
        self.sink_file.close()
    
    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):
        
        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.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._pending_results = {}  # 存储待返回的结果
        self._last_task_time = None  # 记录最后一次接收到任务的时间
        self._eof_sent = False  # 是否已发送 EOF 标记
        
        # 设置 Celery 消费者
        self._setup_celery_consumer()
        self.start_prepare_source()
        self.start_prepare_sink()
    
    def _setup_celery_consumer(self):
        """设置 Celery 消费者配置"""
        # 从环境配置获取 Celery 配置
        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
            }
        })
        
        if self.debug:
            print(f"Celery channel setup for operator: {self.operator_name} v{self.operator_version}")
            print(f"Task name: {self.task_name}")
            print(f"Celery app: {self.celery_app.main}")
            print(f"Broker: {env_celery_config.broker_url}")
            print(f"Backend: {env_celery_config.result_backend}")
            print(f"Queues: {', '.join(self.queue_names)}")
            print(f"Priority order: {' > '.join(self.priorities)}")
            print(f"Visibility timeout: {self.visibility_timeout} seconds")
    
    def start_prepare_source(self):
        """从 Celery 队列接收任务的线程"""
        def _source_worker():
            try:
                # 设置 Celery 消费者
                with self.celery_app.connection() as connection:
                    from kombu import Queue
                    
                    # 创建多个优先级队列对象，按优先级顺序排列
                    # 为每个队列设置 visibility_timeout
                    queues = []
                    for queue_name in self.queue_names:
                        queue = Queue(
                            queue_name,
                            queue_arguments={
                                'x-visibility-timeout': self.visibility_timeout
                            }
                        )
                        queues.append(queue)
                    
                    if self.debug:
                        print(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)):
                            print(f"  {i+1}. {priority}: {queue_name}")
                    
                    def process_task(body, message):
                        """处理接收到的 Celery 任务"""
                        try:
                            if self.debug:
                                print(f"Received Celery task: {body}")
                                print(f"Message properties: {message.properties}")
                            
                            # 获取 Celery 任务 ID 和优先级信息
                            routing_key = message.delivery_info.get('routing_key', 'unknown')
                            celery_task_id = routing_key + '_' + str(id(message))
                            
                            # 从 routing_key 中提取优先级
                            task_priority = "normal"  # 默认优先级
                            for priority in self.priorities:
                                if routing_key.startswith(f"{priority}."):
                                    task_priority = priority
                                    break
                            
                            if self.debug:
                                print(f"Generated task ID: {celery_task_id}")
                                print(f"Task priority: {task_priority} (from queue: {routing_key})")
                            
                            # 解析任务数据 - Celery 任务数据通常在 body 中
                            if self.debug:
                                print(f"Raw body type: {type(body)}")
                                print(f"Raw body content: {body}")
                            
                            # 处理不同的 Celery 消息格式
                            task_data = {}
                            if isinstance(body, dict):
                                # 直接是字典格式
                                task_data = body.copy()
                            elif isinstance(body, (list, tuple)) and len(body) >= 2:
                                # Celery 标准格式: [args, kwargs, ...]
                                args, kwargs = body[0], body[1]
                                if isinstance(kwargs, dict):
                                    task_data = kwargs.copy()
                                elif isinstance(args, list) and len(args) > 0 and isinstance(args[0], dict):
                                    task_data = args[0].copy()
                            
                            # 确保有必需的字段
                            if task_data.get("task_type") != 'desc':
                                if 'artifacts' not in task_data:
                                    task_data['artifacts'] = []
                                if 'options' not in task_data:
                                    task_data['options'] = {}
                            
                            # 设置 task_id 为 Celery 任务 ID，并添加优先级信息
                            task_data['task_id'] = celery_task_id
                            task_data['priority'] = task_priority
                            
                            if self.debug:
                                print(f"Using Celery task ID: {celery_task_id}")
                            
                            # 更新最后任务时间
                            self._last_task_time = time.time()
                            
                            # 等待直到进程映射大小小于最大预取大小
                            with self._condition:
                                while len(self._current_process_map) >= self.max_prefetch_size:
                                    if self._stop_event.is_set():
                                        return
                                    self._condition.wait()
                                
                                # 使用 Celery task_id 作为键
                                self._current_process_map[celery_task_id] = {
                                    'task_data': task_data,
                                    'message': message,
                                    'celery_task_id': celery_task_id
                                }
                            
                            # 将任务放入源队列
                            self.source_queue.put(task_data)
                            
                            # 标记 channel 为就绪
                            if not self._ready_event.is_set():
                                self._ready_event.set()
                                if self.debug:
                                    print("Celery channel is now ready")
                            
                        except Exception as e:
                            if self.debug:
                                print(f"Error processing Celery task: {e}")
                            message.reject()
                    
                    # 创建消费者并设置多个优先级队列
                    # 队列按优先级顺序排列，Celery 会优先处理前面的队列
                    consumer = connection.Consumer(queues=queues, callbacks=[process_task], no_ack=False)
                    
                    # 启动消费
                    consumer.consume()
                    
                    # 初始化最后任务时间
                    self._last_task_time = time.time()
                    
                    # 开始消费循环
                    while not self._stop_event.is_set():
                        try:
                            # 在 job_mode 下检查超时
                            if self.job_mode and not self._eof_sent:
                                current_time = time.time()
                                if (self._last_task_time and 
                                    current_time - self._last_task_time > self.job_mode_timeout):
                                    
                                    if self.debug:
                                        print(f"Job mode timeout: {self.job_mode_timeout}s elapsed since last task")
                                    
                                    # 发送 EOF 标记
                                    self._eof_sent = True
                                    eof_marker = EOFMarker()
                                    self.source_queue.put(eof_marker.model_dump())
                                    
                                    if self.debug:
                                        print("Job mode: Sent EOF marker due to timeout")
                                    break
                            
                            # 使用较短的超时来定期检查 job_mode 条件
                            timeout = 1.0 if self.job_mode else 5.0
                            connection.drain_events(timeout=timeout)
                            
                        except Exception as e:
                            # 只在非超时错误且有实际错误信息时才打印
                            error_str = str(e).strip()
                            if (self.debug and 
                                "timed out" not in error_str.lower() and 
                                "timeout" not in error_str.lower() and
                                error_str):  # 确保错误信息不为空
                                print(f"Celery drain_events error: {error_str}")
                            
                            # 在 job_mode 下，即使是超时也要检查是否需要发送 EOF
                            if self.job_mode and not self._eof_sent:
                                current_time = time.time()
                                if (self._last_task_time and 
                                    current_time - self._last_task_time > self.job_mode_timeout):
                                    
                                    if self.debug:
                                        print(f"Job mode timeout during drain_events: {self.job_mode_timeout}s elapsed")
                                    
                                    # 发送 EOF 标记
                                    self._eof_sent = True
                                    eof_marker = EOFMarker()
                                    self.source_queue.put(eof_marker.model_dump())
                                    
                                    if self.debug:
                                        print("Job mode: Sent EOF marker due to timeout in drain_events")
                                    break
                            
                            continue
                    
            except Exception as e:
                import traceback
                traceback.print_exc()
                print(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():
                    try:
                        # 从 sink 队列获取结果
                        result = self.sink_queue.get(timeout=0.1)
                        
                        celery_task_id = result.get('task_id')
                        if not celery_task_id:
                            self.sink_queue.task_done()
                            continue
                        
                        # 获取对应的 Celery 消息
                        with self._condition:
                            process_info = self._current_process_map.pop(celery_task_id, None)
                            self._condition.notify()  # 通知源线程可以继续接收任务
                        
                        if process_info:
                            message = process_info['message']
                            
                            # 将结果发送回 Celery
                            try:
                                # 获取原始的 Celery correlation_id (客户端任务ID)
                                original_task_id = message.properties.get('correlation_id')
                                
                                if self.debug:
                                    print(f"Processed Celery task {celery_task_id}")
                                    print(f"Original task ID: {original_task_id}")
                                    print(f"Result status: {result.get('status', 'unknown')}")
                                
                                # 将结果存储到 Celery 的 result backend
                                if original_task_id:
                                    # 使用 Celery 的 AsyncResult 来存储结果
                                    from celery.result import AsyncResult
                                    async_result = AsyncResult(original_task_id, app=self.celery_app)
                                    
                                    # 构造 Celery 期望的结果格式
                                    if result.get('status') == 'success':
                                        celery_result = result
                                        if not result.pop('__keep_raw_output__', None):
                                            celery_result['artifacts'] = result.get('artifacts', [])
                                            celery_result['results'] = result.get('results', {})
                                            celery_result['error'] = result.get('error', None)
                                        celery_status = 'SUCCESS'
                                    else:
                                        # 处理错误情况，确保符合 Celery 的错误格式要求
                                        error_info = result.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 要求的异常对象
                                        celery_result = Exception(f"{error_type}: {error_message}")
                                        celery_status = 'FAILURE'
                                    
                                    # 存储结果到 backend
                                    self.celery_app.backend.store_result(
                                        original_task_id,
                                        celery_result,
                                        celery_status
                                    )
                                    
                                    if self.debug:
                                        print(f"Stored result for task {original_task_id}")
                                
                                # 确认消息处理完成
                                message.ack()
                                
                            except Exception as e:
                                if self.debug:
                                    print(f"Error sending result for task {celery_task_id}: {e}")
                                message.reject()
                        
                        self.sink_queue.task_done()
                        
                    except:
                        # 超时或其他异常，继续循环
                        continue
                        
            except Exception as e:
                print(f"Celery sink worker error: {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)
        
        # 在 job_mode 下，立即从进程映射中移除任务，因为不依赖 sink worker
        if self.job_mode:
            celery_task_id = task.get('task_id')
            if celery_task_id:
                with self._condition:
                    process_info = self._current_process_map.pop(celery_task_id, None)
                    if process_info and self.debug:
                        print(f"Job mode: Removed Celery task {celery_task_id}")
                    self._condition.notify()
        else:
            # 在正常模式下，只通知源线程，清理由 sink worker 处理
            with self._condition:
                self._condition.notify()
        
        return task
    
    def graceful_shutdown(self):
        """优雅关闭，释放所有未完成的任务"""
        if self.debug:
            print("Starting graceful shutdown of CeleryChannel...")
        
        # 停止接收新任务
        self._stop_event.set()
        
        # 释放所有未完成的任务
        with self._condition:
            rejected_count = 0
            for task_id, process_info in list(self._current_process_map.items()):
                try:
                    message = process_info['message']
                    message.reject()  # 立即释放任务供其他进程处理
                    rejected_count += 1
                    if self.debug:
                        print(f"Rejected task {task_id} for reprocessing")
                except Exception as e:
                    if self.debug:
                        print(f"Error rejecting task {task_id}: {e}")
            
            # 清空进程映射
            self._current_process_map.clear()
            self._condition.notify_all()
            
            if self.debug:
                print(f"Graceful shutdown completed: rejected {rejected_count} tasks")
    
    def close(self):
        """关闭 channel 并停止所有线程"""
        # 先执行优雅关闭
        if not self._stop_event.is_set():
            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)
        
        super().close()
    
    def is_ready(self):
        """检查 channel 是否就绪"""
        return self._ready_event.is_set()
    
    def wait_ready(self, timeout=None):
        """等待 channel 就绪"""
        return self._ready_event.wait(timeout)