from __future__ import annotations

import json
from dataclasses import dataclass, field
from http import HTTPStatus
from sys import argv
from typing import Any, Callable, Literal
from uuid import uuid4

import click
import openapi_fal_rest.api.billing.get_user_details as get_user_details
from rich.table import Table
from rich_click import RichCommand, RichGroup

import fal
import fal.auth as auth
from fal import _serialization, api, sdk
from fal.console import console
from fal.exceptions import ApplicationExceptionHandler
from fal.logging import get_logger, set_debug_logging
from fal.logging.trace import get_tracer
from fal.rest_client import REST_CLIENT
from fal.sdk import AliasInfo, KeyScope

DEFAULT_HOST = "api.alpha.fal.ai"
HOST_ENVVAR = "FAL_HOST"

DEFAULT_PORT = "443"
PORT_ENVVAR = "FAL_PORT"

DEBUG_ENABLED = False


logger = get_logger(__name__)


@dataclass
class State:
    debug: bool = False
    invocation_id: str = field(default_factory=lambda: str(uuid4()))


def debug_option(*param_decls: str, **kwargs: Any) -> Callable[[click.FC], click.FC]:
    def callback(ctx: click.Context, param: click.Parameter, value: bool) -> None:
        state = ctx.ensure_object(State)
        state.debug = value
        set_debug_logging(value)

    if not param_decls:
        param_decls = ("--debug",)

    kwargs.setdefault("is_flag", True)
    kwargs.setdefault("expose_value", False)
    kwargs.setdefault("callback", callback)
    kwargs.setdefault("help", "Enable detailed errors and verbose logging.")
    return click.option(*param_decls, **kwargs)


class MainGroup(RichGroup):
    """A custom implementation of the top-level group
    (i.e. called on all commands and subcommands).

    This implementation allows for centralized behavior, including
    exception handling.
    """

    _exception_handler = ApplicationExceptionHandler()

    _tracer = get_tracer(__name__)

    def invoke(self, ctx):
        from click.exceptions import Abort, ClickException, Exit

        state = ctx.ensure_object(State)
        qualified_name = " ".join([ctx.info_name] + argv[1:])

        with self._tracer.start_as_current_span(
            qualified_name, attributes={"invocation_id": state.invocation_id}
        ):
            try:
                logger.debug(
                    f"Executing command: {qualified_name}",
                    command=qualified_name,
                )
                return super().invoke(ctx)
            except (EOFError, KeyboardInterrupt, ClickException, Exit, Abort):
                # let click's main handle these
                raise
            except Exception as exception:
                logger.error(exception)
                if state.debug:
                    # Here we supress detailed errors on click lines because they're
                    # mostly decorator calls, irrelevant to the dev's error tracing
                    console.print_exception(suppress=[click])
                    console.print()
                    console.print(
                        "The [markdown.code]invocation_id[/] for this operation is: "
                        f"[white]{state.invocation_id}[/]"
                    )
                else:
                    self._exception_handler.handle(exception)

    def add_command(
        self,
        cmd: RichCommand,
        name: str | None = None,
        aliases: list[str] | None = None,
    ) -> None:
        name = name or cmd.name
        assert name, "Command must have a name"

        if not aliases:
            aliases = []

        if aliases:
            # Add aliases to the help text
            cmd.help = (cmd.help or "") + "\n\nAlias: " + ", ".join([name, *aliases])
            cmd.short_help = (
                (cmd.short_help or "") + "(Alias: " + ", ".join(aliases) + ")"
            )

        super().add_command(cmd, name)
        alias_cmd = AliasCommand(cmd)

        for alias in aliases:
            self.add_command(alias_cmd, alias)


class AliasCommand(RichCommand):
    def __init__(self, wrapped):
        self._wrapped = wrapped

    def __getattribute__(self, __name: str):
        if __name == "_wrapped":
            # To be able to call `self._wrapped` below
            return super().__getattribute__(__name)

        if __name == "hidden":
            return True

        return self._wrapped.__getattribute__(__name)


@click.group(cls=MainGroup)
@click.version_option()
@debug_option()
def cli():
    pass


###### Auth group ######
@click.group(cls=RichGroup)
@debug_option()
def auth_cli():
    pass


@auth_cli.command(name="login")
@debug_option()
def auth_login():
    auth.login()


@auth_cli.command(name="logout")
@debug_option()
def auth_logout():
    auth.logout()


