import json
import os
import ssl
import threading
import time

from pathlib import Path
from queue import Empty

from memos.configs.mem_scheduler import AuthConfig, RabbitMQConfig
from memos.context.context import ContextThread
from memos.dependency import require_python_package
from memos.log import get_logger
from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
from memos.mem_scheduler.general_modules.misc import AutoDroppingQueue
from memos.mem_scheduler.schemas.general_schemas import DIRECT_EXCHANGE_TYPE, FANOUT_EXCHANGE_TYPE


logger = get_logger(__name__)

PLAYGROUND_CHAT_STREAM_PATH = "/product/chat/stream/playground"


class RabbitMQSchedulerModule(BaseSchedulerModule):
    @require_python_package(
        import_name="pika",
        install_command="pip install pika",
        install_link="https://pika.readthedocs.io/en/stable/index.html",
    )
    def __init__(self):
        """
        Initialize RabbitMQ connection settings.
        """
        super().__init__()
        self.auth_config = None

        # RabbitMQ settings
        self.rabbitmq_config: RabbitMQConfig | None = None
        self.rabbit_queue_name = "memos-scheduler"
        self.rabbitmq_exchange_name = "memos-fanout"  # Default, will be overridden by config
        self.rabbitmq_exchange_type = FANOUT_EXCHANGE_TYPE  # Default, will be overridden by config
        self.rabbitmq_playground_chat_exchange_name: str | None = None
        self.rabbitmq_playground_chat_exchange_type = FANOUT_EXCHANGE_TYPE
        self.rabbitmq_connection = None
        self.rabbitmq_channel = None

        # fixed params
        self.rabbitmq_message_cache_max_size = 10  # Max 10 messages
        self.rabbitmq_message_cache = AutoDroppingQueue(
            maxsize=self.rabbitmq_message_cache_max_size
        )
        # Pending outgoing messages to avoid loss when connection is not ready
        self.rabbitmq_publish_cache_max_size = 50
        self.rabbitmq_publish_cache = AutoDroppingQueue(
            maxsize=self.rabbitmq_publish_cache_max_size
        )
        self.rabbitmq_connection_attempts = 3  # Max retry attempts on connection failure
        self.rabbitmq_retry_delay = 5  # Delay (seconds) between retries
        self.rabbitmq_heartbeat = 60  # Heartbeat interval (seconds) for connectio
        self.rabbitmq_conn_max_waiting_seconds = 30
        self.rabbitmq_conn_sleep_seconds = 1

        # Thread management
        self._rabbitmq_io_loop_thread = None  # For IOLoop execution
        self._rabbitmq_stop_flag = False  # Graceful shutdown flag
        # Use RLock because publishing may trigger initialization, which also grabs the lock.
        self._rabbitmq_lock = threading.RLock()
        self._rabbitmq_initializing = False  # Avoid duplicate concurrent initializations

    def is_rabbitmq_connected(self) -> bool:
        """Check if RabbitMQ connection is alive"""
        return (
            self.rabbitmq_connection
            and self.rabbitmq_connection.is_open
            and self.rabbitmq_channel
            and self.rabbitmq_channel.is_open
        )

    def initialize_rabbitmq(
        self, config: dict | None | RabbitMQConfig = None, config_path: str | Path | None = None
    ):
        """
        Establish connection to RabbitMQ using pika.
        """
        with self._rabbitmq_lock:
            if self._rabbitmq_initializing:
                logger.info(
                    "[DIAGNOSTIC] initialize_rabbitmq: initialization already in progress; skipping duplicate call."
                )
                return
            self._rabbitmq_initializing = True
        try:
            # Skip remote initialization in CI/pytest unless explicitly enabled
            enable_env = os.getenv("MEMOS_ENABLE_RABBITMQ", "").lower() == "true"
            in_ci = os.getenv("CI", "").lower() == "true"
            in_pytest = os.getenv("PYTEST_CURRENT_TEST") is not None
            logger.info(
                f"[DIAGNOSTIC] initialize_rabbitmq called. in_ci={in_ci}, in_pytest={in_pytest}, "
                f"MEMOS_ENABLE_RABBITMQ={enable_env}, config_path={config_path}"
            )
            if (in_ci or in_pytest) and not enable_env:
                logger.info(
                    "Skipping RabbitMQ initialization in CI/test environment. Set MEMOS_ENABLE_RABBITMQ=true to enable."
                )
                return

            if self.is_rabbitmq_connected():
                logger.warning("RabbitMQ is already connected. Skipping initialization.")
                return

            from pika.adapters.select_connection import SelectConnection

            if config is not None:
                if isinstance(config, RabbitMQConfig):
                    self.rabbitmq_config = config
                elif isinstance(config, dict):
                    self.rabbitmq_config = AuthConfig.from_dict(config).rabbitmq
                else:
                    logger.error(f"Unsupported config type: {type(config)}")
                    return

            else:
                if config_path is not None and Path(config_path).exists():
                    self.auth_config = AuthConfig.from_local_config(config_path=config_path)
                elif AuthConfig.default_config_exists():
                    self.auth_config = AuthConfig.from_local_config()
                else:
                    self.auth_config = AuthConfig.from_local_env()
                self.rabbitmq_config = self.auth_config.rabbitmq

            if self.rabbitmq_config is None:
                logger.error(
                    "Failed to load RabbitMQ configuration. Please check your config file or environment variables."
                )
                return

            # Load exchange configuration from config
            if self.rabbitmq_config:
                if (
                    hasattr(self.rabbitmq_config, "exchange_name")
                    and self.rabbitmq_config.exchange_name
                ):
                    self.rabbitmq_exchange_name = self.rabbitmq_config.exchange_name
                    logger.info(f"Using configured exchange name: {self.rabbitmq_exchange_name}")
                if (
                    hasattr(self.rabbitmq_config, "exchange_type")
                    and self.rabbitmq_config.exchange_type
                ):
                    self.rabbitmq_exchange_type = self.rabbitmq_config.exchange_type
                    logger.info(f"Using configured exchange type: {self.rabbitmq_exchange_type}")

            env_exchange_name = os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME")
            env_exchange_type = os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_TYPE")
            if env_exchange_name:
                self.rabbitmq_exchange_name = env_exchange_name
                logger.info(f"Using env exchange name override: {self.rabbitmq_exchange_name}")
            if env_exchange_type:
                self.rabbitmq_exchange_type = env_exchange_type
                logger.info(f"Using env exchange type override: {self.rabbitmq_exchange_type}")

            playground_exchange_name = os.getenv(
                "MEMSCHEDULER_RABBITMQ_PLAYGROUND_CHAT_EXCHANGE_NAME", ""
            ).strip()
            playground_exchange_type = os.getenv(
                "MEMSCHEDULER_RABBITMQ_PLAYGROUND_CHAT_EXCHANGE_TYPE",
                FANOUT_EXCHANGE_TYPE,
            ).strip()
            if playground_exchange_name:
                self.rabbitmq_playground_chat_exchange_name = playground_exchange_name
                self.rabbitmq_playground_chat_exchange_type = (
                    playground_exchange_type or FANOUT_EXCHANGE_TYPE
                )
                logger.info(
                    "Using playground chat exchange override: name=%s, type=%s",
                    self.rabbitmq_playground_chat_exchange_name,
                    self.rabbitmq_playground_chat_exchange_type,
                )
            else:
                self.rabbitmq_playground_chat_exchange_name = None
                self.rabbitmq_playground_chat_exchange_type = FANOUT_EXCHANGE_TYPE

            # Start connection process
            parameters = self.get_rabbitmq_connection_param()
            self.rabbitmq_connection = SelectConnection(
                parameters,
                on_open_callback=self.on_rabbitmq_connection_open,
                on_open_error_callback=self.on_rabbitmq_connection_error,
                on_close_callback=self.on_rabbitmq_connection_closed,
            )

            # Start IOLoop in dedicated thread
            self._io_loop_thread = ContextThread(
                target=self.rabbitmq_connection.ioloop.start, daemon=True
            )
            self._io_loop_thread.start()
            logger.info("RabbitMQ connection process started")
        except Exception:
            logger.error("Failed to initialize RabbitMQ connection", exc_info=True)
        finally:
            with self._rabbitmq_lock:
                self._rabbitmq_initializing = False

    def get_rabbitmq_queue_size(self) -> int:
        """Get the current number of messages in the queue.

        Returns:
            int: Number of messages in the queue.
                 Returns -1 if there's an error or no active connection.
        """
        if self.rabbitmq_exchange_type != DIRECT_EXCHANGE_TYPE:
            logger.warning("Queue size can only be checked for direct exchanges")
            return None

        with self._rabbitmq_lock:
            if not self.is_rabbitmq_connected():
                logger.warning("No active connection to check queue size")
                return -1

            # Declare queue passively (only checks existence, doesn't create)
            # Using passive=True prevents accidental queue creation
            result = self.rabbitmq_channel.queue_declare(
                queue=self.rabbit_queue_name,
                durable=True,  # Match the original queue durability setting
                passive=True,  # Only check queue existence, don't create
            )

            if result is None:
                return 0
            # Return the message count from the queue declaration result
            return result.method.message_count

    def get_rabbitmq_connection_param(self):
        import pika

        credentials = pika.PlainCredentials(
            username=self.rabbitmq_config.user_name,
            password=self.rabbitmq_config.password,
            erase_on_connect=self.rabbitmq_config.erase_on_connect,
        )
        if self.rabbitmq_config.port == 5671:
            context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
            context.check_hostname = False
            context.verify_mode = False
            return pika.ConnectionParameters(
                host=self.rabbitmq_config.host_name,
                port=self.rabbitmq_config.port,
                virtual_host=self.rabbitmq_config.virtual_host,
                credentials=credentials,
                ssl_options=pika.SSLOptions(context),
                connection_attempts=self.rabbitmq_connection_attempts,
                retry_delay=self.rabbitmq_retry_delay,
                heartbeat=self.rabbitmq_heartbeat,
            )
        else:
            return pika.ConnectionParameters(
                host=self.rabbitmq_config.host_name,
                port=self.rabbitmq_config.port,
                virtual_host=self.rabbitmq_config.virtual_host,
                credentials=credentials,
                connection_attempts=self.rabbitmq_connection_attempts,
                retry_delay=self.rabbitmq_retry_delay,
                heartbeat=self.rabbitmq_heartbeat,
            )

    # Connection lifecycle callbacks
    def on_rabbitmq_connection_open(self, connection):
        """Called when connection is established."""
        logger.info("[DIAGNOSTIC] RabbitMQ connection opened")
        connection.channel(on_open_callback=self.on_rabbitmq_channel_open)

    def on_rabbitmq_connection_error(self, connection, error):
        """Called if connection fails to open."""
        logger.error(f"Connection failed: {error}")
        self.rabbit_reconnect()

    def on_rabbitmq_connection_closed(self, connection, reason):
        """Called when connection closes."""
        logger.warning(f"Connection closed: {reason}")
        if not self._rabbitmq_stop_flag:
            self.rabbit_reconnect()

    # Channel lifecycle callbacks
    def on_rabbitmq_channel_open(self, channel):
        """Called when channel is ready."""
        self.rabbitmq_channel = channel
        logger.info("[DIAGNOSTIC] RabbitMQ channel opened")

        # Setup primary/direct exchange and optional playground/fanout exchange.
        channel.exchange_declare(
            exchange=self.rabbitmq_exchange_name,
            exchange_type=self.rabbitmq_exchange_type,
            durable=True,
            callback=self.on_rabbitmq_primary_exchange_declared,
        )

    def on_rabbitmq_primary_exchange_declared(self, frame):
        """Called when primary exchange is ready."""
        if self.rabbitmq_playground_chat_exchange_name:
            self.rabbitmq_channel.exchange_declare(
                exchange=self.rabbitmq_playground_chat_exchange_name,
                exchange_type=self.rabbitmq_playground_chat_exchange_type,
                durable=True,
                callback=self.on_rabbitmq_playground_exchange_declared,
            )
            return

        self._rabbitmq_continue_queue_setup()

    def on_rabbitmq_playground_exchange_declared(self, frame):
        """Called when optional playground exchange is ready."""
        self._rabbitmq_continue_queue_setup()

    def _rabbitmq_continue_queue_setup(self):
        """Declare scheduler queue and bind it to the primary exchange."""
        self.rabbitmq_channel.queue_declare(
            queue=self.rabbit_queue_name, durable=True, callback=self.on_rabbitmq_queue_declared
        )

    def on_rabbitmq_queue_declared(self, frame):
        """Called when queue is ready."""
        self.rabbitmq_channel.queue_bind(
            exchange=self.rabbitmq_exchange_name,
            queue=self.rabbit_queue_name,
            routing_key=self.rabbit_queue_name,
            callback=self.on_rabbitmq_bind_ok,
        )

    def on_rabbitmq_bind_ok(self, frame):
        """Final setup step when bind is complete."""
        logger.info("RabbitMQ setup completed")
        # Flush any cached publish messages now that connection is ready
        self._flush_cached_publish_messages()

    def resolve_publish_route(self, message: dict) -> tuple[str, str]:
        api_path = message.get("api_path")
        if api_path == PLAYGROUND_CHAT_STREAM_PATH and self.rabbitmq_playground_chat_exchange_name:
            return self.rabbitmq_playground_chat_exchange_name, ""

        return self.rabbitmq_exchange_name, ""

    def on_rabbitmq_message(self, channel, method, properties, body):
        """Handle incoming messages. Only for test."""
        try:
            print(f"Received message: {body.decode()}\n")
            self.rabbitmq_message_cache.put({"properties": properties, "body": body})
            print(f"message delivery_tag: {method.delivery_tag}\n")
            channel.basic_ack(delivery_tag=method.delivery_tag)
        except Exception as e:
            logger.error(f"Message handling failed: {e}", exc_info=True)

    def wait_for_connection_ready(self):
        start_time = time.time()
        while not self.is_rabbitmq_connected():
            delta_time = time.time() - start_time
            if delta_time > self.rabbitmq_conn_max_waiting_seconds:
                logger.error("Failed to start consuming: Connection timeout")
                return False
            self.rabbit_reconnect()
            time.sleep(self.rabbitmq_conn_sleep_seconds)  # Reduced frequency of checks

    # Message handling
    def rabbitmq_start_consuming(self):
        """Start consuming messages asynchronously."""
        self.wait_for_connection_ready()

        self.rabbitmq_channel.basic_consume(
            queue=self.rabbit_queue_name,
            on_message_callback=self.on_rabbitmq_message,
            auto_ack=False,
        )
        logger.info("Started rabbitmq consuming messages")

    def rabbitmq_publish_message(self, message: dict):
        """
        Publish a message to RabbitMQ.
        """
        import pika

        exchange_name, routing_key = self.resolve_publish_route(message)
        label = message.get("label")

        logger.info(
            "[DIAGNOSTIC] Publishing %s message. api_path=%s Exchange: %s, Routing Key: '%s'.",
            label,
            message.get("api_path"),
            exchange_name,
            routing_key,
        )
        logger.info(f"  - Message Content: {json.dumps(message, indent=2, ensure_ascii=False)}")

        with self._rabbitmq_lock:
            logger.info(
                f"[DIAGNOSTIC] rabbitmq_service.rabbitmq_publish_message invoked. "
                f"is_connected={self.is_rabbitmq_connected()}, exchange={exchange_name}, "
                f"routing_key='{routing_key}', label={label}"
            )
            if not self.is_rabbitmq_connected():
                logger.error(
                    "[DIAGNOSTIC] Cannot publish - no active connection. Caching message for retry. "
                    f"connection_exists={bool(self.rabbitmq_connection)}, "
                    f"channel_exists={bool(self.rabbitmq_channel)}, "
                    f"config_loaded={self.rabbitmq_config is not None}"
                )
                self.rabbitmq_publish_cache.put(message)
                # Best-effort to connect
                self.initialize_rabbitmq(config=self.rabbitmq_config)
                return False

            logger.info(
                f"[DIAGNOSTIC] rabbitmq_service.rabbitmq_publish_message: Attempting to publish message. Exchange: {exchange_name}, Routing Key: {routing_key}, Message Content: {json.dumps(message, indent=2, ensure_ascii=False)}"
            )
            try:
                self.rabbitmq_channel.basic_publish(
                    exchange=exchange_name,
                    routing_key=routing_key,
                    body=json.dumps(message),
                    properties=pika.BasicProperties(
                        delivery_mode=2,  # Persistent
                    ),
                    mandatory=True,
                )
                logger.debug(f"Published message: {message}")
                return True
            except Exception as e:
                logger.error(
                    "[DIAGNOSTIC] RabbitMQ publish error. label=%s item_id=%s exchange=%s "
                    "routing_key=%s error=%s",
                    label,
                    message.get("item_id"),
                    exchange_name,
                    routing_key,
                    e,
                )
                logger.error(f"Failed to publish message: {e}")
                # Cache message for retry on next connection
                self.rabbitmq_publish_cache.put(message)
                self.rabbit_reconnect()
                return False

    # Connection management
    def rabbit_reconnect(self):
        """Schedule reconnection attempt."""
        logger.info("Attempting to reconnect...")
        if self.rabbitmq_connection and not self.rabbitmq_connection.is_closed:
            self.rabbitmq_connection.ioloop.stop()

        # Reset connection state
        self.rabbitmq_channel = None
        self.initialize_rabbitmq()

    def rabbitmq_close(self):
        """Gracefully shutdown connection."""
        with self._rabbitmq_lock:
            self._rabbitmq_stop_flag = True

            # Close channel if open
            if self.rabbitmq_channel and self.rabbitmq_channel.is_open:
                try:
                    self.rabbitmq_channel.close()
                except Exception as e:
                    logger.warning(f"Error closing channel: {e}")

            # Close connection if open
            if self.rabbitmq_connection:
                if self.rabbitmq_connection.is_open:
                    try:
                        self.rabbitmq_connection.close()
                    except Exception as e:
                        logger.warning(f"Error closing connection: {e}")

                # Stop IOLoop if running
                try:
                    self.rabbitmq_connection.ioloop.stop()
                except Exception as e:
                    logger.warning(f"Error stopping IOLoop: {e}")

            # Wait for IOLoop thread to finish
            if self._io_loop_thread and self._io_loop_thread.is_alive():
                self._io_loop_thread.join(timeout=5)
                if self._io_loop_thread.is_alive():
                    logger.warning("IOLoop thread did not terminate cleanly")

        logger.info("RabbitMQ connection closed")

    def _flush_cached_publish_messages(self):
        """Flush cached outgoing messages once connection is available."""
        if self.rabbitmq_publish_cache.empty():
            return

        if not self.is_rabbitmq_connected():
            logger.info(
                "[DIAGNOSTIC] _flush_cached_publish_messages: connection still down; "
                f"pending={self.rabbitmq_publish_cache.qsize()}"
            )
            return

        drained: list[dict] = []
        while True:
            try:
                drained.append(self.rabbitmq_publish_cache.get_nowait())
            except Empty:
                break

        if not drained:
            return

        logger.info(
            f"[DIAGNOSTIC] Flushing {len(drained)} cached RabbitMQ messages after reconnect."
        )
        for cached_msg in drained:
            success = self.rabbitmq_publish_message(cached_msg)
            if not success:
                # Message already re-cached inside publish; avoid tight loop
                logger.error(
                    "[DIAGNOSTIC] Failed to flush cached message; re-queued for next attempt."
                )
                break
