import asyncio
import base64
import json
import os
import ssl
import sys
from pathlib import Path
from typing import Any, Awaitable, Callable, List, Mapping, Optional

import aiohttp
import certifi
from yarl import URL

from ._client import Client
from ._config import _ConfigData, _load, _load_recovery_data, _load_user_config, _save
from ._core import DEFAULT_TIMEOUT
from ._errors import ConfigError
from ._login import (
    AuthNegotiator,
    AuthTokenClient,
    HeadlessNegotiator,
    _AuthToken,
    create_standalone_token,
    logout_from_browser,
)
from ._plugins import PluginManager
from ._rewrite import rewrite_module
from ._server_cfg import _ServerConfig, get_server_config
from ._tracing import _make_trace_config
from ._utils import _ContextManager

if sys.version_info >= (3, 10):
    from importlib.metadata import entry_points
else:
    from importlib_metadata import entry_points

DEFAULT_CONFIG_PATH = "~/.neuro"
CONFIG_ENV_NAME = "NEUROMATION_CONFIG"
PASS_CONFIG_ENV_NAME = "NEURO_PASSED_CONFIG"
DEFAULT_API_URL = URL("https://staging.neu.ro/api/v1")


def _make_session(
    timeout: aiohttp.ClientTimeout, trace_configs: Optional[List[aiohttp.TraceConfig]]
) -> _ContextManager[aiohttp.ClientSession]:
    return _ContextManager[aiohttp.ClientSession](
        __make_session(timeout, trace_configs)
    )


async def __make_session(
    timeout: aiohttp.ClientTimeout, trace_configs: Optional[List[aiohttp.TraceConfig]]
) -> aiohttp.ClientSession:
    from . import __version__

    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
    ssl_context.load_verify_locations(capath=certifi.where())
    connector = aiohttp.TCPConnector(ssl=ssl_context)
    return aiohttp.ClientSession(
        timeout=timeout,
        connector=connector,
        trace_configs=trace_configs,
        headers={"User-Agent": f"NeuroCLI/{__version__} ({sys.platform})"},
    )


