import asyncio
import logging
import multiprocessing
import os
from pathlib import Path
from pprint import pformat, pprint
from setproctitle import setproctitle
import ssl
import sys
from typing import (
    Any,
    AsyncIterator,
    Sequence,
)

from aiohttp import web
import aiotools
import click

from ai.backend.common import config
from ai.backend.common.etcd import AsyncEtcd, ConfigScopes
from ai.backend.common.logging import Logger, BraceStyleAdapter
from ai.backend.common.utils import env_info

from . import __version__ as VERSION
from .config import local_config_iv
from .context import Context
from .api.client import init_client_app
from .api.manager import init_manager_app

log = BraceStyleAdapter(logging.getLogger('ai.backend.storage.server'))


@aiotools.server
async def server_main_logwrapper(loop, pidx, _args):
    setproctitle(f"backend.ai: storage-proxy worker-{pidx}")
    log_endpoint = _args[1]
    logger = Logger(_args[0]['logging'], is_master=False, log_endpoint=log_endpoint)
    with logger:
        async with server_main(loop, pidx, _args):
            yield


@aiotools.server
async def server_main(
    loop: asyncio.AbstractEventLoop,
    pidx: int,
    _args: Sequence[Any],
) -> AsyncIterator[None]:
    config = _args[0]

    etcd_credentials = None
    if config['etcd']['user']:
        etcd_credentials = {
            'user': config['etcd']['user'],
            'password': config['etcd']['password'],
        }
    scope_prefix_map = {
        ConfigScopes.GLOBAL: "",
        ConfigScopes.NODE: f"nodes/storage/{config['storage-proxy']['node-id']}",
    }
    etcd = AsyncEtcd(config['etcd']['addr'],
                     config['etcd']['namespace'],
                     scope_prefix_map,
                     credentials=etcd_credentials)
    ctx = Context(pid=os.getpid(), local_config=config, etcd=etcd)
    client_api_app = await init_client_app(ctx)
    manager_api_app = await init_manager_app(ctx)

    client_ssl_ctx = None
    manager_ssl_ctx = None
    if config['api']['client']['ssl-enabled']:
        client_ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        client_ssl_ctx.load_cert_chain(
            str(config['api']['client']['ssl-cert']),
            str(config['api']['client']['ssl-privkey']),
        )
    if config['api']['manager']['ssl-enabled']:
        manager_ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        manager_ssl_ctx.load_cert_chain(
            str(config['api']['manager']['ssl-cert']),
            str(config['api']['manager']['ssl-privkey']),
        )
    client_api_runner = web.AppRunner(client_api_app)
    manager_api_runner = web.AppRunner(manager_api_app)
    await client_api_runner.setup()
    await manager_api_runner.setup()
    client_service_addr = config['api']['client']['service-addr']
    manager_service_addr = config['api']['manager']['service-addr']
    client_api_site = web.TCPSite(
        client_api_runner,
        str(client_service_addr.host),
        client_service_addr.port,
        backlog=1024,
        reuse_port=True,
        ssl_context=client_ssl_ctx,
    )
    manager_api_site = web.TCPSite(
        manager_api_runner,
        str(manager_service_addr.host),
        manager_service_addr.port,
        backlog=1024,
        reuse_port=True,
        ssl_context=manager_ssl_ctx,
    )
    await client_api_site.start()
    await manager_api_site.start()
    log.info('Started service.')
    try:
        yield
    finally:
        log.info('Shutting down...')
        await manager_api_runner.cleanup()
        await client_api_runner.cleanup()


@click.group(invoke_without_command=True)
@click.option('-f', '--config-path', '--config', type=Path, default=None,
              help="The config file path. "
                   "(default: ./storage-proxy.toml and /etc/backend.ai/storage-proxy.toml)")
@click.option('--debug', is_flag=True,
              help="Enable the debug mode and override the global log level to DEBUG.")
@click.pass_context
def main(cli_ctx, config_path, debug):
    # Determine where to read configuration.
    raw_cfg, cfg_src_path = config.read_from_file(config_path, 'storage-proxy')

    config.override_with_env(raw_cfg, ('etcd', 'namespace'), 'BACKEND_NAMESPACE')
    config.override_with_env(raw_cfg, ('etcd', 'addr'), 'BACKEND_ETCD_ADDR')
    config.override_with_env(raw_cfg, ('etcd', 'user'), 'BACKEND_ETCD_USER')
    config.override_with_env(raw_cfg, ('etcd', 'password'), 'BACKEND_ETCD_PASSWORD')
    if debug:
        config.override_key(raw_cfg, ('debug', 'enabled'), True)

    try:
        local_config = config.check(raw_cfg, local_config_iv)
        local_config['_src'] = cfg_src_path
    except config.ConfigurationError as e:
        print('ConfigurationError: Validation of agent configuration has failed:', file=sys.stderr)
        print(pformat(e.invalid_data), file=sys.stderr)
        raise click.Abort()

    if local_config['debug']['enabled']:
        config.override_key(local_config, ('logging', 'level'), 'DEBUG')
        config.override_key(local_config, ('logging', 'pkg-ns', 'ai.backend'), 'DEBUG')

    # if os.getuid() != 0:
    #     print('Storage agent can only be run as root', file=sys.stderr)
    #     raise click.Abort()

    multiprocessing.set_start_method('spawn')

    if cli_ctx.invoked_subcommand is None:
        local_config['storage-proxy']['pid-file'].write_text(str(os.getpid()))
        log_sockpath = Path(f'/tmp/backend.ai/ipc/storage-proxy-logger-{os.getpid()}.sock')
        log_sockpath.parent.mkdir(parents=True, exist_ok=True)
        log_endpoint = f'ipc://{log_sockpath}'
        local_config['logging']['endpoint'] = log_endpoint
        try:
            logger = Logger(local_config['logging'], is_master=True, log_endpoint=log_endpoint)
            with logger:
                setproctitle('backend.ai: storage-proxy')
                log.info('Backend.AI Storage Proxy', VERSION)
                log.info('Runtime: {0}', env_info())
                log.info('Node ID: {0}', local_config['storage-proxy']['node-id'])
                log_config = logging.getLogger('ai.backend.agent.config')
                if local_config['debug']['enabled']:
                    log_config.debug('debug mode enabled.')
                if 'debug' in local_config and local_config['debug']['enabled']:
                    print('== Storage proxy configuration ==')
                    pprint(local_config)
                if local_config['storage-proxy']['event-loop'] == 'uvloop':
                    import uvloop
                    uvloop.install()
                    log.info('Using uvloop as the event loop backend')
                aiotools.start_server(
                    server_main_logwrapper,
                    use_threading=False,
                    num_workers=local_config['storage-proxy']['num-proc'],
                    args=(local_config, log_endpoint),
                )
                log.info('exit.')
        finally:
            if local_config['storage-proxy']['pid-file'].is_file():
                # check is_file() to prevent deleting /dev/null!
                local_config['storage-proxy']['pid-file'].unlink()
    return 0


if __name__ == "__main__":
    sys.exit(main())
