"""数据采集模块

定时采集队列快照数据（积压指标 + Worker 指标），写入 queue_snapshot 表。
"""

import asyncio
import json
import logging
import time
from collections import defaultdict
from datetime import datetime, timezone, timedelta
from typing import Dict, List, Optional, Set, TYPE_CHECKING

from sqlalchemy.dialects.postgresql import insert
from sqlalchemy import delete

from jettask.db.models.queue_snapshot import QueueSnapshot

if TYPE_CHECKING:
    from jettask import Jettask

logger = logging.getLogger(__name__)


class QueueSnapshotCollector:
    """队列快照采集器

    合并积压指标和 Worker 指标的采集，写入 queue_snapshot 表。

    特性：
    - 使用分布式锁，保证多进程环境下只有一个进程执行采集
    - 每分钟采集一次（可配置）
    - 记录任务级明细和 namespace 级汇总（Worker 全局去重）
    - 自动清理过期数据
    """

    LOCK_KEY_TEMPLATE = "{prefix}:queue_snapshot:lock"
    LOCK_TIMEOUT = 60       
    SNAPSHOT_INTERVAL = 60  
    CLEANUP_INTERVAL = 3600 
    RETENTION_DAYS = 7      
    HEARTBEAT_TIMEOUT = 15.0  

    def __init__(
        self,
        app: "Jettask",
        async_redis_client,
        pg_session_factory,
        namespace: str
    ):
        self.app = app
        self.redis = async_redis_client
        self.pg_session_factory = pg_session_factory
        self.namespace = namespace
        self._running = False
        self._task: Optional[asyncio.Task] = None
        self._last_cleanup_time = 0

        self.workers_registry_key = f"{app.redis_prefix}:REGISTRY:WORKERS"

    async def start(self):
        self._running = True
        logger.info(f"[QueueSnapshot] 启动队列快照采集服务, namespace={self.namespace}")

        while self._running:
            try:
                await self._collect_with_lock()

                current_time = asyncio.get_event_loop().time()
                if current_time - self._last_cleanup_time >= self.CLEANUP_INTERVAL:
                    await self._cleanup_old_snapshots()
                    self._last_cleanup_time = current_time

            except asyncio.CancelledError:
                logger.info("[QueueSnapshot] 采集任务被取消")
                break
            except Exception as e:
                logger.error(f"[QueueSnapshot] 采集失败: {e}", exc_info=True)

            await asyncio.sleep(self.SNAPSHOT_INTERVAL)

        logger.info("[QueueSnapshot] 队列快照采集服务已停止")

    async def stop(self):
        self._running = False
        if self._task and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        logger.info("[QueueSnapshot] 停止队列快照采集服务")

    async def _collect_with_lock(self):
        lock_key = self.LOCK_KEY_TEMPLATE.format(prefix=self.app.redis_prefix)

        lock = self.redis.lock(
            name=lock_key,
            timeout=self.LOCK_TIMEOUT
        )

        acquired = await lock.acquire(blocking=False)
        if not acquired:
            logger.debug("[QueueSnapshot] 未获取到锁，跳过本次采集")
            return

        try:
            await self._do_collect()
        finally:
            try:
                await lock.release()
            except Exception:
                pass  

    async def _do_collect(self):
        time_bucket = self._get_time_bucket()

        backlog_data = await self._collect_backlog_data(time_bucket)

        worker_data = await self._collect_worker_data()

        snapshots = self._merge_snapshots(time_bucket, backlog_data, worker_data)

        if snapshots:
            await self._batch_upsert(snapshots)
            logger.info(
                f"[QueueSnapshot] 写入 {len(snapshots)} 条快照, "
                f"backlog_data={len(backlog_data)} 条, "
                f"worker_data={len(worker_data)} 组"
            )


    async def _collect_backlog_data(self, time_bucket: datetime) -> List[dict]:
        from jettask.messaging.registry import QueueRegistry

        registry = QueueRegistry(
            redis_client=None,
            async_redis_client=self.redis,
            redis_prefix=self.app.redis_prefix
        )
        task_queue_pairs = await registry.get_all_task_queue_pairs()

        if not task_queue_pairs:
            logger.debug("[QueueSnapshot] 没有任务需要采集积压数据")
            return []

        queues = set(pair['queue'] for pair in task_queue_pairs)
        logger.debug(f"[QueueSnapshot] 采集 {len(queues)} 个队列的积压数据")

        snapshots = []
        for queue in queues:
            try:
                backlog_info = await self.app.get_queue_backlog(queue, asyncio=True)
                logger.debug(f"[QueueSnapshot] 队列 {queue} backlog_info: {backlog_info}")

                if 'error' in backlog_info:
                    logger.warning(f"[QueueSnapshot] 获取队列 {queue} 积压失败: {backlog_info.get('error')}")
                    continue

                send_offset = backlog_info.get('send_offset', 0)

                for task_info in backlog_info.get('tasks', []):
                    snapshots.append({
                        'queue': queue,
                        'task_name': task_info['task_name'],
                        'backlog': task_info.get('backlog', 0),
                        'pending_read': task_info.get('pending_read', 0),
                        'send_offset': send_offset,
                        'ack_offset': task_info.get('ack_offset', 0),
                        'read_offset': task_info.get('read_offset', 0),
                    })

            except Exception as e:
                logger.error(f"[QueueSnapshot] 处理队列 {queue} 积压数据失败: {e}")

        return snapshots


    async def _collect_worker_data(self) -> Dict[tuple, Dict[str, Set[str]]]:
        current_time = time.time()

        try:
            all_worker_ids = await self.redis.smembers(self.workers_registry_key)
            if not all_worker_ids:
                return {}

            pipeline = self.redis.pipeline()
            worker_id_list = list(all_worker_ids)
            for worker_id in worker_id_list:
                wid = worker_id.decode('utf-8') if isinstance(worker_id, bytes) else worker_id
                worker_key = f"{self.app.redis_prefix}:WORKER:{wid}"
                pipeline.hgetall(worker_key)
            worker_data_list = await pipeline.execute()

            result: Dict[tuple, Dict[str, Set[str]]] = defaultdict(lambda: {'online': set(), 'offline': set()})

            for worker_id, worker_data in zip(worker_id_list, worker_data_list):
                if not worker_data:
                    continue

                wid = worker_id.decode('utf-8') if isinstance(worker_id, bytes) else worker_id

                last_heartbeat_raw = worker_data.get('last_heartbeat') or worker_data.get(b'last_heartbeat')
                is_alive_raw = worker_data.get('is_alive') or worker_data.get(b'is_alive')

                last_heartbeat = 0
                if last_heartbeat_raw:
                    val = last_heartbeat_raw.decode('utf-8') if isinstance(last_heartbeat_raw, bytes) else last_heartbeat_raw
                    try:
                        last_heartbeat = float(val)
                    except (ValueError, TypeError):
                        pass

                is_alive = True
                if is_alive_raw:
                    val = is_alive_raw.decode('utf-8') if isinstance(is_alive_raw, bytes) else is_alive_raw
                    is_alive = str(val).lower() == 'true'

                is_online = is_alive and (current_time - last_heartbeat) < self.HEARTBEAT_TIMEOUT
                status = 'online' if is_online else 'offline'

                for key, value in worker_data.items():
                    key_str = key.decode('utf-8') if isinstance(key, bytes) else key
                    if not key_str.startswith('group_info:'):
                        continue

                    try:
                        val_str = value.decode('utf-8') if isinstance(value, bytes) else value
                        group_info = json.loads(val_str)
                        queue = group_info.get('queue')
                        task_name = group_info.get('task_name')
                        if queue and task_name:
                            result[(queue, task_name)][status].add(wid)
                    except (json.JSONDecodeError, Exception) as e:
                        logger.debug(f"[QueueSnapshot] 解析 group_info 失败: {key_str}, error: {e}")
                        continue

            logger.debug(f"[QueueSnapshot] 采集到 {len(result)} 个 (queue, task) Worker 组合")
            return result

        except Exception as e:
            logger.error(f"[QueueSnapshot] 采集 Worker 数据失败: {e}", exc_info=True)
            return {}


    def _merge_snapshots(
        self,
        time_bucket: datetime,
        backlog_data: List[dict],
        worker_data: Dict[tuple, Dict[str, Set[str]]]
    ) -> List[dict]:
        backlog_index: Dict[tuple, dict] = {}
        for item in backlog_data:
            key = (item['queue'], item['task_name'])
            backlog_index[key] = item

        all_keys = set(backlog_index.keys())
        for (queue, task_name) in worker_data:
            all_keys.add((queue, task_name))

        snapshots = []
        empty_backlog = {
            'backlog': 0, 'pending_read': 0,
            'send_offset': 0, 'ack_offset': 0, 'read_offset': 0
        }

        for (queue, task_name) in all_keys:
            backlog_info = backlog_index.get((queue, task_name), empty_backlog)
            workers = worker_data.get((queue, task_name), {'online': set(), 'offline': set()})
            online_count = len(workers['online'])
            offline_count = len(workers['offline'])

            snapshots.append({
                'time_bucket': time_bucket,
                'namespace': self.namespace,
                'queue': queue,
                'task_name': task_name,
                'backlog': backlog_info.get('backlog', 0),
                'pending_read': backlog_info.get('pending_read', 0),
                'send_offset': backlog_info.get('send_offset', 0),
                'ack_offset': backlog_info.get('ack_offset', 0),
                'read_offset': backlog_info.get('read_offset', 0),
                'online_count': online_count,
                'offline_count': offline_count,
                'total_count': online_count + offline_count,
            })

        queue_online: Dict[str, Set[str]] = defaultdict(set)
        queue_offline: Dict[str, Set[str]] = defaultdict(set)
        for (queue, task_name), workers in worker_data.items():
            queue_online[queue].update(workers['online'])
            queue_offline[queue].update(workers['offline'])

        queue_backlog: Dict[str, int] = defaultdict(int)
        queue_pending: Dict[str, int] = defaultdict(int)
        queue_send_offset: Dict[str, int] = defaultdict(int)
        queue_ack_offset: Dict[str, int] = defaultdict(int)
        queue_read_offset: Dict[str, int] = defaultdict(int)
        for item in backlog_data:
            q = item['queue']
            queue_backlog[q] += item.get('backlog', 0)
            queue_pending[q] += item.get('pending_read', 0)
            queue_send_offset[q] += item.get('send_offset', 0)
            queue_ack_offset[q] += item.get('ack_offset', 0)
            queue_read_offset[q] += item.get('read_offset', 0)

        all_queues = set(list(queue_online.keys()) + list(queue_offline.keys()) + list(queue_backlog.keys()))
        for queue in all_queues:
            online = len(queue_online.get(queue, set()))
            offline = len(queue_offline.get(queue, set()))
            snapshots.append({
                'time_bucket': time_bucket,
                'namespace': self.namespace,
                'queue': queue,
                'task_name': '',
                'backlog': queue_backlog.get(queue, 0),
                'pending_read': queue_pending.get(queue, 0),
                'send_offset': queue_send_offset.get(queue, 0),
                'ack_offset': queue_ack_offset.get(queue, 0),
                'read_offset': queue_read_offset.get(queue, 0),
                'online_count': online,
                'offline_count': offline,
                'total_count': online + offline,
            })

        total_backlog = sum(item.get('backlog', 0) for item in backlog_data)
        total_pending_read = sum(item.get('pending_read', 0) for item in backlog_data)
        total_send_offset = sum(item.get('send_offset', 0) for item in backlog_data)
        total_ack_offset = sum(item.get('ack_offset', 0) for item in backlog_data)
        total_read_offset = sum(item.get('read_offset', 0) for item in backlog_data)

        all_online: Set[str] = set()
        all_offline: Set[str] = set()
        for workers in worker_data.values():
            all_online.update(workers['online'])
            all_offline.update(workers['offline'])

        ns_online = len(all_online)
        ns_offline = len(all_offline)
        snapshots.append({
            'time_bucket': time_bucket,
            'namespace': self.namespace,
            'queue': '',
            'task_name': '',
            'backlog': total_backlog,
            'pending_read': total_pending_read,
            'send_offset': total_send_offset,
            'ack_offset': total_ack_offset,
            'read_offset': total_read_offset,
            'online_count': ns_online,
            'offline_count': ns_offline,
            'total_count': ns_online + ns_offline,
        })

        return snapshots


    def _get_time_bucket(self) -> datetime:
        now = datetime.now(timezone.utc)
        return now.replace(second=0, microsecond=0)

    async def _batch_upsert(self, snapshots: List[dict]):
        async with self.pg_session_factory() as session:
            stmt = insert(QueueSnapshot).values(snapshots)
            stmt = stmt.on_conflict_do_update(
                index_elements=['time_bucket', 'namespace', 'queue', 'task_name'],
                set_={
                    'backlog': stmt.excluded.backlog,
                    'pending_read': stmt.excluded.pending_read,
                    'send_offset': stmt.excluded.send_offset,
                    'ack_offset': stmt.excluded.ack_offset,
                    'read_offset': stmt.excluded.read_offset,
                    'online_count': stmt.excluded.online_count,
                    'offline_count': stmt.excluded.offline_count,
                    'total_count': stmt.excluded.total_count,
                }
            )
            await session.execute(stmt)
            await session.commit()

    async def _cleanup_old_snapshots(self):
        cutoff = datetime.now(timezone.utc) - timedelta(days=self.RETENTION_DAYS)

        try:
            async with self.pg_session_factory() as session:
                result = await session.execute(
                    delete(QueueSnapshot).where(
                        QueueSnapshot.time_bucket < cutoff
                    )
                )
                await session.commit()
                deleted_count = result.rowcount
                if deleted_count > 0:
                    logger.info(f"[QueueSnapshot] 清理 {deleted_count} 条过期快照（{self.RETENTION_DAYS} 天前）")
        except Exception as e:
            logger.error(f"[QueueSnapshot] 清理过期快照失败: {e}", exc_info=True)
