"""
Azure Service Bus Watcher for Casbin
Place this file alongside your service.py:

abs_auth_rbac_core/
  rbac/
    service.py          ← existing
    watcher.py          ← THIS FILE (new)
    decorator.py        ← existing (NO changes)
    policy.conf         ← existing (NO changes)
    types.py            ← existing (NO changes)

pip install azure-servicebus azure-identity
"""

import json
import logging
import os
import threading
import time
import uuid
from typing import Optional, Callable

from pydantic import BaseModel

logger = logging.getLogger(__name__)


# =============================================================================
# CONFIG SCHEMA (replaces RedisWatcherSchema)
# =============================================================================

class AzureServiceBusWatcherSchema(BaseModel):
    """
    Drop-in replacement for RedisWatcherSchema.
    
    Usage:
        config = AzureServiceBusWatcherSchema(
            connection_string="Endpoint=sb://your-ns.servicebus.windows.net/;...",
            topic_name="casbin-policy-updated",
        )
    """
    # Connection - use ONE of these two:
    connection_string: Optional[str] = None
    fully_qualified_namespace: Optional[str] = None  # For managed identity

    # Topic & Subscription
    topic_name: str = "casbin-policy-updated"
    subscription_name: Optional[str] = None  # Auto-set from RBAC_INSTANCE_ID env var

    # Behavior
    max_wait_time_seconds: int = 5
    max_message_count: int = 10
    message_ttl_seconds: int = 300  # 5 minutes
    auto_create_subscription: bool = True


# =============================================================================
# WATCHER (implements Casbin Watcher interface)
# =============================================================================

