"""Task consumption loop: group init, consume, process, timeout."""
from __future__ import annotations

import json
import logging
import threading
from typing import Any, Callable, Dict, Optional

from actor.schemas import load_schema

logger = logging.getLogger(__name__)


def _validate_task_schema(task: Dict[str, Any]) -> None:
    try:
        import jsonschema
    except Exception as exc:  # pragma: no cover
        raise ValueError("jsonschema library required for task validation") from exc
    schema = load_schema("task_message.schema.json")
    jsonschema.validate(instance=task, schema=schema)


class TaskConsumer:
    """Reads from the tasks stream, de-duplicates, executes, and ACKs."""

    def __init__(self, conn, dedup, publisher, tasks_stream: str, consumer_group: str, instance_id: str, state_key: str) -> None:
        self._conn = conn          # RedisConnection
        self._dedup = dedup        # DeduplicationTracker
        self._publisher = publisher  # ResultPublisher
        self._tasks_stream = tasks_stream
        self._consumer_group = consumer_group
        self._instance_id = instance_id
        self._state_key = state_key
        self._group_exists_logged = False
        self._stop = False

    # -- consumer group --------------------------------------------------

    def initialize_consumer(self) -> None:
        """Create consumer group for tasks stream if it does not exist."""
        try:
            self._conn.client.xgroup_create(
                name=self._tasks_stream,
                groupname=self._consumer_group,
                id="0",
                mkstream=True,
            )
            logger.info("Initialized consumer group '%s' on stream %s", self._consumer_group, self._tasks_stream)
        except Exception as exc:
            if "BUSYGROUP" in str(exc):
                if not self._group_exists_logged:
                    logger.info("Consumer group '%s' already exists on %s", self._consumer_group, self._tasks_stream)
                    self._group_exists_logged = True
                return
            from actor.queue.connection import QueueError
            raise QueueError(f"Failed to create consumer group: {exc}") from exc

    def _ensure_group(self) -> None:
        try:
            self.initialize_consumer()
        except Exception as exc:
            logger.error("Failed to ensure consumer group: %s", exc)

    # -- main loop -------------------------------------------------------

    def consume_tasks(
        self,
        callback: Callable[[Dict[str, Any]], Dict[str, Any]],
        block_ms: int = 5000,
        task_timeout: Optional[int] = 30,
        max_idle_iterations: Optional[int] = None,
        start_id: Optional[str] = None,
    ) -> None:
        """Blocking loop — consumes tasks, executes *callback*, publishes results."""
        idle_count = 0
        while not self._stop:
            try:
                self._ensure_group()
                messages = self._conn.client.xreadgroup(
                    groupname=self._consumer_group,
                    consumername=self._instance_id,
                    streams={self._tasks_stream: start_id or ">"},
                    count=1,
                    block=block_ms,
                )
                if not messages:
                    idle_count += 1
                    self._recover_pending(callback, task_timeout=task_timeout)
                    if max_idle_iterations is not None and idle_count >= max_idle_iterations:
                        logger.debug("Max idle iterations reached; stopping consumer loop.")
                        break
                    continue
                idle_count = 0
                for stream_name, msg_list in messages:
                    for msg_id, msg_data in msg_list:
                        logger.debug("Received task message id=%s stream=%s", msg_id, stream_name)
                        if msg_data.get("type") == "config":
                            logger.info("Received config on tasks stream id=%s; skipping task processing", msg_id)
                            try:
                                self._conn.client.xack(self._tasks_stream, self._consumer_group, msg_id)
                                self._dedup.persist_ack(msg_id)
                            except Exception:
                                pass
                            continue
                        if self._dedup.is_duplicate(msg_id):
                            logger.debug("Duplicate message id=%s ignored", msg_id)
                            try:
                                self._conn.client.xack(self._tasks_stream, self._consumer_group, msg_id)
                                self._dedup.persist_ack(msg_id)
                            except Exception:
                                pass
                            continue
                        self._process_message(msg_id, msg_data, callback, task_timeout=task_timeout)
                        self._dedup.record_seen(msg_id)
            except KeyboardInterrupt:
                logger.info("Queue consumer interrupted; shutting down.")
                self._stop = True
                break
            except Exception as exc:
                logger.error("Queue consumer error: %s", exc)
                if self._conn.is_connection_error(exc):
                    self._conn.reconnect(
                        stop_flag=lambda: self._stop,
                        on_reconnect=self._ensure_group,
                    )
                    continue
                self._conn.reconnect(stop_flag=lambda: self._stop, on_reconnect=self._ensure_group)

    def stop(self) -> None:
        self._stop = True

    # -- pending recovery ------------------------------------------------

    def _recover_pending(
        self,
        callback: Callable[[Dict[str, Any]], Dict[str, Any]],
        task_timeout: Optional[int],
    ) -> None:
        """Claim recent pending messages for this consumer (best-effort)."""
        try:
            max_idle_ms = 5 * 60 * 1000
            pending = self._conn.client.xpending_range(
                name=self._tasks_stream,
                groupname=self._consumer_group,
                min="-",
                max="+",
                count=20,
            )
            ids = []
            for entry in pending:
                idle = getattr(entry, "idle", None)
                msg_id = getattr(entry, "message_id", None)
                if idle is None and isinstance(entry, dict):
                    idle = entry.get("idle")
                    msg_id = entry.get("message_id")
                if idle is None and isinstance(entry, tuple) and len(entry) > 0:
                    msg_id = entry[0]
                if msg_id and idle is not None and idle <= max_idle_ms:
                    ids.append(msg_id)
            if not ids:
                return
            claimed = self._conn.client.xclaim(
                name=self._tasks_stream,
                groupname=self._consumer_group,
                consumername=self._instance_id,
                min_idle_time=0,
                message_ids=ids,
            )
            for msg_id, msg_data in claimed:
                logger.debug("Recovered pending message id=%s", msg_id)
                self._process_message(msg_id, msg_data, callback=callback, task_timeout=task_timeout)
        except Exception as exc:
            logger.warning("Pending recovery skipped: %s", exc)

    # -- message processing ----------------------------------------------

    def _run_with_timeout(
        self,
        fn: Callable[[Dict[str, Any]], Dict[str, Any]],
        task: Dict[str, Any],
        timeout: Optional[int],
    ):
        if timeout is None:
            return fn(task)
        result_container: Dict[str, Any] = {}
        exc_container: Dict[str, Exception] = {}

        def target():
            try:
                result_container["result"] = fn(task)
            except Exception as exc:  # pragma: no cover
                exc_container["exc"] = exc

        # NOTE: daemon threads cannot be cancelled in Python; if the timeout
        # fires, the thread will continue running until the process exits.
        th = threading.Thread(target=target, daemon=True)
        th.start()
        th.join(timeout)
        if th.is_alive():
            logger.warning(
                "Task %s exceeded timeout %ss — daemon thread will leak until process exit",
                task.get("task_id"),
                timeout,
            )
            raise TimeoutError(f"Task {task.get('task_id')} exceeded timeout {timeout}s")
        if "exc" in exc_container:
            raise exc_container["exc"]
        return result_container.get("result")

    def _process_message(
        self,
        msg_id: str,
        msg_data: Dict[str, Any],
        callback: Callable[[Dict[str, Any]], Dict[str, Any]],
        task_timeout: Optional[int] = None,
    ) -> None:
        task_id = msg_data.get("task_id")
        task_type = msg_data.get("task_type")
        payload_raw = msg_data.get("payload", "{}")
        try:
            payload = json.loads(payload_raw)
        except Exception as exc:
            raise ValueError(f"Invalid JSON payload for task_id={task_id}") from exc

        task = {"task_id": task_id, "task_type": task_type, "payload": payload}

        try:
            _validate_task_schema(task)
            logger.info("Processing task task_id=%s task_type=%s", task_id, task_type)
            result = self._run_with_timeout(callback, task, task_timeout)
            self._publisher.publish_result(task_id, result)
        except Exception as exc:
            logger.error("Task execution failed for %s: %s", task_id, exc)
            from actor.cli.utils import sanitize_error_message
            self._publisher.publish_error(task_id, sanitize_error_message(exc))
        finally:
            try:
                self._conn.client.xack(self._tasks_stream, self._consumer_group, msg_id)
                self._dedup.persist_ack(msg_id)
                logger.debug("ACKed message id=%s for task_id=%s", msg_id, task_id)
            except Exception as exc:  # pragma: no cover
                logger.error("Failed to ACK message %s: %s", msg_id, exc)
