ForcePilot/backend/package/yuxi/channels/manager.py

from __future__ import annotations

import asyncio
import time
from collections import defaultdict, deque
from typing import Any

from sqlalchemy import func, select

from yuxi.channels.base import BaseChannelAdapter
from yuxi.channels.infra.broadcast import EventBroadcaster
from yuxi.channels.infra.circuit_breaker import CircuitBreaker
from yuxi.channels.infra.config_watcher import ConfigWatcher
from yuxi.channels.services.context import GatewayRequestContext
from yuxi.channels.services.doctor import ConfigDoctor, DiagnosisIssue
from yuxi.channels.exceptions import ChannelException
from yuxi.channels.services.maintenance import MaintenanceRunner
from yuxi.channels.models import ChannelStatus
from yuxi.channels.registry import ChannelRegistry, _BUILTIN_ADAPTERS, _load_builtin_adapters
from yuxi.channels.router import MessageRouter
from yuxi.channels.services.runtime_state import RuntimeState
from yuxi.channels.services.stats_collector import StatsCollector
from yuxi.channels.services.webhook_registry import WebhookRegistry
from yuxi.channels.services.ws_logger import WsLogger
from yuxi.storage.postgres.manager import pg_manager
from yuxi.storage.postgres.models_channels import ChannelMsgRecord
from yuxi.utils.logging_config import logger

HEALTH_CHECK_INTERVAL = 60
CONFIG_WATCH_INTERVAL = 30.0