@auth_cli.command(name="hello", hidden=True)
@debug_option()
def auth_test():
    """
    To test auth.
    """
    print(f"Hello, {auth.USER.info['name']} - '{auth.USER.info['sub']}'")


###### Key group ######
@click.group(cls=RichGroup)
@click.option("--host", default=DEFAULT_HOST, envvar=HOST_ENVVAR)
@click.option("--port", default=DEFAULT_PORT, envvar=PORT_ENVVAR, hidden=True)
@debug_option()
@click.pass_context
def key_cli(ctx, host: str, port: str):
    ctx.obj = sdk.FalServerlessClient(f"{host}:{port}")


@key_cli.command(name="generate", no_args_is_help=True)
@click.option(
    "--scope",
    default=None,
    required=True,
    type=click.Choice([KeyScope.ADMIN.value, KeyScope.API.value]),
    help="The privilage scope of the key.",
)
@click.option(
    "--alias",
    default=None,
    help="An alias for the key.",
)
@debug_option()
@click.pass_obj
def key_generate(client: sdk.FalServerlessClient, scope: str, alias: str | None):
    with client.connect() as connection:
        parsed_scope = KeyScope(scope)
        result = connection.create_user_key(parsed_scope, alias)
        print(
            f"Generated key id and key secret, with the scope `{scope}`.\n"
            "This is the only time the secret will be visible.\n"
            "You will need to generate a new key pair if you lose access to this "
            "secret."
        )
        print(f"FAL_KEY='{result[1]}:{result[0]}'")


@key_cli.command(name="list")
@debug_option()
@click.pass_obj
def key_list(client: sdk.FalServerlessClient):
    table = Table(title="Keys")
    table.add_column("Key ID")
    table.add_column("Created At")
    table.add_column("Scope")
    table.add_column("Alias")

    with client.connect() as connection:
        keys = connection.list_user_keys()
        for key in keys:
            table.add_row(
                key.key_id, str(key.created_at), str(key.scope.value), key.alias
            )

    console.print(table)


@key_cli.command(name="revoke")
@click.argument("key_id", required=True)
@debug_option()
@click.pass_obj
def key_revoke(client: sdk.FalServerlessClient, key_id: str):
    with client.connect() as connection:
        connection.revoke_user_key(key_id)


##### Function group #####
ALIAS_AUTH_OPTIONS = ["public", "private", "shared"]
ALIAS_AUTH_TYPE = Literal["public", "private", "shared"]


@click.group(cls=RichGroup)
@click.option("--host", default=DEFAULT_HOST, envvar=HOST_ENVVAR)
@click.option("--port", default=DEFAULT_PORT, envvar=PORT_ENVVAR, hidden=True)
@debug_option()
@click.pass_context
def function_cli(ctx, host: str, port: str):
    ctx.obj = api.FalServerlessHost(f"{host}:{port}")


def load_function_from(
    host: api.FalServerlessHost,
    file_path: str,
    function_name: str,
) -> api.IsolatedFunction:
    import runpy

    module = runpy.run_path(file_path)
    if function_name not in module:
        raise api.FalServerlessError(f"Function '{function_name}' not found in module")

    # The module for the function is set to <run_path> when runpy is used, in which
    # case we want to manually include the package it is defined in.
    _serialization.include_package_from_path(file_path)

    target = module[function_name]
    if isinstance(target, type) and issubclass(target, fal.App):
        target = fal.wrap_app(target, host=host)

    if not isinstance(target, api.IsolatedFunction):
        raise api.FalServerlessError(
            f"Function '{function_name}' is not a fal.function or a fal.App"
        )
    return target


