"""
Sandai Operator SDK - Processor implementation

High-performance, asynchronous batch processing operator for building scalable data processing services.
"""

import argparse
import shutil
import time
import signal
import sys
import os
import copy
import asyncio
import uuid
import traceback
from queue import Empty
from pathlib import Path
from typing import (
    Dict,
    List,
    Optional,
    Set,
    Generator,
    Callable,
    Union,
    Any,
    TypeVar,
    get_type_hints,
    get_origin,
    get_args,
    Tuple,
    Iterable,
)
from collections.abc import Generator as ABCGenerator, Iterable as ABCIterable
import requests
from dataclasses import dataclass
import threading
from concurrent.futures import ThreadPoolExecutor
import logging
from urllib.parse import urlparse
import redis
import sentry_sdk
from sentry_sdk.integrations.dedupe import DedupeIntegration

from .channel import BaseChannel, FileChannel, CeleryChannel
from .types import (
    TaskInput,
    TaskOutput,
    TaskError,
    InputArtifact,
    OutputArtifact,
    OptionsType,
    ResultsType,
    EOFMarker,
    ProcessingContext,
)
from pydantic import BaseModel
from .env import (
    get_sentry_dsn,
    get_environment,
    is_development,
    get_celery_config,
    get_default_celery_config,
    get_temp_dir,
    get_log_level,
    get_visibility_timeout,
    get_celery_result_expires,
    get_celery_max_delivery_attempts,
)
from .storage import create_default_storage_provider


class StopException(Exception): ...


class OperatorConfigurationError(ValueError):
    """Raised when an operator is declared with an invalid runtime signature."""


class InputConversionError(ValueError):
    """Raised when a raw task cannot be converted into a TaskInput."""


class OutputValidationError(ValueError):
    """Raised when a user batch function yields invalid output."""


class RuntimeRequirementError(RuntimeError):
    """Retained for backward compatibility with older SDK releases."""


class ProcessingContractError(RuntimeError):
    """Raised when internal pipeline contracts are violated."""


def _origin_matches(annotation: Any, expected: Any) -> bool:
    origin = get_origin(annotation)
    if origin is None:
        return False

    if origin is expected:
        return True

    origin_name = getattr(origin, "__qualname__", getattr(origin, "__name__", None))
    expected_name = getattr(
        expected, "__qualname__", getattr(expected, "__name__", None)
    )
    origin_module = getattr(origin, "__module__", None)
    expected_module = getattr(expected, "__module__", None)
    return origin_name == expected_name and origin_module == expected_module


def _matches_parametrized_model(annotation: Any, expected: Any) -> bool:
    if not isinstance(annotation, type):
        return False

    return any(base is expected for base in getattr(annotation, "__mro__", ()))


SignalHandler = Callable[[int, Optional[signal.Signals]], None]
ProcessFunc = Callable[
    [List[TaskInput[OptionsType]], Any, ProcessingContext],
    Generator[TaskOutput[ResultsType], None, None],
]
EventHandler = Callable[[Any], None]


@dataclass
class PreparedTask:
    task_input: TaskInput[Any]
    persistent_output: bool = False


@dataclass
class OutputEnvelope:
    output: Union[TaskOutput[Any], Dict[str, Any]]
    persistent_output: bool = False
    preformatted: bool = False


REMOTE_URI_SCHEMES = ("oss://", "s3://", "fake://", "http://", "https://")