class ChannelManager:
    def __init__(self, registry: ChannelRegistry | None = None, router: MessageRouter | None = None):
        self._registry = registry or ChannelRegistry()
        self._router = router or MessageRouter(channel_manager=self)
        self._adapters: dict[str, BaseChannelAdapter] = {}
        self._circuit_breakers: dict[str, CircuitBreaker] = defaultdict(lambda: CircuitBreaker())
        self._health_tasks: dict[str, asyncio.Task] = {}
        self._initialized = False
        self._phase: str = "not_started"
        self._rate_limiters: dict[str, deque[float]] = {}
        self._rate_limit_locks: dict[str, asyncio.Lock] = {}
        self._channels_config: dict[str, dict[str, Any]] = {}
        self._watcher: ConfigWatcher | None = None
        self._doctor: ConfigDoctor | None = None
        self._broadcaster: EventBroadcaster | None = None
        self._ctx: GatewayRequestContext | None = None
        self.runtime_state = RuntimeState()
        self._auth_limiter: Any = None
        self._ws_handlers: dict[str, Any] = {}
        self._scheduled_tasks: list[asyncio.Task] = []
        self._ws_logger: WsLogger | None = None
        self._maintenance_runner: MaintenanceRunner | None = None
        self._stats_collector: StatsCollector | None = None
        self._webhook_registry: WebhookRegistry | None = None

    @property
    def phase(self) -> str:
        return self._phase

    @property
    def context(self) -> GatewayRequestContext | None:
        return self._ctx

    @property
    def broadcaster(self) -> EventBroadcaster | None:
        return self._broadcaster

    @property
    def watcher(self) -> ConfigWatcher | None:
        return self._watcher

    @property
    def doctor(self) -> ConfigDoctor | None:
        return self._doctor

    async def load_config(self) -> None:
        if self._phase not in ("not_started",):
            return

        _load_builtin_adapters()
        for channel_id, adapter_cls in _BUILTIN_ADAPTERS.items():
            self._registry.register(channel_id, adapter_cls)

        from yuxi import config as conf

        self._channels_config = getattr(conf, "channels", {})
        self._phase = "config_loaded"
        logger.info("ChannelManager: config loaded")

    async def prepare_bootstrap(self) -> None:
        if self._phase not in ("config_loaded",):
            return

        async with pg_manager.get_async_session_context() as db:
            await self._ensure_virtual_department(db)
            await self._ensure_default_agent_config(db)

        self._phase = "bootstrapped"
        logger.info("ChannelManager: bootstrap prepared")

    async def start_channels(self) -> None:
        if self._phase not in ("config_loaded", "bootstrapped"):
            return

        for channel_id in self._registry.list_channels():
            channel_conf = self._channels_config.get(channel_id, {})
            if channel_conf.get("enabled", False):
                try:
                    await self.start_channel(channel_id, channel_conf)
                except Exception:
                    logger.exception(f"Failed to start channel {channel_id}")

        self._phase = "channels_started"
        logger.info(f"ChannelManager: channels started: {list(self._adapters.keys())}")

    async def start_subscriptions(self) -> None:
        if self._phase not in ("channels_started",):
            return

        self._phase = "fully_running"
        logger.info("ChannelManager: subscriptions started, fully running")

    async def initialize(self) -> None:
        if self._initialized:
            return

        _load_builtin_adapters()
        for channel_id, adapter_cls in _BUILTIN_ADAPTERS.items():
            self._registry.register(channel_id, adapter_cls)

        from yuxi import config as conf

        self._channels_config = getattr(conf, "channels", {})

        for channel_id in self._registry.list_channels():
            channel_conf = self._channels_config.get(channel_id, {})
            if channel_conf.get("enabled", False):
                try:
                    await self.start_channel(channel_id, channel_conf)
                except Exception:
                    logger.exception(f"Failed to start channel {channel_id}")

        async with pg_manager.get_async_session_context() as db:
            await self._ensure_virtual_department(db)
            await self._ensure_default_agent_config(db)

        self._initialized = True
        self._phase = "fully_running"
        logger.info(f"ChannelManager initialized with channels: {list(self._adapters.keys())}")

    async def shutdown(self) -> None:
        for channel_id in list(self._adapters.keys()):
            try:
                await self.stop_channel(channel_id)
            except Exception:
                logger.exception(f"Failed to stop channel {channel_id} during shutdown")

        for task in self._health_tasks.values():
            task.cancel()

        for task in self._scheduled_tasks:
            task.cancel()

        if self._watcher:
            await self._watcher.stop()

        self._initialized = False
        logger.info("ChannelManager shutdown complete")

    async def startup(self) -> None:
        if self._initialized:
            return

        try:
            await self._stage_load_config()
            await self._stage_prepare_bootstrap()
            await self._stage_start_early_runtime()
            await self._stage_create_auth_limiter()
            await self._stage_init_channels()
            await self._stage_create_runtime_state()
            await self._stage_start_runtime_services()
            await self._stage_activate_scheduled_services()
            await self._stage_attach_ws_handlers()
            await self._stage_start_event_subscriptions()
            self._initialized = True
            self._phase = "fully_running"
            logger.info("ChannelManager: 10-stage startup complete")
        except Exception:
            logger.exception(f"ChannelManager startup failed at phase: {self._phase}")
            raise

    async def _stage_load_config(self) -> None:
        if self._phase not in ("not_started",):
            return

        _load_builtin_adapters()
        for channel_id, adapter_cls in _BUILTIN_ADAPTERS.items():
            self._registry.register(channel_id, adapter_cls)

        from yuxi import config as conf

        self._channels_config = getattr(conf, "channels", {})
        self._phase = "config_loaded"
        logger.info("ChannelManager: [1/10] config loaded")

    async def _stage_prepare_bootstrap(self) -> None:
        if self._phase not in ("config_loaded",):
            return

        async with pg_manager.get_async_session_context() as db:
            await self._ensure_virtual_department(db)
            await self._ensure_default_agent_config(db)

        self._phase = "bootstrapped"
        logger.info("ChannelManager: [2/10] bootstrap prepared")

    async def _stage_start_early_runtime(self) -> None:
        if self._phase not in ("bootstrapped",):
            return

        self._broadcaster = EventBroadcaster()

        self._ctx = GatewayRequestContext(
            runtime_config=self._channels_config,
            start_channel=self._ctx_start_channel,
            stop_channel=self._ctx_stop_channel,
            mark_channel_logged_out=self._ctx_mark_channel_logged_out,
            broadcast_fn=self._broadcaster.broadcast,
            node_send_to_session_fn=self._broadcaster.node_send_to_session,
        )

        self._phase = "early_runtime"
        logger.info("ChannelManager: [3/10] early runtime started")

    async def _stage_create_auth_limiter(self) -> None:
        if self._phase not in ("early_runtime",):
            return

        self._phase = "auth_limiter_created"
        logger.info("ChannelManager: [4/10] auth rate limiter created")

    async def _stage_init_channels(self) -> None:
        if self._phase not in ("early_runtime", "auth_limiter_created"):
            return

        for channel_id in self._registry.list_channels():
            channel_conf = self._channels_config.get(channel_id, {})
            if channel_conf.get("enabled", False):
                try:
                    await self.start_channel(channel_id, channel_conf)
                except Exception:
                    logger.exception(f"Failed to start channel {channel_id}")

        self._phase = "channels_started"
        logger.info(f"ChannelManager: [5/10] channels started: {list(self._adapters.keys())}")

    async def _stage_create_runtime_state(self) -> None:
        if self._phase not in ("channels_started",):
            return

        self.runtime_state = RuntimeState(
            started_at=time.monotonic(),
            active_channels=len(self._adapters),
            phase=self._phase,
            node_id="forcepilot-gateway",
            main_node=True,
            services=["doctor", "maintenance", "stats", "webhooks"],
        )
        self._ws_logger = WsLogger(max_entries=1000)

        self._phase = "runtime_state_created"
        logger.info("ChannelManager: [6/10] runtime state created")

    async def _stage_start_runtime_services(self) -> None:
        if self._phase not in ("runtime_state_created",):
            return

        self._doctor = ConfigDoctor(self)
        self._maintenance_runner = MaintenanceRunner(self, self.runtime_state)
        self._stats_collector = StatsCollector(self.runtime_state)
        self._webhook_registry = WebhookRegistry(self)

        self._phase = "runtime_services_started"
        logger.info("ChannelManager: [7/10] runtime services started (doctor+maintenance+stats+webhooks)")

    async def _stage_activate_scheduled_services(self) -> None:
        if self._phase not in ("runtime_services_started",):
            return

        self._watcher = ConfigWatcher(self)
        await self._watcher.watch(interval=CONFIG_WATCH_INTERVAL)

        self._scheduled_tasks.append(asyncio.create_task(self._maintenance_runner.run()))
        self._scheduled_tasks.append(asyncio.create_task(self._stats_collector.run()))
        self._scheduled_tasks.append(asyncio.create_task(self._webhook_registry.run()))

        self._phase = "scheduled_services_active"
        logger.info("ChannelManager: [8/10] scheduled services activated (watcher+maintenance+stats+webhooks)")

    async def _stage_attach_ws_handlers(self) -> None:
        if self._phase not in ("scheduled_services_active",):
            return

        if self._broadcaster:
            self._broadcaster.subscribe_callback("channel.status_change", self._on_channel_status_change)
            self._broadcaster.subscribe_callback("tick", self._on_tick)
            self._broadcaster.subscribe_callback("chat", self._on_chat_event)
            self._broadcaster.subscribe_callback("channel.logout", self._on_channel_logout)
            self._broadcaster.subscribe_callback("config.reload", self._on_config_reload)

        self._phase = "ws_handlers_attached"
        logger.info("ChannelManager: [9/10] websocket handlers attached")

    async def _stage_start_event_subscriptions(self) -> None:
        if self._phase not in ("ws_handlers_attached",):
            return

        self.runtime_state.services = ["doctor", "maintenance", "stats", "webhooks", "watcher", "ws"]

        self._phase = "subscriptions_started"
        logger.info("ChannelManager: [10/10] event subscriptions started")

    async def _on_channel_status_change(self, event: str, payload: Any) -> None:
        logger.debug(f"WS event: {event}")

    async def _on_tick(self, event: str, payload: Any) -> None:
        pass

    async def _on_chat_event(self, event: str, payload: Any) -> None:
        logger.debug(f"WS chat event: {event}")

    async def _on_channel_logout(self, event: str, payload: Any) -> None:
        channel_id = payload.get("channel_id") if isinstance(payload, dict) else None
        account_id = payload.get("account_id") if isinstance(payload, dict) else None
        logger.info(f"Channel logout: {channel_id}/{account_id}")

    async def _on_config_reload(self, event: str, payload: Any) -> None:
        await self.reload_config_now()

    async def _ctx_start_channel(self, channel_id: str, config: dict) -> None:
        await self.start_channel(channel_id, config)

    async def _ctx_stop_channel(self, channel_id: str) -> None:
        await self.stop_channel(channel_id)

    async def _ctx_mark_channel_logged_out(self, channel_id: str, account_id: str) -> None:
        logger.info(f"Channel {channel_id} account {account_id} marked as logged out")

    async def diagnose_channels(self) -> list[DiagnosisIssue]:
        if self._doctor is None:
            return []
        return await self._doctor.diagnose()

    async def auto_fix_channel(self, issue: DiagnosisIssue) -> bool:
        if self._doctor is None:
            return False
        return await self._doctor.auto_fix(issue)

    async def reload_config_now(self, changed_keys: list[str] | None = None) -> dict[str, Any]:
        if self._watcher is None:
            return {}
        return await self._watcher.reload_now(changed_keys)

    async def start_channel(self, channel_id: str, config: dict[str, Any] | None = None) -> None:
        if channel_id in self._adapters:
            logger.warning(f"Channel {channel_id} already running")
            return

        adapter_cls = self._registry.get(channel_id)
        if not adapter_cls:
            raise ValueError(f"No adapter registered for channel {channel_id}")

        config = config or self._channels_config.get(channel_id, {})
        adapter = adapter_cls(config=config)
        adapter.on_message(self._handle_inbound_message)

        pre_connect_result = await adapter.pre_connect()
        if pre_connect_result:
            qr_url = pre_connect_result.get("qr_url")
            if qr_url:
                logger.info(f"Channel {channel_id} requires QR scan: {qr_url}")

        await adapter.connect()
        self._adapters[channel_id] = adapter

        if channel_id not in self._health_tasks:
            self._health_tasks[channel_id] = asyncio.create_task(self._health_check_loop(channel_id))

        logger.info(f"Channel {channel_id} started")

    async def stop_channel(self, channel_id: str) -> None:
        adapter = self._adapters.pop(channel_id, None)
        if not adapter:
            return

        task = self._health_tasks.pop(channel_id, None)
        if task:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

        try:
            await adapter.disconnect()
        except Exception:
            logger.exception(f"Error disconnecting channel {channel_id}")

        logger.info(f"Channel {channel_id} stopped")

    async def restart_channel(self, channel_id: str) -> None:
        config = self._channels_config.get(channel_id, {})
        await self.stop_channel(channel_id)
        await self.start_channel(channel_id, config)

    async def send_outbound(self, channel_id: str, response) -> None:
        adapter = self._adapters.get(channel_id)
        if not adapter:
            raise ChannelException(f"Channel {channel_id} not found", retryable=False)

        cb = self._circuit_breakers[channel_id]
        await cb.call(lambda: adapter.send(response))

    async def get_channel_status(self, channel_id: str | None = None) -> dict:
        if channel_id:
            return await self._get_single_channel_status(channel_id)

        all_channels = {}
        for cid in self._registry.list_channels():
            info = await self._get_single_channel_status(cid)
            stats = info.get("stats") or {}
            all_channels[cid] = {
                "channel_id": cid,
                "channel_type": info.get("channel_type"),
                "display_name": info.get("display_name"),
                "enabled": info.get("enabled", False),
                "status": info.get("status", "not_found"),
                "total_count": stats.get("total_messages", 0),
                "today_count": stats.get("today_messages", 0),
                "active_connections": 1 if info.get("status") == ChannelStatus.CONNECTED.value else 0,
                "health": info.get("health"),
            }

        return {"channels": all_channels}

    async def update_channel_config(self, channel_id: str, config_updates: dict[str, Any]) -> dict:
        adapter = self._adapters.get(channel_id)
        if not adapter:
            raise ChannelException(f"Channel {channel_id} not found", retryable=False)

        for key, value in config_updates.items():
            adapter.config[key] = value

        if channel_id not in self._channels_config:
            self._channels_config[channel_id] = {}
        self._channels_config[channel_id].update(config_updates)

        return {"channel_id": channel_id, "config_updated": True}

    async def test_channel(self, channel_id: str) -> dict:
        adapter = self._adapters.get(channel_id)
        if not adapter:
            return {"channel_id": channel_id, "test_result": "failure", "error": "Channel not running"}

        try:
            start = time.monotonic()
            health = await adapter.health_check()
            latency_ms = (time.monotonic() - start) * 1000

            return {
                "channel_id": channel_id,
                "test_result": "success" if health.status == "healthy" else "degraded",
                "latency_ms": round(latency_ms, 1),
                "health": health.model_dump(),
            }
        except Exception as e:
            return {"channel_id": channel_id, "test_result": "failure", "error": str(e)}

    async def check_rate_limit(self, key: str, max_req: int, window_seconds: int) -> bool:
        lock = self._rate_limit_locks.setdefault(key, asyncio.Lock())
        async with lock:
            now = time.monotonic()
            history = self._rate_limiters.setdefault(key, deque(maxlen=max_req))
            while history and now - history[0] > window_seconds:
                history.popleft()
            if len(history) >= max_req:
                return False
            history.append(now)
            return True

    async def _handle_inbound_message(self, message) -> None:
        try:
            await self._router.route_inbound(message)
        except Exception:
            logger.exception("Error handling inbound message")

    async def _health_check_loop(self, channel_id: str) -> None:
        while channel_id in self._adapters:
            await asyncio.sleep(HEALTH_CHECK_INTERVAL)
            adapter = self._adapters.get(channel_id)
            if not adapter:
                break

            cb = self._circuit_breakers.get(channel_id)
            try:
                health = await adapter.health_check()
                if cb and health.status == "healthy":
                    try:
                        await cb.record_success()
                    except Exception:
                        pass
                elif cb:
                    try:
                        await cb.record_failure()
                    except Exception:
                        pass
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.warning(f"Health check failed for {channel_id}: {e}")
                if cb:
                    try:
                        await cb.record_failure()
                    except Exception:
                        pass

    async def _get_single_channel_status(self, channel_id: str) -> dict:
        adapter = self._adapters.get(channel_id)
        if not adapter:
            adapter_cls = self._registry.get(channel_id)
            if adapter_cls:
                caps = (
                    adapter_cls.capabilities.model_dump()
                    if hasattr(adapter_cls, "capabilities")
                    else {
                        "text_chunk_limit": adapter_cls.text_chunk_limit,
                        "supports_markdown": adapter_cls.supports_markdown,
                        "supports_streaming": adapter_cls.supports_streaming,
                        "max_media_size_mb": adapter_cls.max_media_size_mb,
                    }
                )
                channel_type = adapter_cls.channel_type.value
                return {
                    "channel_id": channel_id,
                    "channel_type": channel_type,
                    "display_name": None,
                    "enabled": False,
                    "status": ChannelStatus.DISABLED.value,
                    "config": {"enabled": False},
                    "capabilities": caps,
                    "health": None,
                    "stats": None,
                }
            return {"channel_id": channel_id, "status": "not_found"}

        health = None
        try:
            health_result = await adapter.health_check()
            health = health_result.model_dump()
        except Exception:
            pass

        caps = (
            type(adapter).capabilities.model_dump()
            if hasattr(type(adapter), "capabilities")
            else {
                "text_chunk_limit": type(adapter).text_chunk_limit,
                "supports_markdown": type(adapter).supports_markdown,
                "supports_streaming": type(adapter).supports_streaming,
                "max_media_size_mb": type(adapter).max_media_size_mb,
            }
        )

        circuit_state = self._circuit_breakers.get(channel_id)
        cb_state = circuit_state.state.value if circuit_state else "unknown"

        channel_type = adapter.channel_type.value
        display_name = adapter.config.get("display_name")

        return {
            "channel_id": channel_id,
            "channel_type": channel_type,
            "display_name": display_name,
            "enabled": adapter.config.get("enabled", False) if hasattr(adapter, "config") else True,
            "status": self._adapter_status(adapter),
            "config": adapter.config,
            "capabilities": caps,
            "health": health,
            "circuit_state": cb_state,
            "stats": await self._get_channel_stats(channel_id),
        }

    async def _get_channel_stats(self, channel_id: str) -> dict:
        try:
            from yuxi.utils.datetime_utils import utc_now_naive

            now = utc_now_naive()
            today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)

            async with pg_manager.get_async_session_context() as session:
                result = await session.execute(
                    select(
                        func.count().label("total"),
                        func.sum(
                            func.cast(
                                (ChannelMsgRecord.status == "success").cast(func.Integer),
                                func.Integer,
                            )
                        ).label("success_count"),
                        func.sum(
                            func.cast(
                                (ChannelMsgRecord.status == "error").cast(func.Integer),
                                func.Integer,
                            )
                        ).label("error_count"),
                    ).where(ChannelMsgRecord.channel_id == channel_id)
                )
                row = result.one()
                total = row.total or 0
                success = row.success_count or 0
                error = row.error_count or 0

                today_result = await session.execute(
                    select(func.count()).where(
                        ChannelMsgRecord.channel_id == channel_id,
                        ChannelMsgRecord.created_at >= today_start,
                    )
                )
                today = today_result.scalar() or 0

                return {
                    "total_messages": total,
                    "today_messages": today,
                    "success_count": int(success),
                    "error_count": int(error),
                    "success_rate": round(success / total, 3) if total > 0 else 0,
                }
        except Exception:
            return {
                "total_messages": 0,
                "today_messages": 0,
                "success_count": 0,
                "error_count": 0,
                "success_rate": 0,
            }

    def _adapter_status(self, adapter: BaseChannelAdapter) -> str:
        return getattr(adapter, "status", "unknown")

    async def _ensure_virtual_department(self, db) -> None:
        from yuxi.storage.postgres.models_business import Department

        result = await db.execute(select(Department).where(Department.id == -1))
        if result.scalar_one_or_none() is None:
            dept = Department(
                id=-1,
                name="\u6e20\u9053\u7f51\u5173",
                description="\u591a\u6e20\u9053\u7f51\u5173\u865a\u62df\u90e8\u95e8",
            )
            db.add(dept)
            await db.commit()
            logger.info("Created virtual department (id=-1)")

    async def _ensure_default_agent_config(self, db) -> None:
        from yuxi.repositories.agent_config_repository import AgentConfigRepository

        repo = AgentConfigRepository(db)
        await repo.get_or_create_default(
            department_id=-1,
            agent_id="ChatbotAgent",
            created_by="system",
        )
        logger.info("Ensured default agent config for ChatbotAgent")


channel_manager = ChannelManager()