@function_cli.command("serve")
@click.option("--alias", default=None)
@click.option(
    "--auth",
    "auth_mode",
    type=click.Choice(ALIAS_AUTH_OPTIONS),
    default="private",
)
@click.argument("file_path", required=True)
@click.argument("function_name", required=True)
@debug_option()
@click.pass_obj
def register_application(
    host: api.FalServerlessHost,
    file_path: str,
    function_name: str,
    alias: str | None,
    auth_mode: ALIAS_AUTH_TYPE,
):
    user_id = _get_user_id()

    isolated_function = load_function_from(host, file_path, function_name)
    gateway_options = isolated_function.options.gateway
    if "serve" not in gateway_options and "exposed_port" not in gateway_options:
        raise api.FalServerlessError(
            "One of `serve` or `exposed_port` options needs to be specified "
            "in the isolated annotation to register a function"
        )
    elif (
        "exposed_port" in gateway_options
        and str(gateway_options.get("exposed_port")) != "8080"
    ):
        raise api.FalServerlessError(
            "Must expose port 8080 for now. This will be configurable in the future."
        )

    id = host.register(
        func=isolated_function.func,
        options=isolated_function.options,
        application_name=alias,
        application_auth_mode=auth_mode,
        metadata=isolated_function.options.host.get("metadata", {}),
    )

    if id:
        gateway_host = remove_http_and_port_from_url(host.url)
        gateway_host = (
            gateway_host.replace("api.", "").replace("alpha.", "").replace("ai", "run")
        )

        if alias:
            console.print(
                f"Registered a new revision for function '{alias}' (revision='{id}')."
            )
            console.print(f"URL: https://{gateway_host}/{user_id}/{alias}")
        else:
            console.print(f"Registered anonymous function '{id}'.")
            console.print(f"URL: https://{gateway_host}/{user_id}/{id}")


@function_cli.command("run")
@click.argument("file_path", required=True)
@click.argument("function_name", required=True)
@debug_option()
@click.pass_obj
def run(host: api.FalServerlessHost, file_path: str, function_name: str):
    isolated_function = load_function_from(host, file_path, function_name)
    # let our exc handlers handle UserFunctionException
    isolated_function.reraise = False
    isolated_function()


@function_cli.command("logs")
@click.option("--lines", default=100)
@click.option("--url", default=None)
@debug_option()
@click.pass_obj
def get_logs(
    host: api.FalServerlessHost, lines: int | None = 100, url: str | None = None
):
    console.print(
        "logs command is deprecated. To see logs, got to fal web page: https://www.fal.ai/dashboard/logs"
    )


##### Alias group #####
@click.group(cls=RichGroup)
@click.option("--host", default=DEFAULT_HOST, envvar=HOST_ENVVAR)
@click.option("--port", default=DEFAULT_PORT, envvar=PORT_ENVVAR, hidden=True)
@debug_option()
@click.pass_context
def alias_cli(ctx, host: str, port: str):
    ctx.obj = api.FalServerlessClient(f"{host}:{port}")


def _alias_table(aliases: list[AliasInfo]):
    table = Table(title="Function Aliases")
    table.add_column("Alias")
    table.add_column("Revision")
    table.add_column("Auth")
    table.add_column("Min Concurrency")
    table.add_column("Max Concurrency")
    table.add_column("Max Multiplexing")
    table.add_column("Keep Alive")
    table.add_column("Active Workers")

    for app_alias in aliases:
        table.add_row(
            app_alias.alias,
            app_alias.revision,
            app_alias.auth_mode,
            str(app_alias.min_concurrency),
            str(app_alias.max_concurrency),
            str(app_alias.max_multiplexing),
            str(app_alias.keep_alive),
            str(app_alias.active_runners),
        )

    return table


@alias_cli.command("set")
@click.argument("alias", required=True)
@click.argument("revision", required=True)
@click.option(
    "--auth",
    "auth_mode",
    type=click.Choice(ALIAS_AUTH_OPTIONS),
    default="private",
)
@debug_option()
@click.pass_obj
def alias_set(
    client: api.FalServerlessClient,
    alias: str,
    revision: str,
    auth_mode: ALIAS_AUTH_TYPE,
):
    with client.connect() as connection:
        connection.create_alias(alias, revision, auth_mode)


@alias_cli.command("delete")
@click.argument("alias", required=True)
@debug_option()
@click.pass_obj
def alias_delete(client: api.FalServerlessClient, alias: str):
    with client.connect() as connection:
        application_id = connection.delete_alias(alias)

        console.print(f"Deleted alias '{alias}' for application '{application_id}'.")


@alias_cli.command("list")
@debug_option()
@click.pass_obj
def alias_list(client: api.FalServerlessClient):
    with client.connect() as connection:
        aliases = connection.list_aliases()
        table = _alias_table(aliases)

    console.print(table)


