# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

from enum import Enum
import importlib
import inspect
import logging
from pathlib import Path
import pkgutil
import shutil
import sys
import tempfile

from msrest.serialization import Model
from msrest.exceptions import HttpOperationError
from msrest.paging import Paged

_LOGGER = logging.getLogger(__name__)


copyright_header = b"""# coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
#
# Code generated by Microsoft (R) AutoRest Code Generator.
# Changes may cause incorrect behavior and will be lost if the code is
# regenerated.
# --------------------------------------------------------------------------

"""

header = copyright_header+b"""from msrest.serialization import Model
from msrest.exceptions import HttpOperationError
"""

paging_header = copyright_header+b"""from msrest.paging import Paged
"""

init_file = """
try:
    from .{} import *
except (SyntaxError, ImportError):
    from .{} import *
from .{} import *
"""

MODEL_PY2_NAME = "_models"
MODEL_PY3_NAME = "_models_py3"
PAGED_NAME = "_paged_models"

def as_file_name(name):
    return name + ".py"

def parse_input(input_parameter):
    """From a syntax like package_name#submodule, build a package name
    and complete module name.
    """
    split_package_name = input_parameter.split('#')
    package_name = split_package_name[0]
    module_name = package_name.replace("-", ".")
    if len(split_package_name) >= 2:
        module_name = ".".join([module_name, split_package_name[1]])
    return package_name, module_name

def solve_mro(models):
    for models_module in models:
        models_path = models_module.__path__[0]
        _LOGGER.info(f"Working on {models_path}")
        if Path(models_path, as_file_name(MODEL_PY3_NAME)).exists():
            _LOGGER.info("Skipping since already patched")
            return

        # Build the new files in a temp folder
        with tempfile.TemporaryDirectory() as temp_folder:
            final_models_path = Path(temp_folder, "models")
            final_models_path.mkdir()
            solve_one_model(models_module, final_models_path)

            # Swith the files
            shutil.rmtree(models_path)
            shutil.move(final_models_path, models_path)

