ForcePilot/backend/package/yuxi/channels/adapters/matrix/accounts.py

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from enum import Enum, auto
from typing import Any

from loguru import logger


class ClientLifecycleState(Enum):
    NONE = auto()
    PREPARED = auto()
    STARTED = auto()
    STOPPED = auto()
    ERROR = auto()


class ActiveClientTracker:
    def __init__(self):
        self._client_map: dict[str, Any] = {}
        self._states: dict[str, ClientLifecycleState] = {}
        self._sync_tasks: dict[str, asyncio.Task] = {}
        self._lock = asyncio.Lock()

    @property
    def active_count(self) -> int:
        return sum(1 for s in self._states.values() if s == ClientLifecycleState.STARTED)

    def register(
        self,
        account_id: str,
        client: Any,
        state: ClientLifecycleState = ClientLifecycleState.PREPARED,
    ) -> None:
        self._client_map[account_id] = client
        self._states[account_id] = state

    def unregister(self, account_id: str) -> None:
        self._client_map.pop(account_id, None)
        self._states.pop(account_id, None)
        self._sync_tasks.pop(account_id, None)

    def get_client(self, account_id: str) -> Any | None:
        return self._client_map.get(account_id)

    def get_state(self, account_id: str) -> ClientLifecycleState:
        return self._states.get(account_id, ClientLifecycleState.NONE)

    def transition(self, account_id: str, target: ClientLifecycleState) -> bool:
        current = self.get_state(account_id)
        valid_transitions = {
            ClientLifecycleState.NONE: {ClientLifecycleState.PREPARED},
            ClientLifecycleState.PREPARED: {ClientLifecycleState.STARTED, ClientLifecycleState.ERROR},
            ClientLifecycleState.STARTED: {ClientLifecycleState.STOPPED, ClientLifecycleState.ERROR},
            ClientLifecycleState.STOPPED: {ClientLifecycleState.PREPARED},
            ClientLifecycleState.ERROR: {ClientLifecycleState.PREPARED, ClientLifecycleState.STOPPED},
        }
        if target not in valid_transitions.get(current, set()):
            logger.warning(f"[Matrix] Invalid state transition for '{account_id}': {current.name} -> {target.name}")
            return False
        self._states[account_id] = target
        return True

    def set_sync_task(self, account_id: str, task: asyncio.Task) -> None:
        self._sync_tasks[account_id] = task

    def list_active(self) -> list[str]:
        return [aid for aid, s in self._states.items() if s == ClientLifecycleState.STARTED]

    def list_all_ids(self) -> list[str]:
        return list(self._client_map.keys())


@asynccontextmanager
async def account_client_context(account_id: str, tracker: ActiveClientTracker, client_factory):
    if not tracker.transition(account_id, ClientLifecycleState.PREPARED):
        raise RuntimeError(f"Cannot prepare client for account '{account_id}'")

    client = None
    try:
        client = client_factory(account_id)
        tracker.register(account_id, client, ClientLifecycleState.PREPARED)

        if not tracker.transition(account_id, ClientLifecycleState.STARTED):
            raise RuntimeError(f"Cannot start client for account '{account_id}'")

        yield client

    except Exception:
        tracker.transition(account_id, ClientLifecycleState.ERROR)
        raise
    finally:
        try:
            tracker.transition(account_id, ClientLifecycleState.STOPPED)
        except Exception:
            pass
        if client is not None:
            try:
                await client.close()
            except Exception:
                pass


class MatrixAccount:
    def __init__(self, account_id: str, config: dict[str, Any]):
        self.account_id = account_id
        self.config = config
        self._selected = False

    @property
    def user_id(self) -> str:
        return self.config.get("user_id", "")

    @property
    def homeserver(self) -> str:
        return self.config.get("homeserver", "https://matrix.org")

    @property
    def is_default(self) -> bool:
        return self.config.get("isDefault", False)

    @property
    def enabled(self) -> bool:
        return self.config.get("enabled", True)

    def to_dict(self) -> dict[str, Any]:
        return {
            "account_id": self.account_id,
            "user_id": self.user_id,
            "homeserver": self.homeserver,
            "enabled": self.enabled,
            "is_default": self.is_default,
        }

    def __repr__(self) -> str:
        return f"MatrixAccount({self.account_id}: {self.user_id} @ {self.homeserver})"


class MultiAccountManager:
    def __init__(self, config: dict[str, Any]):
        self._config = config
        self._accounts: dict[str, MatrixAccount] = {}
        self._default_account_id: str | None = None
        self._load_accounts()

    def _load_accounts(self) -> None:
        accounts_raw = self._config.get("accounts", {})
        for account_id, account_config in accounts_raw.items():
            if isinstance(account_config, dict):
                account = MatrixAccount(account_id, account_config)
                self._accounts[account_id] = account
                if account.is_default:
                    self._default_account_id = account_id

        if not self._default_account_id and self._accounts:
            self._default_account_id = next(iter(self._accounts))

        default_account = self._config.get("defaultAccount")
        if default_account and default_account in self._accounts:
            self._default_account_id = default_account

    def get(self, account_id: str | None = None) -> MatrixAccount | None:
        if account_id:
            return self._accounts.get(account_id)
        return self.default_account

    @property
    def default_account(self) -> MatrixAccount | None:
        if self._default_account_id:
            return self._accounts.get(self._default_account_id)
        return None

    @property
    def default_account_id(self) -> str | None:
        return self._default_account_id

    def select_account(self, account_id: str) -> bool:
        if account_id in self._accounts:
            self._default_account_id = account_id
            return True
        return False

    def list_accounts(self) -> list[MatrixAccount]:
        return list(self._accounts.values())

    def resolve_room_account(self, room_id: str) -> str | None:
        rooms_config = self._config.get("rooms", {})
        room_cfg = rooms_config.get(room_id, {})
        return room_cfg.get("account")

    def promote_to_default(self, account_id: str) -> bool:
        if account_id in self._accounts:
            self._default_account_id = account_id
            self._accounts[account_id]._selected = True
            for aid, acc in self._accounts.items():
                acc.config["isDefault"] = aid == account_id
            return True
        return False

    @property
    def account_count(self) -> int:
        return len(self._accounts)

    @property
    def has_multi_account(self) -> bool:
        return len(self._accounts) > 1