import argparse
import importlib
import pathlib
import logging
import re
import sys
import shutil
from typing import List, Dict

from {{ master_module_name }}.version import __version__ as version

ROOT_PATH = pathlib.Path(__file__).absolute().parent
BOTO3_STUBS_NAME = "{{ boto3_stubs_name }}"
MODULE_NAME = "{{ master_module_name }}"
BOTO3_STUBS_PATH = ROOT_PATH.parent / BOTO3_STUBS_NAME
INIT_PATH = BOTO3_STUBS_PATH / "__init__.pyi"
SESSION_PATH = BOTO3_STUBS_PATH / "session.pyi"
SUBMODULES = [
    {% for service_package in package.service_packages -%}
    dict(
        module_name="{{ service_package.service_name.module_name }}",
        import_name="{{ service_package.service_name.import_name }}",
        boto3_name="{{ service_package.service_name.boto3_name }}",
        class_name="{{ service_package.service_name.class_name }}",
        has_resource={{ "True" if service_package.service_resource else "False" }},
        has_waiter={{ "True" if service_package.waiters else "False" }},
        has_paginator={{ "True" if service_package.paginators else "False" }}
    )
    {{ ",\n" if not loop.last else "\n" -}}
    {% endfor -%}
]

FUNCTION_TEMPLATE = """{overload}def {name}(
    service_name: {service_name_type},
    region_name: str = None,
    api_version: str = None,
    use_ssl: bool = None,
    verify: Union[str, bool] = None,
    endpoint_url: str = None,
    aws_access_key_id: str = None,
    aws_secret_access_key: str = None,
    aws_session_token: str = None,
    config: Config = None,
) -> {return_type}: ..."""

METHOD_TEMPLATE = """    {overload}def {name}(
        self,
        service_name: {service_name_type},
        region_name: str = None,
        api_version: str = None,
        use_ssl: bool = None,
        verify: Union[str, bool] = None,
        endpoint_url: str = None,
        aws_access_key_id: str = None,
        aws_secret_access_key: str = None,
        aws_session_token: str = None,
        config: Config = None,
    ) -> {return_type}: ..."""


def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        MODULE_NAME,
        description=(
            "Service stubs generator for boto3-stubs."
            " Run it after boto3-stubs every update, install or remove services."
            " Use it only with local python installation."
        ),
    )
    parser.add_argument("-v", "--version", action="version", version=version)
    parser.add_argument("-d", "--debug", action="store_true", help="Hide log output")
    parser.add_argument("-q", "--quiet", action="store_true", help="Verbose log output")
    return parser


def get_logger(level: int) -> logging.Logger:
    logger = logging.getLogger("{{ master_module_name }}")
    stream_handler = logging.StreamHandler()
    formatter = logging.Formatter("%(name)s: %(levelname)-8s %(message)s")
    stream_handler.setFormatter(formatter)
    stream_handler.setLevel(level)
    logger.addHandler(stream_handler)
    logger.setLevel(level)
    return logger


