from functools import wraps
from typing import Callable, overload
from .model.descriptors import BotCommand, Command, CommandCallbackContext


class Router:
    def __init__(self):
        self._commands = dict[str, BotCommand]()

    @overload
    def command(self, command: Command): ...

    @overload
    def command(self, command: str, caption: str | dict[str, str] | None = None): ...

    def command(
        self, command: str | Command, caption: str | dict[str, str] | None = None
    ):
        def decorator(func: Callable[[CommandCallbackContext], None]):
            if isinstance(command, str):
                cmd = BotCommand(name=command, handler=func, caption=caption)
            else:
                cmd = BotCommand(handler=func, **command.__dict__)
            self._commands[cmd.name] = cmd

            @wraps(func)
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)

            return wrapper

        return decorator