import base64
import importlib
import os
import sys
import textwrap
from dataclasses import dataclass
from functools import cache, partial
from importlib.metadata import version
from os import getenv
from pathlib import Path
from string import Template
from typing import List, Optional, Union

import click
import grpc
from flyteidl.service.identity_pb2 import UserInfoRequest
from flyteidl.service.identity_pb2_grpc import IdentityServiceStub
from flytekit.clients.auth_helper import get_authenticated_channel
from flytekit.configuration import AuthType, Config, PlatformConfig, get_config_file
from flytekit.configuration.default_images import DefaultImages

_DEFAULT_GCP_SERVERLESS_ENDPOINT: str = "serverless-1.us-east-2.s.union.ai"
_SERVERLESS_ENDPOINTS = {
    "serverless-1.us-east-2.s.union.ai",
    "serverless-gcp.cloud-staging.union.ai",
    "utt-srv-staging-1.cloud-staging.union.ai",
    "serverless-preview.canary.unionai.cloud",
    "utt-srv-canary-1.canary.unionai.cloud",
}

_UNION_CONFIG_ENV_VARS: List[str] = ["UNION_CONFIG", "UNIONAI_CONFIG"]
_UNION_SERVERLESS_ENDPOINT_ENV_VAR: str = "UNION_SERVERLESS_ENDPOINT"
_UNION_API_KEY_ENV_VARS: List[str] = [
    "UNION_API_KEY",
    "UNION_SERVERLESS_API_KEY",
    "UNIONAI_SERVERLESS_API_KEY",
]
_UNION_ENABLE_IMAGE_BUILDER_ENV_VAR: str = "UNION_ENABLE_REGISTER_IMAGE_BUILDER"
_UNION_DEFAULT_CONFIG_PATHS: List[Path] = [
    Path.home() / ".union" / "config.yaml",
    Path.home() / ".uctl" / "config.yaml",
]
_UNION_DEFAULT_CONFIG_PATH: Path = _UNION_DEFAULT_CONFIG_PATHS[0]


@dataclass
class _UnionConfig:
    serverless_endpoint: str = getenv(
        _UNION_SERVERLESS_ENDPOINT_ENV_VAR,
        _DEFAULT_GCP_SERVERLESS_ENDPOINT,
    )
    org: Optional[str] = None
    config: Optional[str] = None
    is_direct_unionai_cli_call: bool = False


_UNION_CONFIG = _UnionConfig()


def _is_serverless_endpoint(endpoint: str) -> str:
    """Check if endpoint is serverless."""
    return endpoint in (endpoint, f"dns:///{endpoint}") and endpoint in _SERVERLESS_ENDPOINTS


def _should_enable_image_builder(endpoint: str) -> bool:
    """Enable image builder when either holds:

    - `endpoint` is a serverless endpoint
    - UNIONAI_ENABLE_REGISTER_IMAGE_BUILDER=1

    """
    is_serverless = _is_serverless_endpoint(endpoint)
    enable_with_env = os.getenv(_UNION_ENABLE_IMAGE_BUILDER_ENV_VAR, "0") == "1"
    return is_serverless or enable_with_env


@dataclass
class AppClientCredentials:
    endpoint: str
    client_id: str
    client_secret: str
    org: str


def _encode_app_client_credentials(app_credentials: AppClientCredentials) -> str:
    """Encode app_credentials with base64."""
    data = (
        f"{app_credentials.endpoint}:{app_credentials.client_id}:{app_credentials.client_secret}:{app_credentials.org}"
    )
    return base64.b64encode(data.encode("utf-8")).decode("utf-8")


def _decode_app_client_credentials(encoded_str: str) -> AppClientCredentials:
    """Decode encoded base64 string into app credentials."""
    endpoint, client_id, client_secret, org = base64.b64decode(encoded_str.encode("utf-8")).decode("utf-8").split(":")
    return AppClientCredentials(endpoint=endpoint, client_id=client_id, client_secret=client_secret, org=org)


@cache
def _unionmeta_byoc_is_installed() -> bool:
    """Check if unionmeta_byoc is installed."""
    try:
        importlib.import_module("unionmeta_byoc")
        return True
    except ModuleNotFoundError:
        return False


def _clean_endpoint(endpoint: str) -> str:
    """Clean endpoint."""
    prefixes = ("dns:///", "http://", "https://")

    for prefix in prefixes:
        if endpoint.startswith(prefix):
            n_prefix = len(prefix)
            endpoint = endpoint[n_prefix:]
            break

    return endpoint.rstrip("/")


def _get_endpoint_for_login(endpoint: Optional[str]) -> str:
    """Get endpoint for login."""
    if _unionmeta_byoc_is_installed():
        # byoc
        if endpoint is None:
            msg = "--host must be passed in for login"
            raise click.ClickException(msg)
        return _clean_endpoint(endpoint)
    else:
        if endpoint is not None:
            return _clean_endpoint(endpoint)
        return _UNION_CONFIG.serverless_endpoint