@alias_cli.command("update")
@click.argument("alias", required=True)
@click.option("--keep-alive", "-k", type=int)
@click.option("--max-multiplexing", "-m", type=int)
@click.option("--max-concurrency", "-c", type=int)
@click.option("--min-concurrency", type=int)
# TODO: add auth_mode
# @click.option(
#     "--auth",
#     "auth_mode",
#     type=click.Choice(ALIAS_AUTH_OPTIONS),
# )
@debug_option()
@click.pass_obj
def alias_update(
    client: api.FalServerlessClient,
    alias: str,
    keep_alive: int | None,
    max_multiplexing: int | None,
    max_concurrency: int | None,
    min_concurrency: int | None,
):
    with client.connect() as connection:
        if (
            keep_alive is None
            and max_multiplexing is None
            and max_concurrency is None
            and min_concurrency is None
        ):
            console.log("No parameters for update were provided, ignoring.")
            return

        alias_info = connection.update_application(
            application_name=alias,
            keep_alive=keep_alive,
            max_multiplexing=max_multiplexing,
            max_concurrency=max_concurrency,
            min_concurrency=min_concurrency,
        )
        table = _alias_table([alias_info])

    console.print(table)


@alias_cli.command("runners")
@click.argument("alias", required=True)
@debug_option()
@click.pass_obj
def alias_list_runners(
    client: api.FalServerlessClient,
    alias: str,
):
    with client.connect() as connection:
        runners = connection.list_alias_runners(alias=alias)

    table = Table(title="Application Runners")
    table.add_column("Runner ID")
    table.add_column("In Flight Requests")
    table.add_column("Expires in")
    table.add_column("Uptime")

    for runner in runners:
        table.add_row(
            runner.runner_id,
            str(runner.in_flight_requests),
            (
                "N/A (active)"
                if not runner.expiration_countdown
                else f"{runner.expiration_countdown}s"
            ),
            f"{runner.uptime} ({runner.uptime.total_seconds()}s)",
        )

    console.print(table)


##### Secrets group #####
@click.group(cls=RichGroup)
@click.option("--host", default=DEFAULT_HOST, envvar=HOST_ENVVAR)
@click.option("--port", default=DEFAULT_PORT, envvar=PORT_ENVVAR, hidden=True)
@debug_option()
@click.pass_context
def secrets_cli(ctx, host: str, port: str):
    ctx.obj = sdk.FalServerlessClient(f"{host}:{port}")


@secrets_cli.command("list")
@debug_option()
@click.pass_obj
def list_secrets(client: api.FalServerlessClient):
    table = Table(title="Secrets")
    table.add_column("Secret Name")
    table.add_column("Created At")

    with client.connect() as connection:
        for secret in connection.list_secrets():
            table.add_row(secret.name, str(secret.created_at))

    console.print(table)


@secrets_cli.command("set")
@click.argument("secret_name", required=True)
@click.argument("secret_value", required=True)
@debug_option()
@click.pass_obj
def set_secret(client: api.FalServerlessClient, secret_name: str, secret_value: str):
    with client.connect() as connection:
        connection.set_secret(secret_name, secret_value)
        console.print(f"Secret '{secret_name}' has set")


@secrets_cli.command("delete")
@click.argument("secret_name", required=True)
@debug_option()
@click.pass_obj
def delete_secret(client: api.FalServerlessClient, secret_name: str):
    with client.connect() as connection:
        connection.delete_secret(secret_name)
        console.print(f"Secret '{secret_name}' has deleted")


# Setup of groups
cli.add_command(auth_cli, name="auth")
cli.add_command(key_cli, name="key", aliases=["keys"])
cli.add_command(function_cli, name="function", aliases=["fn"])
cli.add_command(alias_cli, name="alias", aliases=["aliases"])
cli.add_command(secrets_cli, name="secret", aliases=["secrets"])


def remove_http_and_port_from_url(url):
    # Remove http://
    if "http://" in url:
        url = url.replace("http://", "")

    # Remove https://
    if "https://" in url:
        url = url.replace("https://", "")

    # Remove port information
    url_parts = url.split(":")
    if len(url_parts) > 1:
        url = url_parts[0]

    return url


def _get_user_id() -> str:
    try:
        user_details_response = get_user_details.sync_detailed(
            client=REST_CLIENT,
        )
    except Exception as e:
        raise api.FalServerlessError(f"Error fetching user details: {str(e)}")

    if user_details_response.status_code != HTTPStatus.OK:
        try:
            content = json.loads(user_details_response.content.decode("utf8"))
        except Exception:
            raise api.FalServerlessError(
                f"Error fetching user details: {user_details_response}"
            )
        else:
            raise api.FalServerlessError(content["detail"])
    try:
        full_user_id = user_details_response.parsed.user_id
        _provider, _, user_id = full_user_id.partition("|")
        if not user_id:
            user_id = full_user_id

        return user_id
    except Exception as e:
        raise api.FalServerlessError(f"Could not parse the user data: {e}")