class AzureServiceBusWatcher:
    """
    Casbin Watcher using Azure Service Bus Topics/Subscriptions.
    
    Implements the same interface as casbin-redis-watcher so it's
    a drop-in replacement in RBACService.
    
    Guarantees:
    - At-least-once delivery (messages persist until acknowledged)
    - Durable subscriptions (messages queue during server downtime)
    - Dead letter queue for failed messages
    - Auto-reconnect with exponential backoff
    - Self-notification filtering
    """

    def __init__(self, config: AzureServiceBusWatcherSchema):
        self.config = config
        self.callback: Optional[Callable] = None
        self.mutex = threading.Lock()
        self._running = True

        # Unique stable instance ID per server
        self.instance_id = config.subscription_name or os.environ.get(
            "RBAC_INSTANCE_ID",
            os.environ.get("HOSTNAME", f"casbin-{uuid.uuid4().hex[:8]}")
        )
        self.subscription_name = self.instance_id

        # Initialize client + sender
        self._client = None
        self._sender = None
        self._init_client()

        # Background receiver thread
        self._receiver_thread = threading.Thread(
            target=self._receiver_loop,
            daemon=True,
            name=f"casbin-servicebus-{self.instance_id}",
        )
        self.subscribe_event = threading.Event()

    # =========================================================================
    # INITIALIZATION
    # =========================================================================

    def _init_client(self):
        from azure.servicebus import ServiceBusClient

        if self.config.connection_string:
            self._client = ServiceBusClient.from_connection_string(
                conn_str=self.config.connection_string,
                logging_enable=False,
            )
        elif self.config.fully_qualified_namespace:
            from azure.identity import DefaultAzureCredential
            self._client = ServiceBusClient(
                fully_qualified_namespace=self.config.fully_qualified_namespace,
                credential=DefaultAzureCredential(),
                logging_enable=False,
            )
        else:
            raise ValueError(
                "Provide either 'connection_string' or 'fully_qualified_namespace'"
            )

        self._sender = self._client.get_topic_sender(
            topic_name=self.config.topic_name
        )

    def _ensure_subscription(self):
        """Auto-create the subscription if it doesn't exist."""
        if not self.config.auto_create_subscription:
            return

        try:
            from azure.servicebus.management import ServiceBusAdministrationClient
            import datetime

            if self.config.connection_string:
                admin = ServiceBusAdministrationClient.from_connection_string(
                    conn_str=self.config.connection_string
                )
            else:
                from azure.identity import DefaultAzureCredential
                admin = ServiceBusAdministrationClient(
                    fully_qualified_namespace=self.config.fully_qualified_namespace,
                    credential=DefaultAzureCredential(),
                )

            with admin:
                try:
                    admin.get_subscription(
                        self.config.topic_name, self.subscription_name
                    )
                    logger.info(
                        f"[RBAC-SB] Subscription '{self.subscription_name}' exists"
                    )
                except Exception:
                    admin.create_subscription(
                        topic_name=self.config.topic_name,
                        subscription_name=self.subscription_name,
                        default_message_time_to_live=datetime.timedelta(
                            seconds=self.config.message_ttl_seconds
                        ),
                        dead_lettering_on_message_expiration=True,
                        max_delivery_count=5,
                    )
                    logger.info(
                        f"[RBAC-SB] Created subscription '{self.subscription_name}'"
                    )

        except ImportError:
            logger.warning(
                "[RBAC-SB] Cannot auto-create subscription. "
                "Create it manually in Azure Portal."
            )
        except Exception as e:
            logger.warning(f"[RBAC-SB] Subscription auto-create failed: {e}")

    def start(self):
        """Start the watcher. Call after set_update_callback."""
        logger.info(f"[RBAC-SB] ========================================")
        logger.info(f"[RBAC-SB] Starting Azure Service Bus Watcher")
        logger.info(f"[RBAC-SB] Instance ID:       {self.instance_id}")
        logger.info(f"[RBAC-SB] Topic:             {self.config.topic_name}")
        logger.info(f"[RBAC-SB] Subscription:      {self.subscription_name}")
        logger.info(f"[RBAC-SB] ========================================")

        self._ensure_subscription()
        self._receiver_thread.start()
        self.subscribe_event.wait(timeout=10)

        if self.subscribe_event.is_set():
            logger.info(f"[RBAC-SB] ✅ Watcher STARTED - subscribing to: {self.config.topic_name}/{self.subscription_name}")
        else:
            logger.error(f"[RBAC-SB] ❌ Watcher FAILED to start within timeout")

    # =========================================================================
    # RECEIVER (background thread - receives policy change notifications)
    # =========================================================================

    def _receiver_loop(self):
        backoff = 1

        while self._running:
            try:
                logger.info(f"[RBAC-SB] 🔌 Connecting receiver to topic='{self.config.topic_name}', subscription='{self.subscription_name}'")

                receiver = self._client.get_subscription_receiver(
                    topic_name=self.config.topic_name,
                    subscription_name=self.subscription_name,
                    max_wait_time=self.config.max_wait_time_seconds,
                )

                if not self.subscribe_event.is_set():
                    self.subscribe_event.set()

                backoff = 1
                logger.info(f"[RBAC-SB] ✅ Receiver CONNECTED - listening on: {self.config.topic_name}/{self.subscription_name}")

                with receiver:
                    while self._running:
                        messages = receiver.receive_messages(
                            max_message_count=self.config.max_message_count,
                            max_wait_time=self.config.max_wait_time_seconds,
                        )
                        if messages:
                            logger.info(f"[RBAC-SB] 📩 Received {len(messages)} message(s) on subscription='{self.subscription_name}'")
                        for msg in messages:
                            try:
                                self._process_message(msg, receiver)
                            except Exception as e:
                                logger.error(f"[RBAC-SB] Message processing error on subscription='{self.subscription_name}': {e}")
                                try:
                                    receiver.abandon_message(msg)
                                except Exception:
                                    pass

            except Exception as e:
                if not self._running:
                    break
                if not self.subscribe_event.is_set():
                    self.subscribe_event.set()
                logger.warning(
                    f"[RBAC-SB] Receiver error, reconnecting in {backoff}s: {e}"
                )
                time.sleep(backoff)
                backoff = min(backoff * 2, 30)

    def _process_message(self, msg, receiver):
        try:
            body = str(msg)
            data = json.loads(body)
        except (json.JSONDecodeError, Exception):
            logger.warning(f"[RBAC-SB] ⚠️ Could not parse message on subscription='{self.subscription_name}', completing it")
            receiver.complete_message(msg)
            return

        sender_id = data.get("id", "")
        method = data.get("method", "")

        # Skip self-notifications
        if sender_id == self.instance_id:
            logger.debug(f"[RBAC-SB] ⏭️ Skipping self-notification on subscription='{self.subscription_name}' (method={method})")
            receiver.complete_message(msg)
            return

        logger.info(
            f"[RBAC-SB] 📨 Processing message on subscription='{self.subscription_name}': "
            f"method={method}, from={sender_id}, my_id={self.instance_id}"
        )

        # Execute callback
        with self.mutex:
            if self.callback is not None:
                try:
                    self.callback(body)
                    logger.info(f"[RBAC-SB] ✅ Callback executed successfully on subscription='{self.subscription_name}'")
                except Exception as e:
                    logger.error(f"[RBAC-SB] ❌ Callback FAILED on subscription='{self.subscription_name}': {e}")
                    receiver.abandon_message(msg)
                    return
            else:
                logger.warning(f"[RBAC-SB] ⚠️ No callback set on subscription='{self.subscription_name}'")

        # Acknowledge only after successful callback
        receiver.complete_message(msg)

    # =========================================================================
    # PUBLISHER (sends policy change notifications)
    # =========================================================================

    def _publish(self, msg_dict: dict) -> bool:
        from azure.servicebus import ServiceBusMessage

        for attempt in range(3):
            try:
                if self._sender is None:
                    self._sender = self._client.get_topic_sender(
                        topic_name=self.config.topic_name
                    )
                message = ServiceBusMessage(json.dumps(msg_dict))
                self._sender.send_messages(message)
                return True
            except Exception as e:
                logger.warning(f"[RBAC-SB] Publish attempt {attempt+1}/3 failed: {e}")
                self._sender = None
                time.sleep(0.5)

        logger.error("[RBAC-SB] Publish failed after 3 attempts")
        return False

    # =========================================================================
    # CASBIN WATCHER INTERFACE
    # (same methods as Redis watcher - drop-in compatible)
    # =========================================================================

    def set_update_callback(self, callback):
        with self.mutex:
            self.callback = callback

    def update(self):
        return self._publish({"method": "Update", "id": self.instance_id})

    def update_for_add_policy(self, section, ptype, *params):
        return self._publish({
            "method": "UpdateForAddPolicy", "id": self.instance_id,
            "sec": section, "ptype": ptype, "params": params,
        })

    def update_for_remove_policy(self, section, ptype, *params):
        return self._publish({
            "method": "UpdateForRemovePolicy", "id": self.instance_id,
            "sec": section, "ptype": ptype, "params": params,
        })

    def update_for_remove_filtered_policy(self, section, ptype, field_index, *params):
        return self._publish({
            "method": "UpdateForRemoveFilteredPolicy", "id": self.instance_id,
            "sec": section, "ptype": ptype,
            "params": f"{field_index} {' '.join(str(p) for p in params)}",
        })

    def update_for_save_policy(self, model):
        return self._publish({"method": "UpdateForSavePolicy", "id": self.instance_id})

    def update_for_add_policies(self, section, ptype, *params):
        return self._publish({
            "method": "UpdateForAddPolicies", "id": self.instance_id,
            "sec": section, "ptype": ptype, "params": params,
        })

    def update_for_remove_policies(self, section, ptype, *params):
        return self._publish({
            "method": "UpdateForRemovePolicies", "id": self.instance_id,
            "sec": section, "ptype": ptype, "params": params,
        })

    def close(self):
        self._running = False
        try:
            if self._sender:
                self._sender.close()
            if self._client:
                self._client.close()
        except Exception:
            pass