def _get_env_var(env_vars: List[str]) -> Optional[str]:
    """Get environment variable from a list of env_vars. env_vars that are first"""
    for env_var in env_vars:
        if env_var in os.environ:
            return os.environ[env_var]
    return None


def _get_config_path(config_paths: List[Path] = _UNION_DEFAULT_CONFIG_PATHS) -> Optional[Path]:
    """Get configuration path."""
    for path in config_paths:
        if path.exists():
            return path
    return None


_get_union_config_env_var = partial(_get_env_var, env_vars=_UNION_CONFIG_ENV_VARS)
_get_union_api_env_var = partial(_get_env_var, env_vars=_UNION_API_KEY_ENV_VARS)


def _check_unionmeta_byoc(config: Optional[str]):
    """Check defaults for `pip install union[byoc]`

    - `--config` must be passed in
    - `UNION_CONFIG` must be set.
    """
    if not _unionmeta_byoc_is_installed():
        return

    if config is not None or _UNION_CONFIG.config is not None:
        # Config is passed in directly
        return

    unionai_env_var = _get_union_config_env_var()
    config_path = _get_config_path()

    if _UNION_CONFIG.is_direct_unionai_cli_call and unionai_env_var is None and config_path is None:
        msg = (
            "When using the union CLI, please pass in --config, set UNION_CONFIG or "
            "add a config to ~/.union/config.yaml"
        )

        raise click.ClickException(msg)


def _check_yaml_config_is_empty(config_file: Union[str, os.PathLike]):
    with open(config_file, "r") as f:
        contents = f.read().strip()

    if not contents:
        raise click.ClickException(
            f"Unable to load yaml at {config_file} because it is empty. "
            "Delete file or pass a config that is not empty"
        )


def _decode_union_api_key(serverless_api_value: str) -> Config:
    """Decode unionai serverless api key and return Config."""
    config = Config.auto()
    try:
        app_credentials = _decode_app_client_credentials(serverless_api_value)
    except Exception as e:
        raise ValueError("Unable to read UNION_API_KEY") from e

    if _UNION_CONFIG.org is None and app_credentials.org != "":
        _UNION_CONFIG.org = app_credentials.org

    return config.with_params(
        platform=PlatformConfig(
            endpoint=app_credentials.endpoint,
            insecure=False,
            auth_mode=AuthType.CLIENTSECRET,
            client_id=app_credentials.client_id,
            client_credentials_secret=app_credentials.client_secret,
        )
    )


def _get_config_obj(config: Optional[str] = None, default_to_unionai_semantics: bool = False) -> Config:
    """Get Config object.

    If `config` is not None, then it will be used as the Config file.

    If the `unionai` CLI is called directly or `default_to_unionai_semantics=True`, then the config is set
    in the following order:
        1. `UNION_SERVERLESS_API_KEY` environment variable
        2. `UNION_CONFIG` environment variable
        3. Serverless endpoint (Configured with `UNIONAI_SERVERLESS_ENDPOINT`)
        4. ~/.config/unionai/config.yaml if it exists

    If `pyflyte` CLI is called and `flytekit`'s `get_config_file` did not return a `ConfigFile`,
    then serverless it the default endpoint.
    """
    if config is not None:
        return Config.auto(config)

    elif _UNION_CONFIG.is_direct_unionai_cli_call or default_to_unionai_semantics:
        _check_unionmeta_byoc(config)

        # CLI ran with `unionai` -> Config is _UNION_CONFIG.config and the default is serverless
        if (serverless_api_value := _get_union_api_env_var()) is not None and serverless_api_value != "":
            return _decode_union_api_key(serverless_api_value)
        elif _UNION_CONFIG.config is not None:
            return Config.auto(_UNION_CONFIG.config)
        elif getenv(_UNION_SERVERLESS_ENDPOINT_ENV_VAR) is not None:
            # If UNIONAI_SERVERLESS_ENDPOINT is set in `_UnionConfig`, use it instead
            return Config.for_endpoint(endpoint=_UNION_CONFIG.serverless_endpoint)
        elif (config_path := _get_config_path()) is not None:
            # Check if config is empty:
            _check_yaml_config_is_empty(config_path)
            # Read config from file system
            return Config.auto(str(config_path.absolute()))
        else:
            # Got to serverless
            return Config.for_endpoint(endpoint=_UNION_CONFIG.serverless_endpoint)
    else:
        # CLI ran with `pyflyte`, we use mostly the same semantics as `pyflyte`. The difference is:
        # - byoc => Go to localhost
        # - else => Go to serverless
        cfg_file = get_config_file(_UNION_CONFIG.config)

        if cfg_file is not None:
            return Config.auto(cfg_file)

        # At this point, cfg_file is None. For byoc, we keep the flytekit behavior and go to localhost
        if _unionmeta_byoc_is_installed():
            return Config.auto(cfg_file)

        # Without byoc, then go to serverless
        return Config.for_endpoint(endpoint=_UNION_CONFIG.serverless_endpoint)