@rewrite_module
class Factory:
    def __init__(
        self,
        path: Optional[Path] = None,
        trace_configs: Optional[List[aiohttp.TraceConfig]] = None,
        trace_id: Optional[str] = None,
        trace_sampled: Optional[bool] = None,
    ) -> None:
        if path is None:
            path = Path(os.environ.get(CONFIG_ENV_NAME, DEFAULT_CONFIG_PATH))
        self._path = path.expanduser()
        self._trace_configs = [_make_trace_config()]
        if trace_configs:
            self._trace_configs += trace_configs
        self._trace_id = trace_id
        self._trace_sampled = trace_sampled
        self._plugin_manager = PluginManager()
        for entry_point in entry_points(group="neuro_api"):
            entry_point.load()(self._plugin_manager)

    @property
    def path(self) -> Path:
        return self._path

    @property
    def is_config_present(self) -> bool:
        return (self._path / "db").exists()

    async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client:
        if not self.is_config_present and PASS_CONFIG_ENV_NAME in os.environ:
            await self.login_with_passed_config(timeout=timeout)
        try:
            return await self._get(timeout=timeout)
        except ConfigError as initial_error:
            try:
                await self._try_recover_config(timeout)
            except asyncio.CancelledError:
                raise
            except Exception:
                raise initial_error
            return await self._get(timeout=timeout)

    async def _get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client:
        session = await _make_session(timeout, self._trace_configs)
        try:
            client = Client._create(
                session,
                self._path,
                self._trace_id,
                self._trace_sampled,
                self._plugin_manager,
            )
            await client.config.check_server()
        except (asyncio.CancelledError, Exception):
            await session.close()
            raise
        else:
            return client

    async def _try_recover_config(
        self, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT
    ) -> None:
        recovery_data = _load_recovery_data(self._path)
        async with _make_session(timeout, self._trace_configs) as session:
            config_unauthorized = await get_server_config(session, recovery_data.url)
            old_token = _AuthToken.create("", 0, recovery_data.refresh_token)
            async with AuthTokenClient(
                session,
                url=config_unauthorized.auth_config.token_url,
                client_id=config_unauthorized.auth_config.client_id,
            ) as token_client:
                fresh_token = await token_client.refresh(old_token)
            config_authorized = await get_server_config(
                session, recovery_data.url, token=fresh_token.token
            )
            config = self._gen_config(config_authorized, fresh_token, recovery_data.url)
        await self.logout()  # Drop old broken config
        self._save(config)

        client = await self.get(timeout=timeout)
        await client.config.switch_cluster(recovery_data.cluster_name)
        await client.config.switch_org(recovery_data.org_name)
        await client.close()

    async def login(
        self,
        show_browser_cb: Callable[[URL], Awaitable[None]],
        *,
        url: URL = DEFAULT_API_URL,
        timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
    ) -> None:
        config_file = self._path / "db"
        if config_file.exists():
            raise ConfigError(f"Config at {self._path} already exists. Please logout")
        async with _make_session(timeout, self._trace_configs) as session:
            config_unauthorized = await get_server_config(session, url)

            if config_unauthorized.clusters:
                config_authorized = config_unauthorized
                auth_token = create_standalone_token()
            else:
                negotiator = AuthNegotiator(
                    session, config_unauthorized.auth_config, show_browser_cb
                )
                auth_token = await negotiator.get_token()

                config_authorized = await get_server_config(
                    session, url, token=auth_token.token
                )
            config = self._gen_config(config_authorized, auth_token, url)
        self._save(config)

    async def login_headless(
        self,
        get_auth_code_cb: Callable[[URL], Awaitable[str]],
        *,
        url: URL = DEFAULT_API_URL,
        timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
    ) -> None:
        config_file = self._path / "db"
        if config_file.exists():
            raise ConfigError(f"Config at {self._path} already exists. Please logout")
        async with _make_session(timeout, self._trace_configs) as session:
            config_unauthorized = await get_server_config(session, url)
            negotiator = HeadlessNegotiator(
                session, config_unauthorized.auth_config, get_auth_code_cb
            )
            auth_token = await negotiator.get_token()

            config_authorized = await get_server_config(
                session, url, token=auth_token.token
            )
        config = self._gen_config(config_authorized, auth_token, url)
        self._save(config)

    async def login_with_token(
        self,
        token: str,
        *,
        url: URL = DEFAULT_API_URL,
        timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
    ) -> None:
        config_file = self._path / "db"
        if config_file.exists():
            raise ConfigError(f"Config at {self._path} already exists. Please logout")
        async with _make_session(timeout, self._trace_configs) as session:
            server_config = await get_server_config(session, url, token=token)
        config = self._gen_config(
            server_config, _AuthToken.create_non_expiring(token), url
        )
        self._save(config)

    async def login_with_passed_config(
        self,
        config_data: Optional[str] = None,
        *,
        timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
    ) -> None:
        if config_data is None:
            try:
                config_data = os.environ[PASS_CONFIG_ENV_NAME]
            except KeyError:
                raise ConfigError(
                    f"Config env variable {PASS_CONFIG_ENV_NAME} " "is not present"
                )
        try:
            data = json.loads(base64.b64decode(config_data).decode())
            token = data["token"]
            cluster = data["cluster"]
            url = URL(data["url"])
        except (ValueError, KeyError):
            raise ConfigError(f"Data in passed config is malformed: {config_data}")
        await self.login_with_token(token, url=url, timeout=timeout)
        client = await self.get(timeout=timeout)

        await client.config.switch_cluster(cluster)
        await client.close()

    def _gen_config(
        self, server_config: _ServerConfig, token: _AuthToken, url: URL
    ) -> _ConfigData:
        from . import __version__

        if server_config.clusters:
            cluster_name = next(iter(server_config.clusters))
            org_name = server_config.clusters[cluster_name].orgs[0]
        else:
            cluster_name = None
            org_name = None
        config = _ConfigData(
            auth_config=server_config.auth_config,
            auth_token=token,
            url=url,
            admin_url=server_config.admin_url,
            version=__version__,
            cluster_name=cluster_name,
            org_name=org_name,
            clusters=server_config.clusters,
        )
        return config

    async def logout(
        self,
        show_browser_cb: Callable[[URL], Awaitable[None]] = None,
    ) -> None:
        if show_browser_cb is not None:
            try:
                old_config = _load(self._path)
            except ConfigError:
                pass  # Do not try to logout from auth0 if config is broken
            else:
                await logout_from_browser(old_config.auth_config, show_browser_cb)

        files = ["db", "db-wal", "db-shm"]
        for name in files:
            f = self._path / name
            if f.exists():
                f.unlink()
        if self._path.is_file():
            # Old-styled single file config from 2019
            self._path.unlink()
        else:
            try:
                self._path.rmdir()
            except OSError:
                # Directory Not Empty or Not A Directory
                pass

    async def load_user_config(self) -> Mapping[str, Any]:
        return _load_user_config(self._plugin_manager, self._path)

    def _save(self, config: _ConfigData) -> None:
        _save(config, self._path, False)
