# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
import abc
import json
import threading
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Union

from openviking.pyagfs import AGFSSyncClientProtocol, AsyncAGFSClient
from openviking_cli.utils.logger import get_logger

logger = get_logger(__name__)


@dataclass
class QueueError:
    """Error record."""

    timestamp: datetime
    message: str
    data: Optional[Dict[str, Any]] = None


@dataclass
class QueueStatus:
    """Queue status."""

    pending: int = 0
    in_progress: int = 0
    processed: int = 0
    requeue_count: int = 0
    error_count: int = 0
    errors: List[QueueError] = field(default_factory=list)

    @property
    def has_errors(self) -> bool:
        return self.error_count > 0

    @property
    def is_complete(self) -> bool:
        return self.pending == 0 and self.in_progress == 0


class EnqueueHookBase(abc.ABC):
    """Enqueue hook base class.

    All custom enqueue logic should inherit from this base class.
    Provides on_enqueue method for custom processing before message enqueue.
    """

    @abc.abstractmethod
    async def on_enqueue(self, data: Union[str, Dict[str, Any]]) -> Union[str, Dict[str, Any]]:
        """Called before message enqueue. Can modify data or perform validation."""
        return data


class DequeueHandlerBase(abc.ABC):
    """Dequeue handler base class, supports callback mechanism to report processing results."""

    _success_callback: Optional[Callable[[], None]] = None
    _requeue_callback: Optional[Callable[[], None]] = None
    _error_callback: Optional[Callable[[str, Optional[Dict[str, Any]]], None]] = None

    def set_callbacks(
        self,
        on_success: Callable[[], None],
        on_requeue: Callable[[], None],
        on_error: Callable[[str, Optional[Dict[str, Any]]], None],
    ) -> None:
        """Set callback functions."""
        self._success_callback = on_success
        self._requeue_callback = on_requeue
        self._error_callback = on_error

    def report_success(self) -> None:
        """Report processing success."""
        if self._success_callback:
            self._success_callback()

    def report_requeue(self) -> None:
        """Report that the current message was re-enqueued for later retry."""
        if self._requeue_callback:
            self._requeue_callback()

    def report_error(self, error_msg: str, data: Optional[Dict[str, Any]] = None) -> None:
        """Report processing error."""
        if self._error_callback:
            self._error_callback(error_msg, data)

    @abc.abstractmethod
    async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """Called after message dequeue. Returns None to discard message."""
        if not data:
            return None
        return data


