from __future__ import annotations

import itertools
import time
from functools import partial

import click
from click import ClickException
from flask import current_app
from flask.cli import with_appcontext

try:
    from rich.console import Console, group
    from rich.live import Live
    from rich.pretty import Pretty
    from rich.prompt import Confirm
    from rich.table import Table
    from rich.tree import Tree

    from ._cli import (
        get_filtered_endpoint,
        limiter_theme,
        render_func,
        render_limit,
        render_limits,
        render_storage,
        render_strategy,
    )

    CLI_DEPS_AVAILABLE = True
except ImportError:  # noqa
    CLI_DEPS_AVAILABLE = False
from typing_extensions import TypedDict
from werkzeug.routing import Rule

from ._extension import Limiter
from ._limits import RuntimeLimit
from ._typing import Callable, Generator
from .constants import ConfigVars, HeaderNames


@click.group(help="Flask-Limiter maintenance & utility commands")
def cli() -> None:
    if not CLI_DEPS_AVAILABLE:  # noqa
        raise ClickException(
            "Missing dependencies for flask-limiter cli. Please install the cli extra (pip install flask-limiter[cli])"
        )


@cli.command(help="View the extension configuration")
@with_appcontext
def config() -> None:
    with current_app.test_request_context():
        console = Console(theme=limiter_theme)
        limiters = list(current_app.extensions.get("limiter", set()))
        limiter = limiters and list(limiters)[0]
        if limiter:
            extension_details = Table(title="Flask-Limiter Config")
            extension_details.add_column("Notes")
            extension_details.add_column("Configuration")
            extension_details.add_column("Value")
            extension_details.add_row("Enabled", ConfigVars.ENABLED, Pretty(limiter.enabled))
            extension_details.add_row(
                "Key Function", ConfigVars.KEY_FUNC, render_func(limiter._key_func)
            )
            extension_details.add_row(
                "Key Prefix", ConfigVars.KEY_PREFIX, Pretty(limiter._key_prefix)
            )
            limiter_config = Tree(ConfigVars.STRATEGY)
            limiter_config_values = Tree(render_strategy(limiter.limiter))
            node = limiter_config.add(ConfigVars.STORAGE_URI)
            node.add("Instance")
            node.add("Backend")
            limiter_config.add(ConfigVars.STORAGE_OPTIONS)
            limiter_config.add("Status")
            limiter_config_values.add(render_storage(limiter))
            extension_details.add_row("Rate Limiting Config", limiter_config, limiter_config_values)
            if limiter.limit_manager.application_limits:
                extension_details.add_row(
                    "Application Limits",
                    ConfigVars.APPLICATION_LIMITS,
                    Pretty(
                        [render_limit(limit) for limit in limiter.limit_manager.application_limits]
                    ),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.APPLICATION_LIMITS_PER_METHOD,
                    Pretty(limiter._application_limits_per_method),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.APPLICATION_LIMITS_EXEMPT_WHEN,
                    render_func(limiter._application_limits_exempt_when),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.APPLICATION_LIMITS_DEDUCT_WHEN,
                    render_func(limiter._application_limits_deduct_when),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.APPLICATION_LIMITS_COST,
                    Pretty(limiter._application_limits_cost),
                )
            else:
                extension_details.add_row(
                    "ApplicationLimits Limits",
                    ConfigVars.APPLICATION_LIMITS,
                    Pretty([]),
                )
            if limiter.limit_manager.default_limits:
                extension_details.add_row(
                    "Default Limits",
                    ConfigVars.DEFAULT_LIMITS,
                    Pretty([render_limit(limit) for limit in limiter.limit_manager.default_limits]),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.DEFAULT_LIMITS_PER_METHOD,
                    Pretty(limiter._default_limits_per_method),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.DEFAULT_LIMITS_EXEMPT_WHEN,
                    render_func(limiter._default_limits_exempt_when),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.DEFAULT_LIMITS_DEDUCT_WHEN,
                    render_func(limiter._default_limits_deduct_when),
                )
                extension_details.add_row(
                    None,
                    ConfigVars.DEFAULT_LIMITS_COST,
                    render_func(limiter._default_limits_cost),
                )
            else:
                extension_details.add_row("Default Limits", ConfigVars.DEFAULT_LIMITS, Pretty([]))

            if limiter._meta_limits:
                extension_details.add_row(
                    "Meta Limits",
                    ConfigVars.META_LIMITS,
                    Pretty(
                        [render_limit(limit) for limit in itertools.chain(*limiter._meta_limits)]
                    ),
                )
            if limiter._headers_enabled:
                header_configs = Tree(ConfigVars.HEADERS_ENABLED)
                header_configs.add(ConfigVars.HEADER_RESET)
                header_configs.add(ConfigVars.HEADER_REMAINING)
                header_configs.add(ConfigVars.HEADER_RETRY_AFTER)
                header_configs.add(ConfigVars.HEADER_RETRY_AFTER_VALUE)

                header_values = Tree(Pretty(limiter._headers_enabled))
                header_values.add(Pretty(limiter._header_mapping[HeaderNames.RESET]))
                header_values.add(Pretty(limiter._header_mapping[HeaderNames.REMAINING]))
                header_values.add(Pretty(limiter._header_mapping[HeaderNames.RETRY_AFTER]))
                header_values.add(Pretty(limiter._retry_after))
                extension_details.add_row(
                    "Header configuration",
                    header_configs,
                    header_values,
                )
            else:
                extension_details.add_row(
                    "Header configuration", ConfigVars.HEADERS_ENABLED, Pretty(False)
                )

            extension_details.add_row(
                "Fail on first breach",
                ConfigVars.FAIL_ON_FIRST_BREACH,
                Pretty(limiter._fail_on_first_breach),
            )
            extension_details.add_row(
                "On breach callback",
                ConfigVars.ON_BREACH,
                render_func(limiter._on_breach),
            )

            console.print(extension_details)
        else:
            console.print(
                f"No Flask-Limiter extension installed on {current_app}",
                style="bold red",
            )