def main() -> None:
    parser = get_parser()
    args = parser.parse_args()
    log_level = logging.INFO
    if args.debug:
        log_level = logging.DEBUG
    if args.quiet:
        log_level = logging.CRITICAL
    logger = get_logger(log_level)

    if not BOTO3_STUBS_PATH.is_dir():
        logger.error(
            (
                "Directory %s does not exist, make"
                " sure boto3-stubs are installed to the same location."
            ),
            BOTO3_STUBS_PATH
        )
        sys.exit(1)

    init_client_functions: List[str] = []
    init_resource_functions: List[str] = []
    session_client_functions: List[str] = []
    session_resource_functions: List[str] = []
    imports: List[str] = []
    init_contents = INIT_PATH.read_text()
    session_contents = SESSION_PATH.read_text()
    discovered_submodules: List[Dict] = []
    for submodule in SUBMODULES:
        if importlib.find_loader(submodule["module_name"]) is None:
            continue

        discovered_submodules.append(submodule)
        logger.info(
            "Discovered %s subresource in %s",
            submodule["boto3_name"],
            submodule["module_name"],
        )

    for submodule in discovered_submodules:
        init_client_functions.append(
            FUNCTION_TEMPLATE.format(
                overload="@overload\n" if len(discovered_submodules) > 1 else "",
                name="client",
                service_name_type='Literal["{}"]'.format(submodule["boto3_name"]),
                return_type="{}Client".format(submodule["class_name"]),
            )
        )
        session_client_functions.append(
            METHOD_TEMPLATE.format(
                overload="@overload\n    " if len(discovered_submodules) > 1 else "",
                name="client",
                service_name_type='Literal["{}"]'.format(submodule["boto3_name"]),
                return_type="{}Client".format(submodule["class_name"]),
            )
        )
        imports.append(
            "from {{ master_module_name }}.{} import Client as {}Client".format(
                submodule["import_name"], submodule["class_name"],
            )
        )
        if submodule["has_resource"]:
            init_resource_functions.append(
                FUNCTION_TEMPLATE.format(
                    overload="@overload\n" if len(discovered_submodules) > 1 else "",
                    name="resource",
                    service_name_type='Literal["{}"]'.format(submodule["boto3_name"]),
                    return_type="{}ServiceResource".format(submodule["class_name"]),
                )
            )
            session_resource_functions.append(
                METHOD_TEMPLATE.format(
                    overload="@overload\n    " if len(discovered_submodules) > 1 else "",
                    name="resource",
                    service_name_type='Literal["{}"]'.format(submodule["boto3_name"]),
                    return_type="{}ServiceResource".format(submodule["class_name"]),
                )
            )
            imports.append(
                "from {{ master_module_name }}.{} import ServiceResource as {}ServiceResource".format(
                    submodule["import_name"], submodule["class_name"],
                )
            )

    if not init_client_functions:
        init_client_functions.append(
            FUNCTION_TEMPLATE.format(
                overload="",
                name="client",
                service_name_type="str",
                return_type="Any",
            )
        )
    if not init_resource_functions:
        init_resource_functions.append(
            FUNCTION_TEMPLATE.format(
                overload="",
                name="resource",
                service_name_type="str",
                return_type="Any",
            )
        )
    if not session_client_functions:
        session_client_functions.append(
            METHOD_TEMPLATE.format(
                overload="",
                name="client",
                service_name_type="str",
                return_type="Any",
            )
        )
    if not session_resource_functions:
        session_resource_functions.append(
            METHOD_TEMPLATE.format(
                overload="",
                name="resource",
                service_name_type="str",
                return_type="Any",
            )
        )

    imports_string = "### IMPORTS\n{}\n### ENDIMPORTS".format(
        "\n".join(imports) if imports else "\n"
    )
    init_contents = re.sub(
        r"### IMPORTS.+### ENDIMPORTS", imports_string, init_contents, flags=re.S
    )
    session_contents = re.sub(
        r"### IMPORTS.+### ENDIMPORTS", imports_string, session_contents, flags=re.S
    )

    client_string = "### CLIENT\n{}\n### ENDCLIENT".format(
        "\n".join(init_client_functions) if init_client_functions else "\n"
    )
    session_client_string = "### CLIENT\n{}\n    ### ENDCLIENT".format(
        "\n".join(session_client_functions) if session_client_functions else "\n"
    )
    init_contents = re.sub(
        r"### CLIENT.+### ENDCLIENT", client_string, init_contents, flags=re.S
    )
    session_contents = re.sub(
        r"### CLIENT.+### ENDCLIENT", session_client_string, session_contents, flags=re.S
    )

    resource_string = "### RESOURCE\n{}\n### ENDRESOURCE".format(
        "\n".join(init_resource_functions) if init_resource_functions else "\n"
    )
    session_resource_string = "### RESOURCE\n{}\n    ### ENDRESOURCE".format(
        "\n".join(session_resource_functions) if session_resource_functions else "\n"
    )
    init_contents = re.sub(
        r"### RESOURCE.+### ENDRESOURCE", resource_string, init_contents, flags=re.S
    )
    session_contents = re.sub(
        r"### RESOURCE.+### ENDRESOURCE", session_resource_string, session_contents, flags=re.S
    )

    def write_text(path: pathlib.Path, text: str) -> None:
        logger.debug("Updating %s", path)
        path.write_text(text)

    for submodule in SUBMODULES:
        submodule_path = ROOT_PATH / submodule["import_name"]
        if submodule not in discovered_submodules:
            if submodule_path.exists():
                logger.debug("Removing directory for deleted service %s", submodule_path)
                shutil.rmtree(submodule_path)
            continue

        if not submodule_path.exists():
            logger.debug("Creating directory %s", submodule_path)
            submodule_path.mkdir(exist_ok=True)

        write_text(submodule_path / '__init__.py', 'from {} import *'.format(submodule["module_name"]))
        write_text(submodule_path / 'client.py', 'from {}.client import *'.format(submodule["module_name"]))
        write_text(submodule_path / 'type_defs.py', 'from {}.type_defs import *'.format(submodule["module_name"]))
        if submodule["has_resource"]:
            write_text(
                submodule_path / 'service_resource.py',
                'from {}.service_resource import *'.format(submodule["module_name"]),
            )
        if submodule["has_waiter"]:
            write_text(
                submodule_path / 'waiter.py',
                'from {}.waiter import *'.format(submodule["module_name"]),
            )
        if submodule["has_paginator"]:
            write_text(
                submodule_path / 'paginator.py',
                'from {}.paginator import *'.format(submodule["module_name"]),
            )

    write_text(INIT_PATH, init_contents)
    write_text(SESSION_PATH, session_contents)

    if not discovered_submodules:
        logger.warning("No services submodules discovered, install the ones you use and run this command again")
        logger.info("https://mypy-boto3.readthedocs.io/en/latest/#sub-modules")
        sys.exit()

    names_str = ""
    for i, submodule in enumerate(discovered_submodules):
        name = submodule["boto3_name"]
        if not i:
            names_str = name
            continue
        if i < len(discovered_submodules) - 1:
            names_str = "{}, {}".format(names_str, name)
            continue

        names_str = "{} and {}".format(names_str, name)

    logger.info("You can now use %s type annotations.", names_str)