def _get_organization(platform_config: PlatformConfig, channel: Optional[grpc.Channel] = None) -> str:
    """Get organization based on endpoint."""
    if _UNION_CONFIG.org is not None:
        return _UNION_CONFIG.org
    elif _is_serverless_endpoint(platform_config.endpoint):
        org = _get_user_handle(platform_config, channel)
        _UNION_CONFIG.org = org
        return org
    else:
        # Managed+ users, the org is not required for requests and we set it ""
        # to replicate default flytekit behavior.
        return ""


def _get_user_handle(platform_config: PlatformConfig, channel: Optional[grpc.Channel] = None) -> str:
    """Get user_handle for PlatformConfig."""
    if channel is None:
        channel = get_authenticated_channel(platform_config)

    client = IdentityServiceStub(channel)
    user_info = client.UserInfo(UserInfoRequest())
    user_handle = user_info.additional_claims.fields["userhandle"]
    return user_handle.string_value


def _get_default_image() -> str:
    """Get default image version."""
    cfg_obj = _get_config_obj()

    # TODO: This is only temporary to support GCP endpoints. When the unionai images are public,
    # we will always use unionai images
    endpoint = cfg_obj.platform.endpoint
    if _is_serverless_endpoint(endpoint):
        major, minor = sys.version_info.major, sys.version_info.minor
        unionai_version = version("union")
        if "dev" in unionai_version:
            suffix = "latest"
        else:
            suffix = unionai_version

        return f"cr.union.ai/v1/unionai/union:py{major}.{minor}-{suffix}"

    return DefaultImages().find_image_for()


def _write_config_to_path(endpoint: str, auth_type: str, config_path: Path = _UNION_DEFAULT_CONFIG_PATH):
    """Write config to config directory."""
    config_dir = config_path.parent
    config_dir.mkdir(exist_ok=True, parents=True)

    config_template = Template(
        textwrap.dedent(
            """\
    admin:
      endpoint: $endpoint
      insecure: false
      authType: $auth_type
    logger:
      show-source: true
      level: 0
    union:
      connection:
        host: $endpoint
        insecure: false
      auth:
        type: $auth_type
    """
        )
    )
    config_path.write_text(config_template.substitute(endpoint=f"dns:///{endpoint}", auth_type=auth_type))
    return config_path


def _get_default_project(previous_default: str):
    cfg_obj = _get_config_obj()
    if _is_serverless_endpoint(cfg_obj.platform.endpoint):
        return "default"
    return previous_default


def _get_auth_success_html(endpoint: str) -> str:
    """Get default success html. Return None to use flytekit's default success html."""
    if endpoint.endswith("union.ai") or endpoint.endswith("unionai.cloud"):
        SUCCESS_HTML = textwrap.dedent(
            f"""
        <html>
        <head>
            <title>OAuth2 Authentication to Union Successful</title>
        </head>
        <body style="background:white;font-family:Arial">
            <div style="position: absolute;top:40%;left:50%;transform: translate(-50%, -50%);text-align:center;">
                <div style="margin:auto">
                    <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 300 65" fill="currentColor"
                        style="color:#fdb51e;width:360px;">
                        <title>Union.ai</title>
                        <path d="M32,64.8C14.4,64.8,0,51.5,0,34V3.6h17.6v41.3c0,1.9,1.1,3,3,3h23c1.9,0,3-1.1,3-3V3.6H64V34
                        C64,51.5,49.6,64.8,32,64.8z M69.9,30.9v30.4h17.6V20c0-1.9,1.1-3,3-3h23c1.9,0,3,1.1,3,3v41.3H134V30.9c0-17.5-14.4-30.8-32.1-30.8
                        S69.9,13.5,69.9,30.9z M236,30.9v30.4h17.6V20c0-1.9,1.1-3,3-3h23c1.9,0,3,1.1,3,3v41.3H300V30.9c0-17.5-14.4-30.8-32-30.8
                        S236,13.5,236,30.9L236,30.9z M230.1,32.4c0,18.2-14.2,32.5-32.2,32.5s-32-14.3-32-32.5s14-32.1,32-32.1S230.1,14.3,230.1,32.4
                        L230.1,32.4z M213.5,20.2c0-1.9-1.1-3-3-3h-24.8c-1.9,0-3,1.1-3,3v24.5c0,1.9,1.1,3,3,3h24.8c1.9,0,3-1.1,3-3V20.2z M158.9,3.6
                        h-17.6v57.8h17.6V3.6z"></path>
                    </svg>
                    <h2>You've successfully authenticated to Union!</h2>
                    <p style="font-size:20px;">Return to your terminal for next steps</p>
                </div>
            </div>
        </body>
        </html>
        """  # noqa
        )
        return SUCCESS_HTML
    return None