@cli.command(help="Enumerate details about all routes with rate limits")
@click.option("--endpoint", default=None, help="Endpoint to filter by")
@click.option("--path", default=None, help="Path to filter by")
@click.option("--method", default=None, help="HTTP Method to filter by")
@click.option("--key", default=None, help="Test the limit")
@click.option("--watch/--no-watch", default=False, help="Create a live dashboard")
@with_appcontext
def limits(
    endpoint: str | None = None,
    path: str | None = None,
    method: str = "GET",
    key: str | None = None,
    watch: bool = False,
) -> None:
    with current_app.test_request_context():
        limiters: set[Limiter] = current_app.extensions.get("limiter", set())
        limiter: Limiter | None = list(limiters)[0] if limiters else None
        console = Console(theme=limiter_theme)
        if limiter:
            manager = limiter.limit_manager
            groups: dict[str, list[Callable[..., Tree]]] = {}

            filter_endpoint = get_filtered_endpoint(current_app, console, endpoint, path, method)
            for rule in sorted(
                current_app.url_map.iter_rules(filter_endpoint), key=lambda r: str(r)
            ):
                rule_endpoint = rule.endpoint
                if rule_endpoint == "static":
                    continue
                if len(rule_endpoint.split(".")) > 1:
                    bp_fullname = ".".join(rule_endpoint.split(".")[:-1])
                    groups.setdefault(bp_fullname, []).append(
                        partial(
                            render_limits,
                            current_app,
                            limiter,
                            manager.resolve_limits(current_app, rule_endpoint, bp_fullname),
                            rule_endpoint,
                            bp_fullname,
                            rule,
                            exemption_scope=manager.exemption_scope(
                                current_app, rule_endpoint, bp_fullname
                            ),
                            method=method,
                            test=key,
                        )
                    )
                else:
                    groups.setdefault("root", []).append(
                        partial(
                            render_limits,
                            current_app,
                            limiter,
                            manager.resolve_limits(current_app, rule_endpoint, ""),
                            rule_endpoint,
                            None,
                            rule,
                            exemption_scope=manager.exemption_scope(
                                current_app, rule_endpoint, None
                            ),
                            method=method,
                            test=key,
                        )
                    )

            @group()
            def console_renderable() -> Generator:  # type: ignore
                if limiter and limiter.limit_manager.application_limits and not (endpoint or path):
                    yield render_limits(
                        current_app,
                        limiter,
                        (list(itertools.chain(*limiter._meta_limits)), []),
                        test=key,
                        method=method,
                        label="[gold3]Meta Limits[/gold3]",
                    )
                    yield render_limits(
                        current_app,
                        limiter,
                        (limiter.limit_manager.application_limits, []),
                        test=key,
                        method=method,
                        label="[gold3]Application Limits[/gold3]",
                    )
                for name in groups:
                    if name == "root":
                        group_tree = Tree(f"[gold3]{current_app.name}[/gold3]")
                    else:
                        group_tree = Tree(f"[blue]{name}[/blue]")
                    [group_tree.add(renderable()) for renderable in groups[name]]
                    yield group_tree

            if not watch:
                console.print(console_renderable())
            else:  # noqa
                with Live(
                    console_renderable(),
                    console=console,
                    refresh_per_second=0.4,
                    screen=True,
                ) as live:
                    while True:
                        try:
                            live.update(console_renderable())
                            time.sleep(0.4)
                        except KeyboardInterrupt:
                            break
        else:
            console.print(
                f"No Flask-Limiter extension installed on {current_app}",
                style="bold red",
            )