class NamedQueue:
    """NamedQueue: Operation class for specific named queue, supports status tracking."""

    MAX_ERRORS = 100

    def __init__(
        self,
        agfs: AGFSSyncClientProtocol,
        mount_point: str,
        name: str,
        enqueue_hook: Optional[EnqueueHookBase] = None,
        dequeue_handler: Optional[DequeueHandlerBase] = None,
    ):
        self.name = name
        self.path = f"{mount_point}/{name}"
        self._agfs = agfs
        self._async_agfs = AsyncAGFSClient(agfs)
        self._enqueue_hook = enqueue_hook
        self._dequeue_handler = dequeue_handler
        self._initialized = False

        # Status tracking
        self._lock = threading.Lock()
        self._in_progress = 0
        self._processed = 0
        self._requeue_count = 0
        self._error_count = 0
        self._errors: List[QueueError] = []

        # Inject callbacks to handler
        if self._dequeue_handler:
            self._dequeue_handler.set_callbacks(
                on_success=self._on_process_success,
                on_requeue=self._on_process_requeue,
                on_error=self._on_process_error,
            )

    def _on_dequeue_start(self) -> None:
        """Called on dequeue."""
        with self._lock:
            self._in_progress += 1

    def _on_process_success(self) -> None:
        """Called on processing success."""
        with self._lock:
            self._in_progress -= 1
            self._processed += 1

    def _on_process_requeue(self) -> None:
        """Called when a dequeued message is re-enqueued for later retry."""
        with self._lock:
            self._requeue_count += 1

    def _on_process_error(self, error_msg: str, data: Optional[Dict[str, Any]] = None) -> None:
        """Called on processing failure."""
        with self._lock:
            self._in_progress -= 1
            self._error_count += 1
            self._errors.append(
                QueueError(
                    timestamp=datetime.now(),
                    message=error_msg,
                    data=data,
                )
            )
            if len(self._errors) > self.MAX_ERRORS:
                self._errors = self._errors[-self.MAX_ERRORS :]

    async def get_status(self) -> QueueStatus:
        """Get queue status."""
        pending = await self.size()
        with self._lock:
            return QueueStatus(
                pending=pending,
                in_progress=self._in_progress,
                processed=self._processed,
                requeue_count=self._requeue_count,
                error_count=self._error_count,
                errors=list(self._errors),
            )

    def reset_status(self) -> None:
        """Reset status counters."""
        with self._lock:
            self._in_progress = 0
            self._processed = 0
            self._requeue_count = 0
            self._error_count = 0
            self._errors = []

    def has_dequeue_handler(self) -> bool:
        """Check if dequeue handler exists."""
        return self._dequeue_handler is not None

    async def _ensure_initialized(self):
        """Ensure queue directory is created in AGFS."""
        if not self._initialized:
            try:
                await self._async_agfs.mkdir(self.path)
            except Exception as e:
                if "exist" not in str(e).lower():
                    logger.warning(f"[NamedQueue] Failed to ensure queue {self.name}: {e}")
            self._initialized = True

    async def enqueue(self, data: Union[str, Dict[str, Any]]) -> str:
        """Send message to queue (enqueue)."""
        await self._ensure_initialized()
        enqueue_file = f"{self.path}/enqueue"

        # Execute enqueue hook
        if self._enqueue_hook:
            data = await self._enqueue_hook.on_enqueue(data)

        if isinstance(data, dict):
            data = json.dumps(data)

        msg_id = await self._async_agfs.write(enqueue_file, data.encode("utf-8"))
        return msg_id if isinstance(msg_id, str) else str(msg_id)

    async def ack(self, msg_id: str) -> None:
        """Acknowledge successful processing of a message (deletes it from persistent storage).

        Must be called after the dequeue handler finishes processing a message.
        If not called (e.g. process crashes), the message will be automatically
        re-queued on the next startup via RecoverStale.
        """
        if not msg_id:
            return
        ack_file = f"{self.path}/ack"
        try:
            await self._async_agfs.write(ack_file, msg_id.encode("utf-8"))
        except Exception as e:
            logger.warning(f"[NamedQueue] Ack failed for {self.name} msg_id={msg_id}: {e}")

    async def _read_queue_message(self) -> Optional[Dict[str, Any]]:
        """Read and remove one message from the AGFS queue; return parsed dict or None.

        Normalises the various return types AGFSClient.read() may produce.
        """
        content = await self._async_agfs.read(f"{self.path}/dequeue")
        if not content or content == b"{}":
            return None
        if isinstance(content, bytes):
            raw = content
        elif isinstance(content, str):
            raw = content.encode("utf-8")
        elif hasattr(content, "content") and content.content is not None:
            raw = content.content
        else:
            raw = str(content).encode("utf-8")
        return json.loads(raw.decode("utf-8"))

    async def dequeue(self) -> Optional[Dict[str, Any]]:
        """Dequeue a message, process it, then ack to confirm deletion.

        Flow (at-least-once delivery):
          1. Read from /dequeue  → backend marks message as 'processing' (not deleted yet)
          2. Call on_dequeue()   → actual processing
          3. Call ack()          → backend deletes the message permanently

        If the process crashes between steps 1 and 3, the backend's RecoverStale
        on the next startup resets the message back to 'pending' for retry.
        """
        await self._ensure_initialized()
        try:
            data = await self._read_queue_message()
            if data is None:
                return None
            # Capture message ID before passing data to handler (handler may modify it)
            msg_id = data.get("id", "") if isinstance(data, dict) else ""
            if self._dequeue_handler:
                self._on_dequeue_start()
                data = await self._dequeue_handler.on_dequeue(data)
            # Ack unconditionally after handler returns (success or handled error).
            # If on_dequeue raises, the exception propagates and ack is skipped —
            # the message will be recovered on next startup.
            await self.ack(msg_id)
            return data
        except Exception as e:
            logger.debug(f"[NamedQueue] Dequeue failed for {self.name}: {e}")
            return None

    async def dequeue_raw(self) -> Optional[Dict[str, Any]]:
        """Get and remove message from queue without invoking the handler."""
        await self._ensure_initialized()
        try:
            return await self._read_queue_message()
        except Exception as e:
            logger.debug(f"[NamedQueue] Dequeue raw failed for {self.name}: {e}")
            return None

    async def process_dequeued(self, data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Invoke the dequeue handler on already-fetched raw data.

        NOTE: caller must call _on_dequeue_start() before invoking this method
        so that in_progress is incremented atomically with the dequeue.
        """
        if self._dequeue_handler:
            return await self._dequeue_handler.on_dequeue(data)
        return data

    async def peek(self) -> Optional[Dict[str, Any]]:
        """Peek at head message without removing."""
        await self._ensure_initialized()
        peek_file = f"{self.path}/peek"

        try:
            content = await self._async_agfs.read(peek_file)
            if not content or content == b"{}":
                return None
            if isinstance(content, bytes):
                return json.loads(content.decode("utf-8"))
            elif isinstance(content, str):
                return json.loads(content)
            else:
                return None
        except Exception as e:
            logger.debug(f"[NamedQueue] Peek failed for {self.name}: {e}")
            return None

    async def size(self) -> int:
        """Get queue size."""
        await self._ensure_initialized()
        size_file = f"{self.path}/size"

        try:
            content = await self._async_agfs.read(size_file)
            if not content:
                return 0
            if isinstance(content, bytes):
                return int(content.decode("utf-8").strip())
            elif isinstance(content, str):
                return int(content.strip())
            else:
                return 0
        except Exception as e:
            logger.debug(f"[NamedQueue] Get size failed for {self.name}: {e}")
            return 0

    async def clear(self) -> bool:
        """Clear queue."""
        await self._ensure_initialized()
        clear_file = f"{self.path}/clear"

        try:
            await self._async_agfs.write(clear_file, b"")
            return True
        except Exception as e:
            logger.error(f"[NamedQueue] Clear failed for {self.name}: {e}")
            return False