def solve_one_model(models_module, output_folder):
    """Will build the compacted models in the output_folder"""

    models_classes = [
        (len(model_class.__mro__), model_name, inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items()
        if model_name[0].isupper() and Model in model_class.__mro__
    ]
    # Sort on MRO size first, and then alphabetically
    models_classes.sort(key=lambda x: (x[0], x[1]))

    # Just need the name of exceptions
    exceptions_classes = [
        model_name for model_name, model_class in vars(models_module).items()
        if model_name[0].isupper() and HttpOperationError in model_class.__mro__
    ]

    py2_models_classes = [
        (len_mro, model_name, path.replace("_py3.py", ".py"), None)
        for len_mro, model_name, path, _ in models_classes
    ]

    paged_models_classes = [
        (model_name, inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items()
        if model_name[0].isupper() and Paged in model_class.__mro__
    ]

    enum_models_classes = [
        (model_name, inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items()
        if model_name[0].isupper() and Enum in model_class.__mro__
    ]
    if enum_models_classes:
        # Can't be more than one enum file
        enum_file = Path(enum_models_classes[0][1])
        enum_file_module_name = "_"+enum_file.with_suffix('').name
        shutil.copyfile(enum_file, Path(output_folder, as_file_name(enum_file_module_name)))
    else:
        enum_file_module_name = None

    write_model_file(Path(output_folder, as_file_name(MODEL_PY3_NAME)), models_classes)
    write_model_file(Path(output_folder, as_file_name(MODEL_PY2_NAME)), py2_models_classes)
    write_paging_file(Path(output_folder, as_file_name(PAGED_NAME)), paged_models_classes)
    write_complete_init(
        Path(output_folder, "__init__.py"),
        models_classes,
        exceptions_classes,
        paged_models_classes,
        enum_models_classes,
        enum_file_module_name
    )

def write_model_file(output_file_path, classes_to_write):
    with open(output_file_path, "bw") as write_fd:
        write_fd.write(header)

        for model in classes_to_write:
            _, _, model_file_path, _ = model

            with open(model_file_path, "rb") as read_fd:
                lines = read_fd.readlines()
                # Skip until it's "class XXXX"
                while lines:
                    if lines[0].startswith(b"class "):
                        break
                    lines.pop(0)
                else:
                    raise ValueError("Never found any class definition!")
                # Now I keep everything
                write_fd.write(b'\n')
                write_fd.write(b'\n')
                write_fd.writelines(lines)

def write_paging_file(output_file_path, classes_to_write):
    with open(output_file_path, "bw") as write_fd:
        write_fd.write(paging_header)

        for model in classes_to_write:
            _, model_file_path, _ = model

            with open(model_file_path, "rb") as read_fd:
                # Skip the first 15 lines (based on Autorest deterministic behavior)
                # If we want this less random, look for the first line starts with "class"
                lines = read_fd.readlines()[14:]
                write_fd.write(b'\n')
                write_fd.write(b'\n')
                write_fd.writelines(lines)

def write_init(output_file_path, model_file_name, model_file_name_py2, paging_file_name, enum_file_name):
    with open(output_file_path, "bw") as write_fd:
        write_fd.write(copyright_header)

        write_fd.write(init_file.format(
            model_file_name,
            model_file_name_py2,
            paging_file_name,
        ).encode('utf8'))
        if enum_file_name:
            write_fd.write(
                "from .{} import *".format(enum_file_name).encode('utf8')
            )

def write_complete_init(output_file_path, models, exceptions_classes, paging_models, enum_models, enum_file_module_name):
    with open(output_file_path, "bw") as write_fd:
        write_fd.write(copyright_header)

        write_fd.write(b"\ntry:\n")
        # Write py3 import
        for _, model_name, _, _ in models:
            write_fd.write(f"    from .{MODEL_PY3_NAME} import {model_name}\n".encode("utf-8"))
        # Py 3 exceptions
        for model_name in exceptions_classes:
            write_fd.write(f"    from .{MODEL_PY3_NAME} import {model_name}\n".encode("utf-8"))

        write_fd.write(b"except (SyntaxError, ImportError):\n")

        # Write py2 import
        for _, model_name, _, _ in models:
            write_fd.write(f"    from .{MODEL_PY2_NAME} import {model_name}\n".encode("utf-8"))
        # Py 2 exceptions
        for model_name in exceptions_classes:
            write_fd.write(f"    from .{MODEL_PY2_NAME} import {model_name}\n".encode("utf-8"))

        # Write paged import
        for model_name, _, _ in paging_models:
            write_fd.write(f"from .{PAGED_NAME} import {model_name}\n".encode("utf-8"))

        if enum_models:
            # Write enum model import
            for model_name, _, _ in enum_models:
                write_fd.write(f"from .{enum_file_module_name} import {model_name}\n".encode("utf-8"))

        write_fd.write(b"\n\n__all__=[\n")
        # Write all classes name
        for _, model_name, _, _ in models:
            write_fd.write(f"    '{model_name}',\n".encode("utf-8"))
        for model_name in exceptions_classes:
            write_fd.write(f"    '{model_name}',\n".encode("utf-8"))
        for model_name, _, _ in paging_models:
            write_fd.write(f"    '{model_name}',\n".encode("utf-8"))
        if enum_models:
            for model_name, _, _ in enum_models:
                write_fd.write(f"    '{model_name}',\n".encode("utf-8"))

        write_fd.write(b"]\n")


def find_models_to_change(module_name):
    """Will figure out if the package is a multi-api one,
    and understand what to generate.
    """
    main_module = importlib.import_module(module_name)
    try:
        models_module = main_module.models
        models_module.__path__
        # It didn't fail, that's a single API package
        return [models_module]
    except AttributeError:
        # This means I loaded the fake module "models"
        # and it's multi-api, load all models
        return [
            importlib.import_module('.'+label+'.models', main_module.__name__)
            for (_, label, ispkg) in pkgutil.iter_modules(main_module.__path__)
            if ispkg and label != 'aio'
        ]


def find_autorest_generated_folder(module_prefix="azure.mgmt"):
    """Find all Autorest generated code in that module prefix.

    This actually looks for a "models" package only. We could be smarter if necessary.
    """
    _LOGGER.info(f"Looking for Autorest generated package in {module_prefix}")
    result = []
    try:
        _LOGGER.debug(f"Try {module_prefix}")
        importlib.import_module(".models", module_prefix)
        # If not exception, we found it
        _LOGGER.info(f"Found {module_prefix}")
        result.append(module_prefix)
    except ModuleNotFoundError:
        # No model, might dig deeper
        prefix_module = importlib.import_module(module_prefix)
        for _, sub_package, ispkg in pkgutil.iter_modules(prefix_module.__path__, module_prefix+"."):
            if ispkg:
                result += find_autorest_generated_folder(sub_package)
    return result


def main(prefix="azure.mgmt"):
    packages = find_autorest_generated_folder(prefix)
    for autorest_package in packages:
        models_module = find_models_to_change(autorest_package)
        solve_mro(models_module)

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main(sys.argv[1] if len(sys.argv) >= 2 else "azure.mgmt")