@cli.command(help="Clear limits for a specific key")
@click.option("--endpoint", default=None, help="Endpoint to filter by")
@click.option("--path", default=None, help="Path to filter by")
@click.option("--method", default=None, help="HTTP Method to filter by")
@click.option("--key", default=None, required=True, help="Key to reset the limits for")
@click.option("-y", is_flag=True, help="Skip prompt for confirmation")
@with_appcontext
def clear(
    key: str,
    endpoint: str | None = None,
    path: str | None = None,
    method: str = "GET",
    y: bool = False,
) -> None:
    with current_app.test_request_context():
        limiters = list(current_app.extensions.get("limiter", set()))
        limiter: Limiter | None = limiters[0] if limiters else None
        console = Console(theme=limiter_theme)
        if limiter:
            manager = limiter.limit_manager
            filter_endpoint = get_filtered_endpoint(current_app, console, endpoint, path, method)

            class Details(TypedDict):
                rule: Rule
                limits: tuple[list[RuntimeLimit], ...]

            rule_limits: dict[str, Details] = {}
            for rule in sorted(
                current_app.url_map.iter_rules(filter_endpoint), key=lambda r: str(r)
            ):
                rule_endpoint = rule.endpoint
                if rule_endpoint == "static":
                    continue
                if len(rule_endpoint.split(".")) > 1:
                    bp_fullname = ".".join(rule_endpoint.split(".")[:-1])
                    rule_limits[rule_endpoint] = Details(
                        rule=rule,
                        limits=manager.resolve_limits(current_app, rule_endpoint, bp_fullname),
                    )
                else:
                    rule_limits[rule_endpoint] = Details(
                        rule=rule,
                        limits=manager.resolve_limits(current_app, rule_endpoint, ""),
                    )
            application_limits = None
            if not filter_endpoint:
                application_limits = limiter.limit_manager.application_limits

            if not y:  # noqa
                if application_limits:
                    console.print(
                        render_limits(
                            current_app,
                            limiter,
                            (application_limits, []),
                            label="Application Limits",
                            test=key,
                        )
                    )
                for endpoint, details in rule_limits.items():
                    if details["limits"]:
                        console.print(
                            render_limits(
                                current_app,
                                limiter,
                                details["limits"],
                                endpoint,
                                rule=details["rule"],
                                test=key,
                            )
                        )
            if y or Confirm.ask(f"Proceed with resetting limits for key: [danger]{key}[/danger]?"):
                if application_limits:
                    node = Tree("Application Limits")
                    for limit in application_limits:
                        limiter.limiter.clear(
                            limit.limit,
                            key,
                            limit.scope_for("", method),
                        )
                        node.add(f"{render_limit(limit)}: [success]Cleared[/success]")
                    console.print(node)
                for endpoint, details in rule_limits.items():
                    if details["limits"]:
                        node = Tree(endpoint)
                        default, decorated = details["limits"]
                        for limit in default + decorated:
                            if (
                                limit.per_method
                                and details["rule"]
                                and details["rule"].methods
                                and not method
                            ):
                                for rule_method in details["rule"].methods:
                                    limiter.limiter.clear(
                                        limit.limit,
                                        key,
                                        limit.scope_for(endpoint, rule_method),
                                    )
                            else:
                                limiter.limiter.clear(
                                    limit.limit,
                                    key,
                                    limit.scope_for(endpoint, method),
                                )
                            node.add(f"{render_limit(limit)}: [success]Cleared[/success]")
                        console.print(node)
        else:
            console.print(
                f"No Flask-Limiter extension installed on {current_app}",
                style="bold red",
            )


if __name__ == "__main__":  # noqa
    cli()