from __future__ import annotations

import contextlib
import sys
from collections.abc import AsyncIterator
from pathlib import Path
from pprint import pformat
from typing import TYPE_CHECKING, Any, Self

import click

if TYPE_CHECKING:
    from ai.backend.common.etcd import AsyncEtcd
    from ai.backend.logging import AbstractLogger
    from ai.backend.logging.types import LogLevel
    from ai.backend.manager.config.bootstrap import BootstrapConfig
    from ai.backend.manager.config.unified import ManagerUnifiedConfig

    from .types import RedisConnectionSet


class CLIContext:
    _bootstrap_config: BootstrapConfig | None
    _logger: AbstractLogger

    def __init__(self, log_level: LogLevel, config_path: Path | None = None) -> None:
        self.config_path = config_path
        self.log_level = log_level
        self._bootstrap_config = None

    async def get_bootstrap_config(self) -> BootstrapConfig:
        from ai.backend.common.config import find_config_file
        from ai.backend.common.exception import ConfigurationError
        from ai.backend.manager.config.bootstrap import BootstrapConfig

        # Lazy-load the configuration only when requested.
        try:
            if self._bootstrap_config is None:
                if self.config_path is None:
                    self.config_path = find_config_file("manager")

                self._bootstrap_config = await BootstrapConfig.load_from_file(
                    self.config_path, self.log_level
                )
        except ConfigurationError as e:
            print(
                "ConfigurationError: Could not read or validate the manager local config:",
                file=sys.stderr,
            )
            print(pformat(e.invalid_data), file=sys.stderr)
            raise click.Abort() from e
        return self._bootstrap_config

    def __enter__(self) -> Self:
        from ai.backend.logging import LocalLogger

        # The "start-server" command is injected by ai.backend.cli from the entrypoint
        # and it has its own multi-process-aware logging initialization.
        # If we duplicate the local logging with it, the process termination may hang.
        click_ctx = click.get_current_context()
        if click_ctx.invoked_subcommand != "start-server":
            self._logger = LocalLogger(log_level=self.log_level)
            self._logger.__enter__()
        return self

    def __exit__(self, *exc_info: Any) -> None:
        click_ctx = click.get_current_context()
        if click_ctx.invoked_subcommand != "start-server":
            self._logger.__exit__()


@contextlib.asynccontextmanager
async def etcd_ctx(cli_ctx: CLIContext) -> AsyncIterator[AsyncEtcd]:
    from ai.backend.common.etcd import AsyncEtcd, ConfigScopes
    from ai.backend.common.exception import ConfigurationError

    bootstrap_config = await cli_ctx.get_bootstrap_config()
    etcd_config_data = bootstrap_config.etcd.to_dataclass()
    creds = None
    if etcd_config_data.user:
        if not etcd_config_data.password:
            raise ConfigurationError({
                "etcd": "password is required when user is set",
            })

        creds = {
            "user": etcd_config_data.user,
            "password": etcd_config_data.password,
        }
    scope_prefix_map = {
        ConfigScopes.GLOBAL: "",
        # TODO: provide a way to specify other scope prefixes
    }
    async with AsyncEtcd(
        [addr.to_legacy() for addr in etcd_config_data.addrs],
        etcd_config_data.namespace,
        scope_prefix_map,
        credentials=creds,
    ) as etcd:
        yield etcd


@contextlib.asynccontextmanager
async def config_ctx(cli_ctx: CLIContext) -> AsyncIterator[ManagerUnifiedConfig]:
    from ai.backend.common.etcd import AsyncEtcd
    from ai.backend.manager.config.loader.legacy_etcd_loader import LegacyEtcdLoader
    from ai.backend.manager.config.unified import ManagerUnifiedConfig

    # scope_prefix_map is created inside ConfigServer

    bootstrap_config = await cli_ctx.get_bootstrap_config()
    etcd_config_data = bootstrap_config.etcd.to_dataclass()
    async with AsyncEtcd.create_from_config(etcd_config_data) as etcd:
        etcd_loader = LegacyEtcdLoader(etcd)
        redis_config = await etcd_loader.load()
        unified_config = ManagerUnifiedConfig(**redis_config)
    yield unified_config


@contextlib.asynccontextmanager
async def redis_ctx(cli_ctx: CLIContext) -> AsyncIterator[RedisConnectionSet]:
    from ai.backend.common.clients.valkey_client.valkey_image.client import ValkeyImageClient
    from ai.backend.common.clients.valkey_client.valkey_live.client import ValkeyLiveClient
    from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
    from ai.backend.common.clients.valkey_client.valkey_stream.client import ValkeyStreamClient
    from ai.backend.common.configs.redis import RedisConfig
    from ai.backend.common.defs import (
        REDIS_IMAGE_DB,
        REDIS_LIVE_DB,
        REDIS_STATISTICS_DB,
        REDIS_STREAM_DB,
        RedisRole,
    )
    from ai.backend.common.etcd import AsyncEtcd
    from ai.backend.manager.config.loader.legacy_etcd_loader import LegacyEtcdLoader

    from .types import RedisConnectionSet

    bootstrap_config = await cli_ctx.get_bootstrap_config()
    etcd_config_data = bootstrap_config.etcd.to_dataclass()
    async with AsyncEtcd.create_from_config(etcd_config_data) as etcd:
        loader = LegacyEtcdLoader(etcd, config_prefix="config/redis")
        raw_redis_config = await loader.load()
        redis_config = RedisConfig(**raw_redis_config)
        redis_profile_target = redis_config.to_redis_profile_target()

    valkey_live_client = await ValkeyLiveClient.create(
        redis_profile_target.profile_target(RedisRole.LIVE).to_valkey_target(),
        db_id=REDIS_LIVE_DB,
        human_readable_name="mgr_cli.live",
    )
    valkey_stat_client = await ValkeyStatClient.create(
        redis_profile_target.profile_target(RedisRole.STATISTICS).to_valkey_target(),
        db_id=REDIS_STATISTICS_DB,
        human_readable_name="mgr_cli.stat",
    )
    valkey_image_client = await ValkeyImageClient.create(
        redis_profile_target.profile_target(RedisRole.IMAGE).to_valkey_target(),
        db_id=REDIS_IMAGE_DB,
        human_readable_name="mgr_cli.image",
    )
    valkey_stream_client = await ValkeyStreamClient.create(
        redis_profile_target.profile_target(RedisRole.STREAM).to_valkey_target(),
        db_id=REDIS_STREAM_DB,
        human_readable_name="mgr_cli.stream",
    )
    yield RedisConnectionSet(
        live=valkey_live_client,
        stat=valkey_stat_client,
        image=valkey_image_client,
        stream=valkey_stream_client,
    )
    await valkey_stream_client.close()
    await valkey_image_client.close()
    await valkey_stat_client.close()
    await valkey_live_client.close()