class BatchProcessor:

    def __init__(
        self,
        name: str,
        version: str,
        *,
        max_concurrency=1,
        max_batch_size=1,
        max_prefetch_size=1,
        prepare_concurrency=None,
        output_concurrency=None,
        visibility_timeout=None,
        max_delivery_attempts=None,
        ignore_nongil_check=False,
    ):
        self.name = name
        self.version = version
        self.running = False
        self._stop_event = threading.Event()

        self.process_func: Optional[ProcessFunc] = None
        self.options_type: Optional[type] = None
        self.results_type: Optional[type] = None
        self.pre_start_handlers: List[EventHandler] = []
        self.post_stop_handlers: List[EventHandler] = []

        log_level = get_log_level()
        self.logger = logging.getLogger(__name__)
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
            )
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
            self.logger.setLevel(getattr(logging, log_level))

        self.config = {}
        self.development_mode = is_development()

        self.max_batch_size = max_batch_size
        self.max_concurrency = max_concurrency
        self.max_prefetch_size = max_prefetch_size
        self.prepare_concurrency = prepare_concurrency
        self.output_concurrency = output_concurrency
        if not ignore_nongil_check:
            self.check_non_gil()

        if visibility_timeout is not None:
            resolved_visibility_timeout = visibility_timeout
        else:
            resolved_visibility_timeout = get_visibility_timeout()

        if max_delivery_attempts is not None:
            resolved_max_delivery_attempts = max_delivery_attempts
        else:
            resolved_max_delivery_attempts = get_celery_max_delivery_attempts()

        self.config["channel.celery"] = {
            "visibility_timeout": resolved_visibility_timeout,
            "max_delivery_attempts": resolved_max_delivery_attempts,
        }

        self.temp_dir = get_temp_dir()
        self._temp_dir_path = None
        self._processing_context = None

        self._process_executor = None
        self._prepare_executor = None
        self._output_executor = None
        self._channel_executor = None
        self._function_description = None

        self._redis_client = None
        self._cancelled_keys = set()
        self._storage_provider = None
        self._get_storage_provider()

    def check_non_gil(self):
        return None

    def on_batch(
        self,
        *,
        max_prefetch_size=1,
        max_batch_size=1,
        max_concurrency=1,
        prepare_concurrency=None,
        output_concurrency=None,
        visibility_timeout=None,
        max_delivery_attempts=None,
    ):
        self.max_batch_size = max_batch_size
        self.max_prefetch_size = max_prefetch_size
        self.max_concurrency = max_concurrency
        self.prepare_concurrency = prepare_concurrency
        self.output_concurrency = output_concurrency

        if visibility_timeout is not None:
            if "channel.celery" not in self.config:
                self.config["channel.celery"] = {}
            self.config["channel.celery"]["visibility_timeout"] = visibility_timeout

        if max_delivery_attempts is not None:
            if "channel.celery" not in self.config:
                self.config["channel.celery"] = {}
            self.config["channel.celery"]["max_delivery_attempts"] = max_delivery_attempts

        def decorator(func: ProcessFunc):
            self.process_func = func
            self._extract_type_info(func)
            self._prepare_function_description(func)
            return func

        return decorator

    def _extract_type_info(self, func):
        type_hints = get_type_hints(func)
        self.logger.info("Type hints: %s", type_hints)

        batch_inputs_type = type_hints.get("batch_inputs")
        if batch_inputs_type is not None:
            self.options_type = self._extract_task_input_model(batch_inputs_type)

        return_type = type_hints.get("return")
        if return_type is not None:
            self.results_type = self._extract_task_output_model(return_type)

    def _extract_task_input_model(self, annotation: Any) -> Optional[type[BaseModel]]:
        container_origin = get_origin(annotation)
        if container_origin not in (list, List):
            raise OperatorConfigurationError(
                "batch_inputs must be annotated as list[TaskInput[OptionsModel]]"
            )

        args = get_args(annotation)
        if len(args) != 1:
            raise OperatorConfigurationError(
                "batch_inputs must contain exactly one TaskInput type parameter"
            )

        task_input_annotation = args[0]
        if _origin_matches(task_input_annotation, TaskInput):
            task_args = get_args(task_input_annotation)
        elif _matches_parametrized_model(task_input_annotation, TaskInput):
            task_args = get_args(task_input_annotation)
            if not task_args:
                metadata = getattr(task_input_annotation, "__pydantic_generic_metadata__", {})
                task_args = tuple(metadata.get("args", ()))
        else:
            raise OperatorConfigurationError(
                "batch_inputs must be annotated as list[TaskInput[OptionsModel]]"
            )

        if len(task_args) != 1:
            raise OperatorConfigurationError(
                "TaskInput annotation must include exactly one options model"
            )

        return self._validate_model_type(task_args[0], "options")

    def _extract_task_output_model(self, annotation: Any) -> Optional[type[BaseModel]]:
        task_output_annotation = self._unwrap_task_output_annotation(annotation)
        if task_output_annotation is None:
            raise OperatorConfigurationError(
                "process function return type must yield TaskOutput[ResultsModel]"
            )

        task_args = get_args(task_output_annotation)
        if not task_args and _matches_parametrized_model(task_output_annotation, TaskOutput):
            metadata = getattr(task_output_annotation, "__pydantic_generic_metadata__", {})
            task_args = tuple(metadata.get("args", ()))
        if len(task_args) != 1:
            raise OperatorConfigurationError(
                "TaskOutput annotation must include exactly one results model"
            )

        return self._validate_model_type(task_args[0], "results")

    def _unwrap_task_output_annotation(self, annotation: Any) -> Optional[Any]:
        if _origin_matches(annotation, TaskOutput) or _matches_parametrized_model(
            annotation, TaskOutput
        ):
            return annotation

        origin = get_origin(annotation)
        if origin is None:
            return None

        if origin in (Generator, Iterable, ABCGenerator, ABCIterable):
            args = get_args(annotation)
            if not args:
                return None
            candidate = args[0]
            if _origin_matches(candidate, TaskOutput) or _matches_parametrized_model(
                candidate, TaskOutput
            ):
                return candidate
            return None

        if origin in (list, List, tuple):
            args = get_args(annotation)
            if len(args) == 1 and (
                _origin_matches(args[0], TaskOutput)
                or _matches_parametrized_model(args[0], TaskOutput)
            ):
                return args[0]

        return None

    def _validate_model_type(
        self, annotation: Any, label: str
    ) -> Optional[type[BaseModel]]:
        if annotation in (Any, None):
            return None
        if not isinstance(annotation, type) or not issubclass(annotation, BaseModel):
            raise OperatorConfigurationError(
                f"{label} model must inherit from pydantic.BaseModel, got {annotation!r}"
            )
        return annotation

    def _prepare_function_description(self, func):
        self._function_description = self._build_function_description()
        if func and func.__doc__:
            self._function_description["doc"] = func.__doc__

        self.logger.info("Function description prepared for desc queries")

    def _build_function_description(self) -> Dict[str, Any]:
        return {
            "doc": "No documentation available",
            "options_schema": self._build_model_schema(self.options_type, "Options"),
            "results_schema": self._build_model_schema(self.results_type, "Results"),
        }

    def _build_model_schema(
        self, model_type: Optional[type[BaseModel]], fallback_name: str
    ) -> Optional[Dict[str, Any]]:
        if model_type is None:
            return None

        try:
            return model_type.model_json_schema()
        except Exception as exc:
            self.logger.warning("Failed to generate %s schema: %s", fallback_name, exc)
            return {
                "type": "object",
                "description": f"{fallback_name} schema generation failed",
            }

    def _handle_desc_query(self, raw_input: Dict) -> Dict:
        if self._function_description:
            desc_output = copy.deepcopy(self._function_description)
        else:
            desc_output = {
                "doc": "No documentation available",
                "options_schema": None,
                "results_schema": None,
                "error": "Function description not available",
            }
        desc_output["task_id"] = raw_input["task_id"]
        desc_output["status"] = "success"
        desc_output["__keep_raw_output__"] = True
        return desc_output

    def _should_skip_due_to_deadline(self, raw_input: Dict) -> bool:
        if "deadline" not in raw_input or not raw_input["deadline"]:
            return False

        current_time = time.time()
        deadline = raw_input.get("deadline")

        if deadline is not None and current_time > deadline:
            self.logger.warning(
                f"Task {raw_input.get('task_id', 'unknown')} deadline exceeded before processing: {deadline} < {current_time} (exceeded by {current_time - deadline:.2f}s)"
            )
            return True

        return False

    def _create_deadline_error_output(self, raw_input: Dict) -> Dict:
        current_time = time.time()
        deadline = raw_input.get("deadline", 0)
        task_id = raw_input.get("task_id", "unknown")

        error_output = {
            "task_id": task_id,
            "status": "error",
            "error": {
                "error_type": "TaskDeadlineExceededError",
                "error_message": f"Task deadline exceeded before processing: {deadline} < {current_time}",
                "error_details": {
                    "deadline": deadline,
                    "current_time": current_time,
                    "exceeded_by": current_time - deadline,
                    "skipped_before_processing": True,
                },
            },
            "artifacts": [],
            "results": None,
            "__keep_raw_output__": True,
        }

        return error_output

    def _should_skip_due_to_cancel(self, raw_input: Dict) -> bool:
        cancel_key = raw_input.get("cancel_key")
        if not cancel_key:
            return False

        if cancel_key in self._cancelled_keys:
            self.logger.warning(
                f"Task {raw_input.get('task_id', 'unknown')} cancelled (key {cancel_key} already known to be set)"
            )
            return True

        if self._redis_client and self._check_cancel_key_in_redis(cancel_key):
            self._cancelled_keys.add(cancel_key)
            self.logger.warning(
                f"Task {raw_input.get('task_id', 'unknown')} cancelled (key {cancel_key} found in Redis)"
            )
            return True

        return False

    def _check_cancel_key_in_redis(self, cancel_key: str) -> bool:
        if self._redis_client:
            return bool(self._redis_client.exists(cancel_key))
        return False

    def _create_cancel_error_output(self, raw_input: Dict) -> Dict:
        cancel_key = raw_input.get("cancel_key", "")
        task_id = raw_input.get("task_id", "unknown")

        error_output = {
            "task_id": task_id,
            "status": "error",
            "error": {
                "error_type": "TaskCancelledError",
                "error_message": f"Task was cancelled via cancel_key: {cancel_key}",
                "error_details": {
                    "cancel_key": cancel_key,
                    "skipped_before_processing": True,
                },
            },
            "artifacts": [],
            "results": None,
            "__keep_raw_output__": True,
        }

        return error_output

    def _setup_redis_client(self, channel):
        conf = channel.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.logger.warning(
                    "Redis cancellation check source %s is unavailable: %s (%s)",
                    label,
                    candidate_url,
                    exc,
                )
                continue

            self._redis_client = client
            self.logger.info(
                "Redis client setup successful from %s URL: %s", label, candidate_url
            )
            return

        self._redis_client = None
        self.logger.warning(
            "Redis cancellation checks disabled because neither result backend nor broker uses redis: backend=%s broker=%s",
            result_backend,
            broker_url,
        )

    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.logger.warning(
                        "Invalid Redis database in URL %s, falling back to db=0",
                        connection_url,
                    )

        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 pre_start(self, func: EventHandler):
        self.pre_start_handlers.append(func)
        return func

    def post_stop(self, func: EventHandler):
        self.post_stop_handlers.append(func)
        return func

    def _setup_temp_directory(self):
        if self.temp_dir:
            self._temp_dir_path = Path(self.temp_dir)
            self._temp_dir_path.mkdir(parents=True, exist_ok=True)
        else:
            self._temp_dir_path = Path("./var/artifacts")
            self._temp_dir_path.mkdir(parents=True, exist_ok=True)

        self._processing_context = ProcessingContext(self._temp_dir_path)

        self.logger.info(f"Temporary directory: {self._temp_dir_path}")
        return self._temp_dir_path

    def _cleanup_temp_directory(self):
        if self._temp_dir_path and self._temp_dir_path.exists():
            try:
                if self.temp_dir:
                    self.logger.info(
                        f"Preserved user-specified temporary directory: {self._temp_dir_path}"
                    )
                else:
                    self.logger.info(
                        f"Temporary directory preserved: {self._temp_dir_path}"
                    )

            except Exception as e:
                self.logger.error(f"Error accessing temporary directory: {e}")

    def _generate_temp_filename(
        self, original_uri: str, task_id: str, file_type: str = "inputs"
    ) -> str:
        if original_uri.startswith("file://"):
            parsed = urlparse(original_uri)
            original_path = parsed.path or parsed.netloc
        elif original_uri.startswith("http"):
            original_path = urlparse(original_uri).path
        else:
            original_path = original_uri

        ext = Path(original_path).suffix
        unique_filename = f"{uuid.uuid4()}{ext}"
        task_dir = self._temp_dir_path / task_id / file_type
        task_dir.mkdir(parents=True, exist_ok=True)
        temp_file_path = str(task_dir / unique_filename)
        self.logger.debug(f"Generated temp file path: {temp_file_path}")

        return temp_file_path

    def _get_storage_provider(self):
        if self._storage_provider is None:
            # 从配置中获取 OSS 签名 URL 过期时间
            oss_signed_url_expires = self.config.get("storage.oss.signed_url_expires")

            # 创建存储提供者，传递 OSS 配置
            kwargs = {}
            if oss_signed_url_expires is not None:
                kwargs["signed_url_expires"] = oss_signed_url_expires
            
            self._storage_provider = create_default_storage_provider(
                logger=self.logger, **kwargs
            )
            self.logger.info(
                f"Initialized storage provider: {self._storage_provider.storage_type}"
            )

            # 如果是 OSS 存储，记录签名 URL 过期时间
            if hasattr(self._storage_provider, "signed_url_expires"):
                self.logger.info(
                    f"OSS signed URL expires in: {self._storage_provider.signed_url_expires} seconds"
                )

        return self._storage_provider

    async def _upload_to_storage(self, local_file_path: str, ephemeral: bool=True) -> Tuple[str, Dict]:
        storage_provider = self._get_storage_provider()
        return await storage_provider.upload_file(local_file_path, ephemeral=ephemeral)

    async def _download_to_path(self, uri: str, temp_file_path: str) -> str:
        storage_provider = self._get_storage_provider()
        return await storage_provider.download_file(uri, temp_file_path)

    def _cleanup_local_file(self, local_file_path: str):
        try:
            file_path = Path(local_file_path)
            if file_path.exists():
                file_path.unlink()
                self.logger.debug(f"Cleaned up local file: {local_file_path}")
        except Exception as e:
            self.logger.warning(f"Failed to cleanup local file {local_file_path}: {e}")

    def _cleanup_task_directory(self, task_id: str):
        try:
            task_dir = self._temp_dir_path / task_id
            if task_dir.exists():
                shutil.rmtree(task_dir, ignore_errors=True)
                self.logger.debug(f"Cleaned up task directory: {task_dir}")

        except Exception as e:
            self.logger.warning(f"Failed to cleanup task directory for {task_id}: {e}")

    def _uri_to_path(self, uri: str, task_id: str, file_type: str = "inputs") -> str:

        if not self._temp_dir_path:
            raise RuntimeError(
                "Temporary directory not initialized. Call _setup_temp_directory() first."
            )

        temp_file_path = self._generate_temp_filename(uri, task_id, file_type)

        try:
            if uri.startswith("file://"):
                parsed = urlparse(uri)

                if parsed.netloc and parsed.path:
                    relative_path = parsed.netloc + "/" + parsed.path.lstrip("/")
                    source_path = str(Path.cwd() / relative_path)
                elif parsed.netloc == "" and parsed.path.startswith("/"):
                    source_path = parsed.path
                else:
                    source_path = parsed.path or parsed.netloc
                    if not Path(source_path).is_absolute():
                        source_path = str(Path.cwd() / source_path)

                shutil.copy2(source_path, temp_file_path)
                self.logger.debug(
                    f"Copied local file {source_path} to {temp_file_path}"
                )

            elif uri.startswith(("http://", "https://")):
                self.logger.debug(f"Downloading {uri} to {temp_file_path}")
                response = requests.get(uri, stream=True)
                response.raise_for_status()

                with open(temp_file_path, "wb") as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)

                self.logger.debug(f"Downloaded {uri} to {temp_file_path}")
            elif uri.startswith(("oss://", "s3://", "fake://")):
                self.logger.debug(f"Downloading storage file {uri} to {temp_file_path}")
                asyncio.run(self._download_to_path(uri, temp_file_path))
                self.logger.debug(f"Downloaded storage file {uri} to {temp_file_path}")
            else:
                if not Path(uri).is_absolute():
                    source_path = str(Path.cwd() / uri)
                else:
                    source_path = uri

                shutil.copy2(source_path, temp_file_path)
                self.logger.debug(
                    f"Copied local file {source_path} to {temp_file_path}"
                )

            return temp_file_path

        except Exception as e:
            if Path(temp_file_path).exists():
                try:
                    Path(temp_file_path).unlink()
                except:
                    pass

            self.logger.error(f"Failed to process URI {uri}: {e}")
            raise

    async def convert_input_to_taskinput(self, input: Dict) -> TaskInput[Any]:
        normalized_input = copy.deepcopy(input)
        task_id = normalized_input.get("task_id")
        if not task_id:
            raise InputConversionError("Missing required field 'task_id' in input")

        artifacts = []
        for artifact in normalized_input.get("artifacts") or []:
            if not isinstance(artifact, dict):
                raise InputConversionError(
                    f"Artifact for task {task_id} must be a dict, got {type(artifact)!r}"
                )
            if not artifact.get("uri"):
                raise InputConversionError(
                    f"Artifact for task {task_id} is missing required field 'uri'"
                )

            converted_artifact = dict(artifact)
            try:
                converted_artifact["uri"] = await asyncio.get_event_loop().run_in_executor(
                    self._prepare_executor,
                    self._uri_to_path,
                    converted_artifact["uri"],
                    task_id,
                    "inputs",
                )
            except Exception as exc:
                raise InputConversionError(
                    f"Failed to convert artifact URI {artifact.get('uri')}: {exc}"
                ) from exc

            artifacts.append(InputArtifact(**converted_artifact))

        normalized_input["artifacts"] = artifacts

        raw_options = normalized_input.get("options", {})
        if raw_options is None:
            raw_options = {}

        if self.options_type:
            if not isinstance(raw_options, dict):
                raise InputConversionError(
                    f"Options for task {task_id} must be a dict compatible with {self.options_type.__name__}"
                )
            try:
                normalized_input["options"] = self.options_type(**raw_options)
            except Exception as exc:
                raise InputConversionError(
                    f"Failed to convert options to {self.options_type.__name__}: {exc}"
                ) from exc
        else:
            normalized_input["options"] = raw_options

        try:
            return TaskInput[Any](**normalized_input)
        except Exception as exc:
            raise InputConversionError(f"Failed to convert input to TaskInput: {exc}") from exc

    def create_error_output(
        self,
        task_id: str,
        error_type: str,
        error_message: str,
        error_details: Optional[Dict[str, Any]] = None,
        traceback_str: Optional[str] = None,
    ) -> TaskOutput[Any]:
        return TaskOutput[Any](
            task_id=task_id,
            status="error",
            error=TaskError(
                error_type=error_type,
                error_message=error_message,
                error_details=error_details,
                traceback=traceback_str,
            ),
        )

    def _validate_process_output(
        self, output: Any, expected_task_ids: Set[str]
    ) -> TaskOutput[Any]:
        if not isinstance(output, TaskOutput):
            raise OutputValidationError(
                f"Process function must yield TaskOutput instances, got {type(output)!r}"
            )
        if output.task_id not in expected_task_ids:
            raise OutputValidationError(
                f"Process function yielded unknown task_id {output.task_id!r}"
            )
        return output

    def _is_local_artifact_path(self, uri: str) -> bool:
        if not uri or uri.startswith(REMOTE_URI_SCHEMES):
            return False
        return Path(uri).exists()

    async def convert_taskoutput_to_output(self, output: TaskOutput[Any], ephemeral: bool=True):
        upload_errors = []

        for artifact in output.artifacts or []:
            local_file_path = artifact.uri

            try:
                if self._is_local_artifact_path(local_file_path):
                    uri, meta = await self._upload_to_storage(local_file_path, ephemeral)
                    artifact.uri = uri
                    artifact.meta = {**(artifact.meta or {}), "storage": meta}

                    preserve_local = self.config.get(
                        "channel.file.preserve_local_files", False
                    )
                    if preserve_local:
                        self.logger.debug(
                            f"Preserving local file (channel.file.preserve_local_files=True): {local_file_path}"
                        )
                    else:
                        await asyncio.get_event_loop().run_in_executor(
                            self._output_executor,
                            self._cleanup_local_file,
                            local_file_path,
                        )
                        self.logger.debug(
                            f"Uploaded and cleaned up: {local_file_path} -> {uri}"
                        )
                else:
                    self.logger.debug(
                        f"Artifact URI is not a local file, keeping as-is: {local_file_path}"
                    )
            except Exception as e:
                error_msg = f"Failed to upload artifact {local_file_path}: {e}"
                self.logger.error(error_msg)
                upload_errors.append(error_msg)
                if self._stop_event.is_set():
                    raise StopException(error_msg)

        if upload_errors:
            output.status = "error"
            output.error = TaskError(
                error_type="UploadError",
                error_message=f"Failed to upload {len(upload_errors)} artifact(s)",
                error_details={
                    "failed_uploads": upload_errors,
                    "total_artifacts": len(output.artifacts or []),
                },
            )
            output.artifacts = []

        task_id = output.task_id
        preserve_local = self.config.get("channel.file.preserve_local_files", False)
        if not preserve_local:
            await asyncio.get_event_loop().run_in_executor(
                self._output_executor, self._cleanup_task_directory, task_id
            )
        else:
            self.logger.debug(
                f"Preserving task directory (channel.file.preserve_local_files=True): {self._temp_dir_path / task_id}"
            )

        return await self._create_simple_output(output)

    async def process_batch(
        self, task_inputs: List[TaskInput[Any]]
    ) -> List[TaskOutput[Any]]:
        operator_config = self.config.get("operator.config", {})

        try:
            loop = asyncio.get_event_loop()
            return await loop.run_in_executor(
                self._process_executor,
                lambda: self._safe_process_func(task_inputs, operator_config),
            )
        except Exception as e:
            sentry_sdk.capture_exception(e)
            traceback_str = traceback.format_exc()
            self.logger.error(f"Batch processing failed completely: {e}")

            error_outputs = []
            for task_input in task_inputs:
                error_output = self.create_error_output(
                    task_id=task_input.task_id,
                    error_type="BatchProcessingError",
                    error_message=str(e),
                    error_details={"batch_size": len(task_inputs)},
                    traceback_str=traceback_str,
                )
                error_outputs.append(error_output)

            return error_outputs

    def _safe_process_func(
        self, task_inputs: List[TaskInput[Any]], operator_config: Dict
    ) -> List[TaskOutput[Any]]:

        input_task_map = {task.task_id: task for task in task_inputs}
        expected_task_ids = set(input_task_map.keys())
        yielded_results_by_task_id: Dict[str, TaskOutput[Any]] = {}

        try:
            for output in self.process_func(
                task_inputs, operator_config, self._processing_context
            ):
                validated_output = self._validate_process_output(output, expected_task_ids)
                if validated_output.task_id in yielded_results_by_task_id:
                    raise OutputValidationError(
                        f"Process function yielded duplicate task_id {validated_output.task_id!r}"
                    )
                yielded_results_by_task_id[validated_output.task_id] = validated_output
        except Exception as batch_error:
            traceback_str = traceback.format_exc()
            self.logger.error(f"Batch processing failed: {batch_error}")
            self.logger.debug(f"Batch processing traceback: {traceback_str}")

            for task_input in task_inputs:
                if task_input.task_id not in yielded_results_by_task_id:
                    yielded_results_by_task_id[task_input.task_id] = self.create_error_output(
                        task_id=task_input.task_id,
                        error_type=type(batch_error).__name__,
                        error_message=f"Batch processing failed: {str(batch_error)}",
                        error_details=self._extract_task_context(task_input),
                        traceback_str=traceback_str,
                    )

        missing_task_ids = expected_task_ids - set(yielded_results_by_task_id.keys())
        for missing_task_id in missing_task_ids:
            task_input = input_task_map[missing_task_id]
            yielded_results_by_task_id[missing_task_id] = self.create_error_output(
                task_id=missing_task_id,
                error_type="NoResultError",
                error_message="Process function did not yield result for this task",
                error_details=self._extract_task_context(task_input),
            )
            self.logger.warning(
                f"Task {missing_task_id} had no corresponding output, created fallback error output"
            )

        yielded_results = [yielded_results_by_task_id[task_input.task_id] for task_input in task_inputs]

        if len(yielded_results_by_task_id) != len(task_inputs):
            self.logger.warning(
                f"Result count mismatch: expected {len(task_inputs)}, got {len(yielded_results_by_task_id)}"
            )

        return yielded_results

    def _extract_task_context(self, task_input: TaskInput[Any]) -> Dict[str, Any]:
        error_details = {
            "task_id": task_input.task_id,
            "artifacts_count": len(task_input.artifacts) if task_input.artifacts else 0,
        }

        try:
            if hasattr(task_input.options, "model_dump"):
                error_details["options"] = task_input.options.model_dump()
            else:
                error_details["options"] = str(task_input.options)
        except Exception:
            error_details["options"] = "Failed to serialize options"

        try:
            if task_input.artifacts:
                error_details["artifacts"] = [
                    {"uri": artifact.uri, "slot": artifact.slot}
                    for artifact in task_input.artifacts[:3]  # 只记录前3个，避免过长
                ]
                if len(task_input.artifacts) > 3:
                    error_details["artifacts_truncated"] = (
                        f"... and {len(task_input.artifacts) - 3} more"
                    )
        except Exception:
            error_details["artifacts"] = "Failed to serialize artifacts"

        return error_details

    def run(
        self,
        signal_handler: Optional[SignalHandler] = None,
        job_mode: bool = False,
    ):
        sentry_dsn = get_sentry_dsn()
        if sentry_dsn:
            sentry_sdk.init(
                   dsn=sentry_dsn,
                   send_default_pii=True,
                   disabled_integrations=[DedupeIntegration()],
            )
        asyncio.run(self._async_run(signal_handler, job_mode))

    def _log_exception(self, message: str, exc: Exception):
        self.logger.error(f"{message}: {exc}", exc_info=True)

    async def _async_run(
        self, signal_handler: Optional[SignalHandler] = None, job_mode: bool = False
    ):
        if self.process_func is None:
            raise OperatorConfigurationError(
                "No process function defined, use @on_batch to define one"
            )

        operator_config = self.config.get("operator.config", {})
        self.logger.info(f"Using operator_config: {operator_config}")

        self.logger.info(f"Found {len(self.pre_start_handlers)} pre start handlers")
        for index, handler in enumerate(self.pre_start_handlers):
            self.logger.info(f"{index}) before pre start handler")
            handler(operator_config)
            self.logger.info(f"{index}) after pre start handler")

        self.running = True
        self._stop_event.clear()
        channel = None

        def enhanced_signal_handler(signum, frame):
            self.logger.warning(
                f"\nReceived signal {signum}, initiating graceful shutdown..."
            )
            self._stop_event.set()

        signal_handler = signal_handler or enhanced_signal_handler
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

        self._setup_temp_directory()

        self._process_executor = ThreadPoolExecutor(
            max_workers=self.max_concurrency,
            thread_name_prefix=f"{self.name}-executor-pool",
        )
        self._prepare_executor = ThreadPoolExecutor(
            max_workers=self._prepare_worker_count(),
            thread_name_prefix=f"{self.name}-prepare-pool",
        )
        self._output_executor = ThreadPoolExecutor(
            max_workers=self._output_worker_count(),
            thread_name_prefix=f"{self.name}-output-pool",
        )
        self._channel_executor = ThreadPoolExecutor(
            max_workers=self._channel_worker_count(),
            thread_name_prefix=f"{self.name}-channel-pool",
        )

        self.logger.info(f"Operator processing service started")
        self.logger.info(
            f"Config: max_batch_size={self.max_batch_size}, max_concurrency={self.max_concurrency}"
        )
        self.logger.info(
            "Created executors: compute=%s prepare=%s output=%s channel=%s",
            self.max_concurrency,
            self._prepare_worker_count(),
            self._output_worker_count(),
            self._channel_worker_count(),
        )

        default_channel = "file" if is_development() else "celery"
        channel_type = self.config.get("channel.type", default_channel)

        if channel_type == "celery":
            try:
                env_celery_config = get_celery_config()
                if not env_celery_config:
                    env_celery_config = get_default_celery_config()
                    self.logger.warning(
                        "Celery not configured in environment, using default configuration"
                    )

                user_celery_config = self.config.get("channel.celery", {})

                final_celery_config = {
                    "broker_url": user_celery_config.get(
                        "broker_url", env_celery_config.broker_url
                    ),
                    "result_backend": user_celery_config.get(
                        "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,
                    "priorities": user_celery_config.get(
                        "priorities", ["high", "normal", "low", "very_low"]
                    ),
                    "visibility_timeout": user_celery_config.get(
                        "visibility_timeout", get_visibility_timeout()
                    ),
                    "job_mode_timeout": user_celery_config.get("job_mode_timeout", 10),
                    "result_expires": user_celery_config.get(
                        "result_expires", get_celery_result_expires()
                    ),
                    "max_delivery_attempts": user_celery_config.get(
                        "max_delivery_attempts", get_celery_max_delivery_attempts()
                    ),
                }

                channel = CeleryChannel(
                    operator_name=self.name,
                    operator_version=self.version,
                    priorities=final_celery_config["priorities"],
                    visibility_timeout=final_celery_config["visibility_timeout"],
                    max_prefetch_size=self.max_concurrency * self.max_batch_size
                    + self.max_prefetch_size,
                    debug=self.development_mode,
                    job_mode=job_mode,
                    job_mode_timeout=final_celery_config["job_mode_timeout"],
                    result_expires=final_celery_config["result_expires"],
                    max_delivery_attempts=final_celery_config["max_delivery_attempts"],
                    broker_url=final_celery_config["broker_url"],
                    result_backend=final_celery_config["result_backend"],
                    task_serializer=final_celery_config["task_serializer"],
                    result_serializer=final_celery_config["result_serializer"],
                    accept_content=final_celery_config["accept_content"],
                )

                queue_names = [
                    f"{priority}.{self.name}.{self.version}"
                    for priority in final_celery_config["priorities"]
                ]
                self.logger.info(
                    f"Using Celery channel with priority queues: {', '.join(queue_names)}"
                )
                self.logger.info(f"Broker: {final_celery_config['broker_url']}")
                self.logger.info(f"Backend: {final_celery_config['result_backend']}")
                self.logger.info(
                    f"Visibility timeout: {final_celery_config['visibility_timeout']} seconds"
                )

                self._setup_redis_client(channel)
            except Exception as e:
                self.logger.error(f"Failed to create Celery channel: {e}")
                raise
        else:
            source_jsonl_path = self.config.get(
                "channel.file.input_path", "input.jsonl"
            )
            sink_jsonl_path = self.config.get(
                "channel.file.output_path", "output.jsonl"
            )
            channel = FileChannel(
                source_jsonl_path=source_jsonl_path,
                sink_jsonl_path=sink_jsonl_path,
                max_prefetch_size=self.max_concurrency * self.max_batch_size
                + self.max_prefetch_size,
                job_mode=job_mode,
            )
            self.logger.info(
                f"Using File channel: {source_jsonl_path} -> {sink_jsonl_path}"
            )

        input_queue = asyncio.Queue(
            maxsize=self.max_prefetch_size + self.max_concurrency + self.max_batch_size
        )

        try:
            if job_mode:
                await self._run_job_mode(channel, input_queue)
            else:
                await self._run_continuous_mode(channel, input_queue)
        except KeyboardInterrupt:
            self.logger.info("Received interrupt signal, stopping service...")
        except Exception as e:
            self._log_exception("Service runtime exception", e)
        finally:
            try:
                channel.sync()
                self.logger.info("Channel sync completed")
            except Exception as e:
                self.logger.error(f"Error during channel sync: {e}")
            try:
                channel.graceful_shutdown()
                self.logger.info("Channel graceful shutdown completed")
            except Exception as e:
                self.logger.error(f"Error during channel graceful shutdown: {e}")
            channel.close()
            self.stop()

    async def _run_continuous_mode(
        self, channel: BaseChannel, input_queue: asyncio.Queue
    ):
        await self._run_pipeline(
            channel,
            input_queue,
            asyncio.create_task(self._task_puller_loop(channel, input_queue)),
        )

    async def _run_job_mode(self, channel: BaseChannel, input_queue: asyncio.Queue):
        self.logger.info(
            "Running in job mode - will exit after processing all messages"
        )
        self.logger.info("Waiting for channel to be ready...")
        ready = await asyncio.get_event_loop().run_in_executor(
            self._channel_executor, lambda: channel.wait_ready(timeout=10.0)
        )
        if not ready:
            self.logger.warning("Channel not ready within timeout, proceeding anyway")
        else:
            self.logger.info("Channel is ready, starting to process tasks")

        await self._run_pipeline(
            channel,
            input_queue,
            asyncio.create_task(self._task_puller_job(channel, input_queue)),
        )
        self.logger.info("Job mode completed - all tasks processed")

    def _prepare_worker_count(self) -> int:
        return max(1, self.prepare_concurrency or self.max_concurrency)

    def _compute_worker_count(self) -> int:
        return max(1, self.max_concurrency)

    def _output_worker_count(self) -> int:
        return max(1, self.output_concurrency or self.max_concurrency)

    def _channel_worker_count(self) -> int:
        return max(1, self.max_concurrency)

    def _pipeline_queue_size(self) -> int:
        return self.max_prefetch_size + self.max_concurrency + self.max_batch_size

    async def _run_pipeline(
        self,
        channel: BaseChannel,
        input_queue: asyncio.Queue,
        input_task: asyncio.Task,
    ):
        prepared_queue = asyncio.Queue(maxsize=self._pipeline_queue_size())
        output_queue = asyncio.Queue(maxsize=self._pipeline_queue_size())

        prepare_worker_count = self._prepare_worker_count()
        compute_worker_count = self._compute_worker_count()
        output_worker_count = self._output_worker_count()

        prepare_tasks = [
            asyncio.create_task(
                self._prepare_task_loop(input_queue, prepared_queue, output_queue)
            )
            for _ in range(prepare_worker_count)
        ]
        compute_tasks = [
            asyncio.create_task(self._compute_batch_loop(prepared_queue, output_queue))
            for _ in range(compute_worker_count)
        ]
        output_tasks = [
            asyncio.create_task(self._output_loop(output_queue, channel))
            for _ in range(output_worker_count)
        ]

        try:
            await input_task
        finally:
            await input_queue.join()
            await self._stop_worker_stage(input_queue, prepare_tasks, "prepare")

            await prepared_queue.join()
            await self._stop_worker_stage(prepared_queue, compute_tasks, "compute")

            await output_queue.join()
            await self._stop_worker_stage(output_queue, output_tasks, "output")

    async def _stop_worker_stage(
        self,
        queue: asyncio.Queue,
        worker_tasks: List[asyncio.Task],
        stage_name: str,
    ):
        for _ in worker_tasks:
            await queue.put(None)

        results = await asyncio.gather(*worker_tasks, return_exceptions=True)
        for result in results:
            if isinstance(result, Exception) and not isinstance(result, StopException):
                self.logger.error(f"{stage_name} stage worker failed: {result}")

    async def _task_puller_job(self, channel: BaseChannel, input_queue: asyncio.Queue):
        while True:
            try:
                _input = await asyncio.get_event_loop().run_in_executor(
                    self._channel_executor, lambda: channel.pull(timeout=1.0)
                )

                if _input.get("marker_type") == "eof":
                    self.logger.info("Received EOF marker, stopping input reader")
                    break

                await input_queue.put(_input)
            except Empty:
                await asyncio.sleep(0.1)
            except Exception as e:
                self.logger.error(f"Error reading input: {e}")
                break

            if self._stop_event.is_set():
                raise StopException("force exit by stop event")

    async def _task_puller_loop(self, channel: BaseChannel, input_queue: asyncio.Queue):
        while self.running and not self._stop_event.is_set():
            try:
                _input = await asyncio.get_event_loop().run_in_executor(
                    self._channel_executor, lambda: channel.pull(timeout=1.0)
                )

                if _input.get("task_type") == "stop":
                    reason = _input.get("reason", "No reason provided")
                    self.logger.info(
                        f"Received Stop marker, initiating graceful shutdown. Reason: {reason}"
                    )
                    self._stop_event.set()

                if self._stop_event.is_set():
                    self.logger.info("Stop event detected, discarding pulled input")
                    break
                await input_queue.put(_input)
                self.logger.debug(f"Queued input: {_input.get('id', 'unknown')}")
            except Empty:
                await asyncio.sleep(0.1)
                continue
            except Exception as e:
                if not self._stop_event.is_set():
                    self.logger.error(f"Error pulling input: {e}")
                await asyncio.sleep(1.0)

        self.logger.info("Task puller stopped")
        if self._stop_event.is_set():
            raise StopException("force exit by stop event")

    async def _prepare_task_loop(
        self,
        input_queue: asyncio.Queue,
        prepared_queue: asyncio.Queue,
        output_queue: asyncio.Queue,
    ):
        while True:
            raw_input = None
            try:
                raw_input = await input_queue.get()
                if raw_input is None:
                    input_queue.task_done()
                    break

                output_envelope = None
                prepared_task = None

                if raw_input.get("task_type") == "desc":
                    desc_output = self._handle_desc_query(raw_input)
                    output_envelope = OutputEnvelope(desc_output, preformatted=True)
                    self.logger.info(
                        f"Directly handled desc query: {raw_input.get('task_id', 'unknown')}"
                    )
                elif self._should_skip_due_to_deadline(raw_input):
                    deadline_error_output = self._create_deadline_error_output(raw_input)
                    output_envelope = OutputEnvelope(
                        deadline_error_output, preformatted=True
                    )
                    self.logger.warning(
                        f"Skipped task {raw_input.get('task_id', 'unknown')} due to deadline"
                    )
                elif self._should_skip_due_to_cancel(raw_input):
                    cancel_error_output = self._create_cancel_error_output(raw_input)
                    output_envelope = OutputEnvelope(
                        cancel_error_output, preformatted=True
                    )
                    self.logger.warning(
                        f"Skipped task {raw_input.get('task_id', 'unknown')} due to cancellation"
                    )
                else:
                    try:
                        converted_input = await self.convert_input_to_taskinput(raw_input)
                    except InputConversionError as exc:
                        task_id = raw_input.get("task_id", "unknown")
                        error_output = self.create_error_output(
                            task_id=task_id,
                            error_type="InputConversionError",
                            error_message=str(exc),
                            error_details={"raw_input": raw_input},
                        )
                        output_envelope = OutputEnvelope(error_output)
                        self.logger.warning(
                            f"Created error output for failed input conversion: {task_id}"
                        )
                    else:
                        prepared_task = PreparedTask(
                            task_input=converted_input,
                            persistent_output=bool(
                                (converted_input.meta or {}).get(
                                    "persistent-output", False
                                )
                            ),
                        )

                if output_envelope is not None:
                    await output_queue.put([output_envelope])
                elif prepared_task is not None:
                    await prepared_queue.put(prepared_task)
            except Exception as e:
                self._log_exception("Error in prepare loop", e)
                await asyncio.sleep(1.0)
            finally:
                if raw_input is not None:
                    input_queue.task_done()

        self.logger.info("Prepare loop stopped")

    async def _compute_batch_loop(
        self,
        prepared_queue: asyncio.Queue,
        output_queue: asyncio.Queue,
    ):
        while True:
            batch_tasks = []

            try:
                prepared_task = await prepared_queue.get()
                if prepared_task is None:
                    prepared_queue.task_done()
                    break

                batch_tasks.append(prepared_task)

                while len(batch_tasks) < self.max_batch_size:
                    try:
                        prepared_task = await asyncio.wait_for(
                            prepared_queue.get(), timeout=1.0
                        )
                        if prepared_task is None:
                            prepared_queue.task_done()
                            break
                        batch_tasks.append(prepared_task)
                    except asyncio.TimeoutError:
                        break

                output_batch = await self._execute_batch(batch_tasks)
                if output_batch:
                    await output_queue.put(output_batch)
            except Exception as e:
                self._log_exception("Error in compute loop", e)
                await asyncio.sleep(1.0)
            finally:
                for _ in batch_tasks:
                    prepared_queue.task_done()

        self.logger.info("Compute loop stopped")

    async def _execute_batch(
        self, batch_tasks: List[PreparedTask]
    ) -> List[OutputEnvelope]:
        if not batch_tasks:
            return []

        batch_inputs = [batch_task.task_input for batch_task in batch_tasks]
        self.logger.info(f"Processing batch of {len(batch_inputs)} valid tasks")

        try:
            batch_outputs = await self.process_batch(batch_inputs)
            if len(batch_outputs) != len(batch_tasks):
                raise ProcessingContractError(
                    "process_batch returned a different number of outputs than inputs "
                    f"(expected {len(batch_tasks)}, got {len(batch_outputs)})"
                )
            output_batch = [
                OutputEnvelope(_output, batch_task.persistent_output)
                for batch_task, _output in zip(batch_tasks, batch_outputs)
            ]
            self.logger.info(
                f"Batch processing completed: {len(batch_outputs)} outputs"
            )
            return output_batch
        except Exception as e:
            traceback_str = traceback.format_exc()
            self.logger.error(f"Batch processing failed: {e}")
            return [
                OutputEnvelope(
                    self.create_error_output(
                        task_id=batch_task.task_input.task_id,
                        error_type="ProcessingError",
                        error_message=str(e),
                        error_details={"batch_size": len(batch_inputs)},
                        traceback_str=traceback_str,
                    )
                )
                for batch_task in batch_tasks
            ]

    async def _output_loop(
        self,
        output_queue: asyncio.Queue,
        channel: BaseChannel,
    ):
        while True:
            output_batch = None
            try:
                output_batch = await output_queue.get()
                if output_batch is None:
                    output_queue.task_done()
                    break

                await self._deliver_output_batch(output_batch, channel)
            except Exception as e:
                self._log_exception("Error in output loop", e)
                await asyncio.sleep(1.0)
            finally:
                if output_batch is not None:
                    output_queue.task_done()

        self.logger.info("Output loop stopped")

    async def _deliver_output_batch(
        self,
        output_batch: List[OutputEnvelope],
        channel: BaseChannel,
    ):
        if not output_batch:
            return

        final_outputs: List[Optional[Dict[str, Any]]] = [None] * len(output_batch)
        convert_positions = []
        convert_tasks = []

        for index, envelope in enumerate(output_batch):
            if envelope.preformatted:
                final_outputs[index] = envelope.output
            else:
                convert_positions.append(index)
                convert_tasks.append(
                    asyncio.create_task(
                        self.convert_taskoutput_to_output(
                            envelope.output, not envelope.persistent_output
                        )
                    )
                )

        converted_outputs = await asyncio.gather(*convert_tasks, return_exceptions=True)

        for index, result in zip(convert_positions, converted_outputs):
            envelope = output_batch[index]
            task_id = self._get_output_task_id(envelope.output)

            if isinstance(result, Exception):
                self.logger.error(
                    f"Output conversion failed for task {task_id}: {result}"
                )
                final_outputs[index] = {
                    "task_id": task_id,
                    "status": "error",
                    "error": {
                        "error_type": "OutputConversionError",
                        "error_message": str(result),
                    },
                }
            else:
                final_outputs[index] = result

        push_tasks = []
        for output in final_outputs:
            if output is None:
                continue
            push_tasks.append(
                asyncio.get_event_loop().run_in_executor(
                    self._channel_executor, channel.push, output
                )
            )

        push_results = await asyncio.gather(*push_tasks, return_exceptions=True)
        for result in push_results:
            if isinstance(result, Exception):
                self.logger.error(f"Failed to push output: {result}")

        self.logger.info(f"Completed processing {len(output_batch)} tasks")

    def _get_output_task_id(self, output: Union[TaskOutput[Any], Dict[str, Any]]) -> str:
        if isinstance(output, TaskOutput):
            return output.task_id
        return output.get("task_id", "unknown")

    async def _create_simple_output(self, task_output: TaskOutput[Any]) -> Dict:
        return {
            "task_id": task_output.task_id,
            "status": task_output.status,
            "results": (
                task_output.results.model_dump() if task_output.results else {}
            ),
            "error": task_output.error.model_dump() if task_output.error else None,
            "artifacts": [
                artifact.model_dump() for artifact in (task_output.artifacts or [])
            ],
        }

    def stop(self):
        self.logger.info("Initiating graceful shutdown...")
        self._stop_event.set()
        self.running = False

        operator_config = self.config.get("operator.config", {})
        self.logger.info(
            f"Calling {len(self.post_stop_handlers)} post_stop handlers with operator_config: {operator_config}"
        )
        for index, post_stop_handler in enumerate(self.post_stop_handlers):
            self.logger.info(f"{index}) before post_stop_handler call")
            try:
                post_stop_handler(operator_config)
                self.logger.debug(f"{index}) post_stop_handler called successfully")
            except Exception as e:
                self.logger.error(f"{index}) error when post_stop_handler called: {e}")

        if self._process_executor:
            self.logger.info("Shutting down process executor pool ...")
            try:
                self._process_executor.shutdown(wait=True, cancel_futures=True)
                self.logger.info("Executor process pool shutdown successfully")
            except Exception as e:
                self.logger.error(f"Error shutting down process executor pool: {e}")
            finally:
                self._process_executor = None

        if self._prepare_executor:
            self.logger.info("Shutting down prepare executor pool ...")
            try:
                self._prepare_executor.shutdown(wait=True, cancel_futures=True)
                self.logger.info("Prepare executor pool shutdown successfully")
            except Exception as e:
                self.logger.error(f"Error shutting down prepare executor pool: {e}")
            finally:
                self._prepare_executor = None

        if self._output_executor:
            self.logger.info("Shutting down output executor pool ...")
            try:
                self._output_executor.shutdown(wait=True, cancel_futures=True)
                self.logger.info("Output executor pool shutdown successfully")
            except Exception as e:
                self.logger.error(f"Error shutting down output executor pool: {e}")
            finally:
                self._output_executor = None

        if self._channel_executor:
            self.logger.info("Shutting down channel executor pool ...")
            try:
                self._channel_executor.shutdown(wait=True, cancel_futures=True)
                self.logger.info("Channel executor pool shutdown successfully")
            except Exception as e:
                self.logger.error(f"Error shutting down channel executor pool: {e}")
            finally:
                self._channel_executor = None

        self._cleanup_temp_directory()
        self.logger.info("Operator processor stopped gracefully")
