from typing import Optional
import click
import functools
from lightning_cloud import env
from lightning_cloud.login import Auth
from lightning_cloud.openapi import (
    ApiClient,
    LightningappInstanceServiceApi,
    LightningappV2ServiceApi,
    Configuration,
    AuthServiceApi,
    ClusterServiceApi,
    ProjectsServiceApi,
    LightningworkServiceApi,
    SecretServiceApi,
    SSHPublicKeyServiceApi,
)
from lightning_cloud.openapi.rest import ApiException
from lightning_cloud.source_code.logs_socket_api import LightningLogsSocketAPI


def create_swagger_client(check_context=True):
    """
    Create the swagger client to use the autogenerated code

    Parameters
    ----------
    check_context: bool
        If true, check if the context is set. It's only false for APIs that
        doesn't need the context information i.e login
    """
    if check_context and not env.CONTEXT:
        raise RuntimeError(
            "Default cluster is not found. Try logging in again!")
    url = env.LIGHTNING_CLOUD_URL
    configuration = Configuration()
    configuration.host = url
    configuration.debug = env.DEBUG
    # for custom certs we need to hint urllib to use one of them if available
    # (requests package would use any of them). these two are also use
    # during artifacts to actually reverted and not use the custom certs
    # if present in /grid-cli/grid/cli/cli/artifacts.py
    configuration.ssl_ca_cert = env.SSL_CA_CERT
    api_client = ApiClient(configuration)
    api_client.default_headers["Authorization"] = Auth().authenticate()
    api_client.user_agent = f"Grid-CLI-{env.VERSION}"
    return api_client


class GridRestClient(
        LightningLogsSocketAPI,
        LightningappInstanceServiceApi,
        LightningappV2ServiceApi,
        AuthServiceApi,
        ClusterServiceApi,
        ProjectsServiceApi,
        LightningworkServiceApi,
        SecretServiceApi,
        SSHPublicKeyServiceApi,
):
    def __init__(self, api_client: Optional[ApiClient] = None):
        api_client = api_client if api_client else create_swagger_client()
        api_client.request = request_auth_warning_wrapper(api_client.request)
        super().__init__(api_client)


def request_auth_warning_wrapper(func):
    @functools.wraps(func)
    def wrap(*args, **kwargs):
        try:
            response = func(*args, **kwargs)
            return response
        except ApiException as err:
            if err.status == 401:
                raise click.ClickException(
                    "Authentication failed. Please run `lightning login`.")
            raise err

    return wrap
