import json
import pathlib
from typing import Any, Dict, List, Optional

import typer
import yaml
from rich.console import Console
from rich.table import Table
from typing_extensions import Annotated

from _nebari.config import read_configuration
from _nebari.keycloak import create_user as keycloak_create_user
from _nebari.keycloak import export_keycloak_users
from _nebari.keycloak import list_groups as keycloak_list_groups
from _nebari.keycloak import list_users as keycloak_list_users
from nebari.hookspecs import hookimpl


def _print_data_table(
    data: List[Dict[str, Any]],
    headers: List[str],
    title: str = "",
    count: Optional[int] = None,
    output_format: str = "rich",
):
    """
    Print tabular data in one of four formats:
      1) Rich table (output_format="rich", default)
      2) Plain text table (output_format="plain")
      3) YAML-like output (output_format="yaml")
      4) JSON output (output_format="json")

    Parameters
    ----------
    data : list of dict
        Each dict represents one row of data. The dict keys should match
        the strings in `headers`.
    headers : list of str
        The list of dictionary keys that define both the column order and
        the column headers. Example: ["username", "email", "groups"].
    title : str, optional
        A title to display (if using "rich") or to include in YAML/JSON output.
    count : int, optional
        If provided, a count displayed in the title (Rich),
        or printed in plain/YAML/JSON output.
    output_format : {"rich", "plain", "yaml", "json"}, default="rich"
        The format of the output.
    """

    if output_format == "yaml":
        rows = []
        for row in data:
            row_dict = {}
            for header in headers:
                row_dict[header] = row.get(header, "")
            rows.append(row_dict)

        # PyYAML to safely dump this dict as valid YAML
        print(yaml.safe_dump(rows, sort_keys=False))
        return

    if output_format == "json":
        import json

        # Build a structured dict to represent everything
        json_obj = {"title": title, "count": count, "data": []}
        for row in data:
            row_subset = {header: row.get(header, "") for header in headers}
            json_obj["data"].append(row_subset)
        # Print with indentation for readability
        print(json.dumps(json_obj, indent=2))
        return

    if output_format == "plain":
        if title:
            print(title)
        if count is not None:
            print(f"Count: {count}")

        # Print header row
        print(" | ".join(header.capitalize() for header in headers))
        print("-" * 60)

        # Print each row
        for row in data:
            row_cells = [str(row.get(header, "")) for header in headers]
            print(" | ".join(row_cells))
        return

    console = Console()
    # Append count to the title if provided
    if count:
        table_title = f"{title} (Count: {count})" if title else f"Count: {count}"
    else:
        table_title = title

    table = Table(title=table_title)
    for header in headers:
        table.add_column(header.capitalize(), style="bold")

    for row in data:
        row_cells = [str(row.get(h, "")) for h in headers]
        table.add_row(*row_cells)

    console.print(table)


@hookimpl
def nebari_subcommand(cli: typer.Typer):
    app_keycloak = typer.Typer(
        add_completion=False,
        no_args_is_help=True,
        rich_markup_mode="rich",
        context_settings={"help_option_names": ["-h", "--help"]},
    )

    cli.add_typer(
        app_keycloak,
        name="keycloak",
        help="Interact with the Nebari Keycloak identity and access management tool.",
        rich_help_panel="Additional Commands",
    )

    @app_keycloak.command(name="add-user")
    def add_user(
        user: str = typer.Option(
            ..., "--user", help="Username for the new user, must be unique"
        ),
        password: str = typer.Option(
            None,
            "-p",
            "--password",
            help="Provide a password for the user. If left blank, a password must be generated by the admin through the Keycloak interface.",
        ),
        config_filename: pathlib.Path = typer.Option(
            ...,
            "-c",
            "--config",
            help="nebari configuration file path",
        ),
        groups: Annotated[
            List[str],
            typer.Option(
                "-g",
                "--groups",
                help="Provide existing groups to add user to. Can be specified multiple times for multiple groups",
            ),
        ] = None,
        email: str = typer.Option(
            None,
            "-e",
            "--email",
            help="Email address for the new user. If not provided, a default email will be generated.",
        ),
    ):
        """Add a user to Keycloak. User will be automatically added to the [italic]analyst[/italic] group."""
        from nebari.plugins import nebari_plugin_manager

        config_schema = nebari_plugin_manager.config_schema
        config = read_configuration(config_filename, config_schema)

        if not email:
            # Generate a default email address if not provided, since its a requirement
            # for Grafana
            email = f'{user}@{config.domain or "example.com"}'

        keycloak_create_user(
            config,
            username=user,
            password=password,
            groups=groups,
            email=email,
        )

    @app_keycloak.command(name="list-users")
    def list_users(
        config_filename: pathlib.Path = typer.Option(
            ...,
            "-c",
            "--config",
            help="nebari configuration file path",
        )
    ):
        """List the users in Keycloak."""
        from nebari.plugins import nebari_plugin_manager

        config_schema = nebari_plugin_manager.config_schema
        config = read_configuration(config_filename, config_schema)
        users = keycloak_list_users(config)
        _print_data_table(
            users["users"],
            ["username", "email", "groups"],
            title="Keycloak Users",
            count=users["count"],
        )

    @app_keycloak.command(name="list-groups")
    def list_groups(
        config_filename: pathlib.Path = typer.Option(
            ...,
            "-c",
            "--config",
            help="nebari configuration file path",
        )
    ):
        """List the groups and their current roles in Keycloak."""
        from nebari.plugins import nebari_plugin_manager

        config_schema = nebari_plugin_manager.config_schema
        config = read_configuration(config_filename, config_schema)
        groups = keycloak_list_groups(config)
        _print_data_table(
            groups,
            ["name", "roles"],
            output_format="yaml",
        )

    @app_keycloak.command(name="export-users")
    def export_users(
        config_filename: pathlib.Path = typer.Option(
            ...,
            "-c",
            "--config",
            help="nebari configuration file path",
        ),
        realm: str = typer.Option(
            "nebari",
            "--realm",
            help="realm from which users are to be exported",
        ),
    ):
        """Export the users in Keycloak."""
        from nebari.plugins import nebari_plugin_manager

        config_schema = nebari_plugin_manager.config_schema
        config = read_configuration(config_filename, config_schema=config_schema)
        r = export_keycloak_users(config, realm=realm)
        print(json.dumps(r, indent=4))
