"""
Sandai Operator SDK - Processor implementation

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

import argparse
import json
import tempfile
import shutil
import time
import signal
import sys
import os
import inspect
import copy
import asyncio
import uuid
import traceback
from queue import Empty
from pathlib import Path
from typing import (
    Dict,
    List,
    Optional,
    Generator,
    Callable,
    Union,
    Any,
    TypeVar,
    get_type_hints,
    get_origin,
    get_args,
    Tuple,
)
import requests
from dataclasses import dataclass, asdict
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
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,
)
from .storage import create_default_storage_provider
from typing import TypeVar


class StopException(Exception): ...


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


class BatchProcessor:

    def __init__(
        self,
        name: str,
        version: str,
        *,
        max_concurrency=1,
        max_batch_size=1,
        max_prefetch_size=1,
        visibility_timeout=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
        if not ignore_nongil_check:
            self.check_non_gil()

        if visibility_timeout is not None:
            self.config["channel.celery"] = {"visibility_timeout": visibility_timeout}
        else:
            env_visibility_timeout = get_visibility_timeout()
            self.config["channel.celery"] = {
                "visibility_timeout": env_visibility_timeout
            }

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

        self._process_executor = None
        self._io_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):
        assert sys.version_info >= (3, 14), "Please use python >= 3.14"
        if hasattr(sys, "_is_gil_enabled"):
            assert not sys._is_gil_enabled(), "Please disable GIL for this operator"

    def on_batch(
        self,
        *,
        max_prefetch_size=1,
        max_batch_size=1,
        max_concurrency=1,
        visibility_timeout=None,
    ):
        self.max_batch_size = max_batch_size
        self.max_prefetch_size = max_prefetch_size
        self.max_concurrency = max_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

        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(f"Type hints: {type_hints}")

        for param_name, param_type in type_hints.items():
            type_str = str(param_type)
            self.logger.info(f"Parameter {param_name}: {type_str}")

            # Look for TaskInput[SomeType] pattern
            if "TaskInput[" in type_str and param_name == "batch_inputs":
                import re

                match = re.search(r"TaskInput\[([^]]+)\]", type_str)
                if match:
                    type_name = match.group(1)
                    self.logger.info(f"Found TaskInput type parameter: {type_name}")
                    # Try to get the actual type from the function's globals
                    func_globals = func.__globals__
                    if type_name in func_globals:
                        self.options_type = func_globals[type_name]
                        self.logger.info(f"Resolved options type: {self.options_type}")

            # Look for TaskOutput[SomeType] pattern in return type
            if param_name == "return" and "TaskOutput[" in type_str:
                import re

                match = re.search(r"TaskOutput\[([^]]+)\]", type_str)
                if match:
                    type_name = match.group(1)
                    self.logger.info(f"Found TaskOutput type parameter: {type_name}")
                    # Try to get the actual type from the function's globals
                    func_globals = func.__globals__
                    if type_name in func_globals:
                        self.results_type = func_globals[type_name]
                        self.logger.info(f"Resolved results type: {self.results_type}")

    def _prepare_function_description(self, func):
        doc = func.__doc__ if func else None

        options_schema = None
        if self.options_type:
            try:
                if hasattr(self.options_type, "model_json_schema"):
                    options_schema = self.options_type.model_json_schema()
                else:
                    options_schema = {
                        "type": "object",
                        "description": f"Options type: {self.options_type.__name__}",
                    }
            except Exception as e:
                self.logger.warning(f"Failed to generate options schema: {e}")
                options_schema = {
                    "type": "object",
                    "description": "Schema generation failed",
                }

        results_schema = None
        if self.results_type:
            try:
                if hasattr(self.results_type, "model_json_schema"):
                    results_schema = self.results_type.model_json_schema()
                else:
                    results_schema = {
                        "type": "object",
                        "description": f"Results type: {self.results_type.__name__}",
                    }
            except Exception as e:
                self.logger.warning(f"Failed to generate results schema: {e}")
                results_schema = {
                    "type": "object",
                    "description": "Schema generation failed",
                }

        self._function_description = {
            "doc": doc or "No documentation available",
            "options_schema": options_schema,
            "results_schema": results_schema,
        }

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

    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):
        broker_url = channel.celery_app.conf.broker_url
        parsed = urlparse(broker_url)

        self._redis_client = redis.Redis(
            host=parsed.hostname or "localhost",
            port=parsed.port or 6379,
            username=parsed.username or None,
            password=parsed.password or None,
            db=int(parsed.path.lstrip("/")) if parsed.path else 0,
            decode_responses=True,
        )
        self._redis_client.ping()
        self.logger.info(f"Redis client setup successful: {broker_url}")

    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(**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)

    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://"):
                self.logger.debug(f"Downloading OSS file {uri} to {temp_file_path}")
                storage_provider = self._get_storage_provider()
                # 在线程池中运行异步下载
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                try:
                    loop.run_until_complete(
                        storage_provider.download_file(uri, temp_file_path)
                    )
                finally:
                    loop.close()
                self.logger.debug(f"Downloaded OSS file {uri} to {temp_file_path}")

            elif uri.startswith("s3://"):
                self.logger.debug(f"Downloading S3 file {uri} to {temp_file_path}")
                storage_provider = self._get_storage_provider()
                # 在线程池中运行异步下载
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                try:
                    loop.run_until_complete(
                        storage_provider.download_file(uri, temp_file_path)
                    )
                finally:
                    loop.close()
                self.logger.debug(f"Downloaded S3 file {uri} to {temp_file_path}")

            elif uri.startswith("fake://"):
                self.logger.debug(f"Downloading fake file {uri} to {temp_file_path}")
                storage_provider = self._get_storage_provider()
                # 在线程池中运行异步下载
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                try:
                    loop.run_until_complete(
                        storage_provider.download_file(uri, temp_file_path)
                    )
                finally:
                    loop.close()
                self.logger.debug(f"Downloaded fake 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) -> Optional[TaskInput[Any]]:
        try:
            if not input.get("task_id"):
                self.logger.error(f"Missing required field 'task_id' in input: {input}")
                return None

            task_id = input["task_id"]
            for artifact in input.get("artifacts") or []:
                try:
                    artifact["uri"] = await asyncio.get_event_loop().run_in_executor(
                        self._io_executor,
                        self._uri_to_path,
                        artifact["uri"],
                        task_id,
                        "inputs",
                    )
                except Exception as e:
                    self.logger.error(
                        f"Failed to convert artifact URI {artifact.get('uri')}: {e}"
                    )
                    return None

            if self.options_type and "options" in input:
                try:
                    input["options"] = self.options_type(**input["options"])
                except Exception as e:
                    self.logger.error(
                        f"Failed to convert options to {self.options_type}: {e}"
                    )
                    return None

            return TaskInput[Any](**input)

        except Exception as e:
            self.logger.error(f"Failed to convert input to TaskInput: {e}")
            import traceback

            self.logger.error(f"Traceback: {traceback.format_exc()}")
            return None

    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,
            ),
        )

    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 Path(local_file_path).exists():
                    uri, meta = await self._upload_to_storage(local_file_path, ephemeral)
                    artifact.uri = uri
                    artifact.meta = {**(artifact.meta or {}), "storage": meta}

                    await asyncio.get_event_loop().run_in_executor(
                        self._io_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}"
                traceback.print_exc()
                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
        await asyncio.get_event_loop().run_in_executor(
            self._io_executor, self._cleanup_task_directory, task_id
        )

        return output.model_dump()

    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}
        yielded_results = []
        yielded_task_ids = set()

        try:
            for output in self.process_func(
                task_inputs, operator_config, self._processing_context
            ):
                yielded_results.append(output)
                yielded_task_ids.add(output.task_id)
        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_task_ids:
                    error_output = 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,
                    )
                    yielded_results.append(error_output)
                    yielded_task_ids.add(task_input.task_id)

        missing_task_ids = set(input_task_map.keys()) - yielded_task_ids
        for missing_task_id in missing_task_ids:
            task_input = input_task_map[missing_task_id]
            error_output = 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),
            )
            yielded_results.append(error_output)
            self.logger.warning(
                f"Task {missing_task_id} had no corresponding output, created fallback error output"
            )

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

        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))

    async def _async_run(
        self, signal_handler: Optional[SignalHandler] = None, job_mode: bool = False
    ):
        assert (
            self.process_func
        ), "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._io_executor = ThreadPoolExecutor(
            max_workers=self.max_concurrency * 2,
            thread_name_prefix=f"{self.name}-io-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(
            f"Created dedicated ThreadPoolExecutor with {self.max_concurrency} workers"
        )

        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),
                }

                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"],
                )

                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:
            traceback.print_exc()
            self.logger.error(f"Service runtime exception: {e}")
        finally:
            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
    ):
        input_task = asyncio.create_task(self._task_puller_loop(channel, input_queue))

        batch_tasks = [
            asyncio.create_task(self._batch_consume_loop(input_queue, channel))
            for _ in range(self.max_concurrency)
        ]

        all_tasks = [input_task] + batch_tasks
        await asyncio.gather(*all_tasks, return_exceptions=True)

    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._io_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")

        input_task = asyncio.create_task(self._task_puller_job(channel, input_queue))
        batch_tasks = [
            asyncio.create_task(self._batch_consume_job(input_queue, channel))
            for _ in range(self.max_concurrency)
        ]

        await input_task
        await asyncio.gather(*batch_tasks, return_exceptions=True)
        self.logger.info("Job mode completed - all tasks processed")

    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._io_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")

        for _ in range(self.max_concurrency):
            await input_queue.put(None)

    async def _batch_consume_job(
        self, input_queue: asyncio.Queue, channel: BaseChannel
    ):
        while True:
            try:
                raw_inputs = []
                convert_tasks = []
                should_exit = False
                fastpath_outputs = []
                effect_task_count = 0

                while effect_task_count < self.max_batch_size:
                    try:
                        raw_input = await input_queue.get()
                        if raw_input is None:
                            input_queue.task_done()
                            if convert_tasks or fastpath_outputs:
                                should_exit = True
                                break
                            else:
                                self.logger.info(
                                    "No more task in channel, exit consume job loop"
                                )
                                return

                        if raw_input.get("task_type") == "desc":
                            desc_output = self._handle_desc_query(raw_input)
                            fastpath_outputs.append(desc_output)
                            self.logger.info(
                                f"Directly handled desc query in job mode: {raw_input.get('task_id', 'unknown')} (not counted in batch)"
                            )
                        elif self._should_skip_due_to_deadline(raw_input):
                            deadline_error_output = self._create_deadline_error_output(
                                raw_input
                            )
                            fastpath_outputs.append(deadline_error_output)
                            self.logger.warning(
                                f"Skipped task {raw_input.get('task_id', 'unknown')} due to deadline in job mode (not counted in batch)"
                            )
                        elif self._should_skip_due_to_cancel(raw_input):
                            cancel_error_output = self._create_cancel_error_output(
                                raw_input
                            )
                            fastpath_outputs.append(cancel_error_output)
                            self.logger.warning(
                                f"Skipped task {raw_input.get('task_id', 'unknown')} due to cancellation in job mode (not counted in batch)"
                            )
                        else:
                            raw_inputs.append(raw_input)
                            convert_task = asyncio.create_task(
                                self.convert_input_to_taskinput(raw_input)
                            )
                            convert_tasks.append(convert_task)
                            effect_task_count += 1

                        input_queue.task_done()
                    except Exception as e:
                        self.logger.error(f"Error getting task from queue: {e}")
                        break

                push_tasks = []
                for fastpath_output in fastpath_outputs:
                    push_task = asyncio.get_event_loop().run_in_executor(
                        self._io_executor, channel.push, fastpath_output
                    )
                    push_tasks.append(push_task)

                if push_tasks:
                    await asyncio.gather(*push_tasks)

                converted_inputs = await asyncio.gather(*convert_tasks)

                valid_inputs = []
                failed_inputs = []
                for raw_input, converted_input in zip(raw_inputs, converted_inputs):
                    if converted_input is not None:
                        valid_inputs.append(converted_input)
                    else:
                        failed_inputs.append(raw_input)

                await self._process_batch(valid_inputs, failed_inputs, channel)

                if should_exit:
                    break
            except Exception as e:
                traceback.print_exc()
                self.logger.error(f"Error in job mode processor: {e}")
                break
        self.logger.info("Batch consume job stopped")

    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._io_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 _batch_consume_loop(
        self, input_queue: asyncio.Queue, channel: BaseChannel
    ):
        while self.running and not self._stop_event.is_set():
            try:
                raw_inputs = []
                convert_tasks = []
                fastpath_outputs = []

                effect_task_count = 0
                while (
                    effect_task_count < self.max_batch_size
                    and not self._stop_event.is_set()
                ):
                    try:
                        raw_input = await asyncio.wait_for(
                            input_queue.get(), timeout=1.0
                        )
                        if raw_input.get("task_type") == "desc":
                            desc_output = self._handle_desc_query(raw_input)
                            fastpath_outputs.append(desc_output)
                            self.logger.info(
                                f"Directly handled desc query: {raw_input.get('task_id', 'unknown')} (not counted in batch)"
                            )
                        elif self._should_skip_due_to_deadline(raw_input):
                            deadline_error_output = self._create_deadline_error_output(
                                raw_input
                            )
                            fastpath_outputs.append(deadline_error_output)
                            self.logger.warning(
                                f"Skipped task {raw_input.get('task_id', 'unknown')} due to deadline (not counted in batch)"
                            )
                        elif self._should_skip_due_to_cancel(raw_input):
                            cancel_error_output = self._create_cancel_error_output(
                                raw_input
                            )
                            fastpath_outputs.append(cancel_error_output)
                            self.logger.warning(
                                f"Skipped task {raw_input.get('task_id', 'unknown')} due to cancellation (not counted in batch)"
                            )
                        else:
                            raw_inputs.append(raw_input)
                            convert_task = asyncio.create_task(
                                self.convert_input_to_taskinput(raw_input)
                            )
                            convert_tasks.append(convert_task)
                            effect_task_count += 1
                        input_queue.task_done()
                    except asyncio.TimeoutError:
                        break

                if self._stop_event.is_set():
                    break

                push_tasks = []
                for fastpath_output in fastpath_outputs:
                    push_task = asyncio.get_event_loop().run_in_executor(
                        self._io_executor, channel.push, fastpath_output
                    )
                    push_tasks.append(push_task)

                if push_tasks:
                    await asyncio.gather(*push_tasks)

                if not convert_tasks:
                    continue
                converted_inputs = await asyncio.gather(*convert_tasks)

                valid_inputs = []
                failed_inputs = []
                for raw_input, converted_input in zip(raw_inputs, converted_inputs):
                    if converted_input is not None:
                        valid_inputs.append(converted_input)
                    else:
                        failed_inputs.append(raw_input)

                await self._process_batch(valid_inputs, failed_inputs, channel)
            except Exception as e:
                if not self._stop_event.is_set():
                    traceback.print_exc()
                    self.logger.error(f"Error in consume loop: {e}")
                await asyncio.sleep(1.0)

        self.logger.info("Batch consume loop stopped")

    async def _process_batch(
        self,
        batch_inputs: List[TaskInput[Any]],
        failed_inputs: List[Dict],
        channel: BaseChannel,
    ):
        if not batch_inputs and not failed_inputs:
            return

        all_outputs = []

        for failed_input in failed_inputs:
            task_id = failed_input.get("task_id", "unknown")
            error_output = self.create_error_output(
                task_id=task_id,
                error_type="InputConversionError",
                error_message="Failed to convert input to TaskInput",
                error_details={"raw_input": failed_input},
            )
            all_outputs.append((error_output, False))
            self.logger.warning(
                f"Created error output for failed input conversion: {task_id}"
            )

        if batch_inputs:
            self.logger.info(f"Processing batch of {len(batch_inputs)} valid tasks")

            try:
                batch_outputs = await self.process_batch(batch_inputs)
                all_outputs.extend([(_output, (_input.meta or {}).get("persistent-output", False)) for _input, _output in zip(batch_inputs, batch_outputs)])
                self.logger.info(
                    f"Batch processing completed: {len(batch_outputs)} outputs"
                )
            except Exception as e:
                traceback_str = traceback.format_exc()
                self.logger.error(f"Batch processing failed: {e}")
                for task_input in batch_inputs:
                    error_output = self.create_error_output(
                        task_id=task_input.task_id,
                        error_type="ProcessingError",
                        error_message=str(e),
                        error_details={"batch_size": len(batch_inputs)},
                        traceback_str=traceback_str,
                    )
                    all_outputs.append((error_output, False))

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

        convert_tasks = []
        for output, persistent in all_outputs:
            convert_task = asyncio.create_task(
                self.convert_taskoutput_to_output(output, not persistent)
            )
            convert_tasks.append(convert_task)

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

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

        push_tasks = []
        for output in final_outputs:
            push_task = asyncio.get_event_loop().run_in_executor(
                self._io_executor, channel.push, output
            )
            push_tasks.append(push_task)

        await asyncio.gather(*push_tasks, return_exceptions=True)
        total_tasks = len(batch_inputs) + len(failed_inputs)
        self.logger.info(
            f"Completed processing {total_tasks} tasks ({len(batch_inputs)} valid, {len(failed_inputs)} failed)"
        )

    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._io_executor:
            self.logger.info("Shutting down io executor poll ...")
            try:
                self._io_executor.shutdown(wait=True, cancel_futures=True)
                self.logger.info("IO process pool shutdown successfully")
            except Exception as e:
                self.logger.error(f"Error shutting down io executor pool: {e}")
            finally:
                self._io_executor = None

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