import io
import os
import re
import shutil
import subprocess
import sys
from contextlib import contextmanager
from pathlib import Path
from types import FunctionType, ModuleType
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Type,
    Union,
)

import build
import dill
import toml
from build.__main__ import build_package, build_package_via_sdist
from confit import Cli
from dill._dill import save_function as dill_save_function
from dill._dill import save_module as dill_save_module
from dill._dill import save_type as dill_save_type

try:
    import importlib_metadata
except ImportError:  # pragma: no cover
    import importlib.metadata as importlib_metadata
from loguru import logger
from typing_extensions import Literal

import edsnlp

py_version = f"{sys.version_info.major}.{sys.version_info.minor}"


def get_package(obj: Type):
    # Retrieve the __package__ attribute of the module of a type, if possible.
    # And returns the package version as well
    try:
        if isinstance(obj, ModuleType):
            module_name = obj.__name__
        else:
            module_name = obj.__module__
        if module_name == "__main__":
            raise Exception(f"Could not find package of {obj}")
        module = __import__(module_name, fromlist=["__package__"])
        package = module.__package__.split(".")[0]
        try:
            version = importlib_metadata.version(package)
        except (importlib_metadata.PackageNotFoundError, ValueError):
            return None
        return package, version
    except (ImportError, AttributeError):
        raise Exception(f"Cound not find package of type {obj}")


def save_type(pickler, obj, *args, **kwargs):
    package_name = get_package(obj)
    if package_name is not None:
        pickler.packages.add(package_name)
    dill_save_type(pickler, obj, *args, **kwargs)


def save_function(pickler, obj, *args, **kwargs):
    package_name = get_package(obj)
    if package_name is not None:
        pickler.packages.add(package_name)
    return dill_save_function(pickler, obj, *args, **kwargs)


def save_module(pickler, obj, *args, **kwargs):
    package_name = get_package(obj)
    if package_name is not None:
        pickler.packages.add(package_name)
    return dill_save_module(pickler, obj, *args, **kwargs)


class PackagingPickler(dill.Pickler):
    dispatch = dill.Pickler.dispatch.copy()

    dispatch[FunctionType] = save_function
    dispatch[type] = save_type
    dispatch[ModuleType] = save_module

    def __init__(self, *args, **kwargs):
        self.file = io.BytesIO()
        super().__init__(self.file, *args, **kwargs)
        self.packages = set()


def get_deep_dependencies(obj):
    pickler = PackagingPickler(byref=True)
    pickler.dump(obj)
    return sorted(pickler.packages)


app = Cli(pretty_exceptions_show_locals=False, pretty_exceptions_enable=False)


def snake_case(s):
    # From https://www.w3resource.com/python-exercises/string/python-data-type-string-exercise-97.php  # noqa E501
    return "_".join(
        re.sub(
            "([A-Z][a-z]+)", r" \1", re.sub("([A-Z]+)", r" \1", s.replace("-", " "))
        ).split()
    ).lower()


class ModuleName(str):
    def __new__(cls, *args, **kwargs):
        raise NotImplementedError("ModuleName is only meant for typing.")

    @classmethod
    def __get_validators__(self):
        yield self.validate

    @classmethod
    def validate(cls, value, config=None):
        if not isinstance(value, str):
            raise TypeError("string required")

        if not re.match(
            r"^([A-Z0-9]|[A-Z0-9][A-Z0-9._-]*[A-Z0-9])$", value, flags=re.IGNORECASE
        ):
            raise ValueError("invalid identifier")
        return value


if TYPE_CHECKING:
    ModuleName = str  # noqa F811

POETRY_SNIPPET = """\
from poetry.core.masonry.builders.sdist import SdistBuilder
from poetry.factory import Factory
from poetry.core.masonry.utils.module import ModuleOrPackageNotFound
import sys
# Initialize the Poetry object for the current project
poetry = Factory().create_poetry("__root_dir__")

# Initialize the builder
try:
    builder = SdistBuilder(poetry, None, None)
    # Get the list of files to include
    files = builder.find_files_to_add()
except ModuleOrPackageNotFound:
    if not poetry.package.packages:
        print([])
        sys.exit(0)

print([
    {k: v for k, v in {
    "include": getattr(include, '_include'),
    "from": getattr(include, 'source', None),
    "formats": getattr(include, 'formats', None),
    }.items() if v}
    for include in builder._module.includes
])


# Print the list of files
for file in files:
    print(file.path)
"""

INIT_PY = """
# -----------------------------------------
# This section was autogenerated by edsnlp
# -----------------------------------------

import edsnlp
from pathlib import Path
from typing import Optional, Dict, Any

__version__ = {__version__}

def load(
    overrides: Optional[Dict[str, Any]] = None,
) -> edsnlp.Pipeline:
    artifacts_path = Path(__file__).parent / "{artifacts_dir}"
    model = edsnlp.load(artifacts_path, overrides=overrides)
    return model
"""


# def parse_authors_as_dicts(authors):
#     authors = [authors] if isinstance(authors, str) else authors
#     return [
#         dict(zip(("name", "email"), re.match(r"(.*) <(.*)>", author).groups()))
#         if isinstance(author, str)
#         else author
#         for author in authors
#     ]


def parse_authors_as_strings(authors):
    authors = [authors] if isinstance(authors, str) else authors
    return [
        author if isinstance(author, str) else f"{author['name']} <{author['email']}>"
        for author in authors
    ]


class PoetryPackager:
    def __init__(
        self,
        pyproject: Optional[Dict[str, Any]],
        pipeline: Union[Path, "edsnlp.Pipeline"],
        version: Optional[str],
        name: Optional[ModuleName],
        root_dir: Path = ".",
        build_dir: Path = "build",
        dist_dir: Path = "dist",
        artifacts_name: ModuleName = "artifacts",
        dependencies: Optional[Sequence[Tuple[str, str]]] = None,
        metadata: Optional[Dict[str, Any]] = {},
    ):
        self.poetry_bin_path = (
            subprocess.run(["which", "poetry"], stdout=subprocess.PIPE)
            .stdout.decode()
            .strip()
        )
        self.version = version
        self.name = name
        self.pyproject = pyproject
        self.root_dir = root_dir.resolve()
        self.dependencies = dependencies
        self.pipeline = pipeline
        self.artifacts_name = artifacts_name
        self.dist_dir = (
            dist_dir if Path(dist_dir).is_absolute() else self.root_dir / dist_dir
        )

        with self.ensure_pyproject(metadata):
            python_executable = (
                Path(self.poetry_bin_path).read_text().split("\n")[0][2:]
            )
            result = subprocess.run(
                [
                    *python_executable.split(),
                    "-c",
                    POETRY_SNIPPET.replace("__root_dir__", str(self.root_dir)),
                ],
                stdout=subprocess.PIPE,
                cwd=self.root_dir,
            )
            if result.returncode != 0:
                raise Exception()
            out = result.stdout.decode().strip().split("\n")

        self.poetry_packages = eval(out[0])
        self.build_dir = (
            build_dir if Path(build_dir).is_absolute() else root_dir / build_dir
        ) / self.name
        self.file_paths = [self.root_dir / file_path for file_path in out[1:]]

        logger.info(f"root_dir: {self.root_dir}")
        logger.info(f"build_dir: {self.build_dir}")
        logger.info(f"artifacts_name: {self.artifacts_name}")
        logger.info(f"name: {self.name}")

    @contextmanager
    def ensure_pyproject(self, metadata):
        """Generates a Poetry based pyproject.toml"""
        metadata = dict(metadata)
        new_pyproject = self.pyproject is None
        if "authors" in metadata:
            metadata["authors"] = parse_authors_as_strings(metadata["authors"])
        try:
            if new_pyproject:
                self.pyproject = {
                    "build-system": {
                        "requires": ["poetry-core>=1.0.0"],
                        "build-backend": "poetry.core.masonry.api",
                    },
                    "tool": {
                        "poetry": {
                            **metadata,
                            "name": self.name,
                            "version": self.version or "0.1.0",
                            "dependencies": {
                                "python": f">={py_version},<4.0",
                                **{
                                    dep_name: f"^{dep_version}"
                                    for dep_name, dep_version in self.dependencies
                                },
                            },
                        },
                    },
                }
                (self.root_dir / "pyproject.toml").write_text(
                    toml.dumps(self.pyproject)
                )
            else:
                self.name = (
                    self.pyproject["tool"]["poetry"]["name"]
                    if self.name is None
                    else self.name
                )
                for key, value in metadata.items():
                    pyproject_value = self.pyproject["tool"]["poetry"].get(key)
                    if pyproject_value != metadata[key]:
                        raise ValueError(
                            f"Field {key} in pyproject.toml doesn't match the one "
                            f"passed as argument, you should remove it from the "
                            f"metadata parameter. Avoid using metadata if you already "
                            f"have a pyproject.toml file.\n"
                            f"pyproject.toml:\n {pyproject_value}\n"
                            f"metadata:\n {value}"
                        )
            yield
        except Exception:
            if new_pyproject:
                os.remove(self.root_dir / "pyproject.toml")
            raise

    def list_files_to_add(self):
        # Extract python from the shebang in the poetry executable
        return self.file_paths

    def build(
        self,
        distributions: Sequence[str] = (),
        config_settings: Optional[build.ConfigSettingsType] = None,
        isolation: bool = True,
        skip_dependency_check: bool = False,
    ):
        logger.info(f"Building package {self.name}")

        if distributions:
            build_call = build_package
        else:
            build_call = build_package_via_sdist
            distributions = ["wheel"]
        build_call(
            srcdir=self.build_dir,
            outdir=self.dist_dir,
            distributions=distributions,
            config_settings=config_settings,
            isolation=isolation,
            skip_dependency_check=skip_dependency_check,
        )

    def update_pyproject(self):
        # Replacing project name
        old_name = self.pyproject["tool"]["poetry"]["name"]
        self.pyproject["tool"]["poetry"]["name"] = self.name
        logger.info(
            f"Replaced project name {old_name!r} with {self.name!r} in poetry based "
            f"project"
        )

        if self.version is not None:
            old_version = self.pyproject["tool"]["poetry"]["version"]
            self.pyproject["tool"]["poetry"]["version"] = self.version
            logger.info(
                f"Replaced project version {old_version!r} with {self.version!r} in "
                f"poetry based project"
            )

        # Adding artifacts to include in pyproject.toml
        snake_name = snake_case(self.name.lower())
        included = self.pyproject["tool"]["poetry"].setdefault("include", [])
        included.append(f"{snake_name}/{self.artifacts_name}/**")

        packages = list(self.poetry_packages)
        packages.append({"include": snake_name})
        self.pyproject["tool"]["poetry"]["packages"] = packages

    def make_src_dir(self):
        snake_name = snake_case(self.name.lower())
        package_dir = self.build_dir / snake_name
        shutil.rmtree(package_dir, ignore_errors=True)
        os.makedirs(package_dir, exist_ok=True)
        build_artifacts_dir = package_dir / self.artifacts_name
        for file_path in self.list_files_to_add():
            new_file_path = self.build_dir / Path(file_path).relative_to(self.root_dir)
            if isinstance(self.pipeline, Path) and self.pipeline in file_path.parents:
                raise Exception(
                    f"Pipeline ({self.artifacts_name}) is already "
                    "included in the package's data, you should "
                    "remove it from the pyproject.toml metadata."
                )
            os.makedirs(new_file_path.parent, exist_ok=True)
            logger.info(f"COPY {file_path}" f" TO {new_file_path}")
            shutil.copy(file_path, new_file_path)

        self.update_pyproject()

        # Write pyproject.toml
        (self.build_dir / "pyproject.toml").write_text(toml.dumps(self.pyproject))

        if isinstance(self.pipeline, Path):
            # self.pipeline = edsnlp.load(self.pipeline)
            shutil.copytree(
                self.pipeline,
                build_artifacts_dir,
            )
        else:
            self.pipeline.to_disk(build_artifacts_dir)
        os.makedirs(package_dir, exist_ok=True)
        with open(package_dir / "__init__.py", mode="a") as f:
            f.write(
                INIT_PY.format(
                    __version__=repr(self.version),
                    artifacts_dir=os.path.relpath(build_artifacts_dir, package_dir),
                )
            )


@app.command(name="package")
def package(
    pipeline: Union[Path, "edsnlp.Pipeline"],
    name: Optional[ModuleName] = None,
    root_dir: Path = ".",
    build_dir: Path = "build",
    dist_dir: Path = "dist",
    artifacts_name: ModuleName = "artifacts",
    check_dependencies: bool = False,
    project_type: Optional[Literal["poetry", "setuptools"]] = None,
    version: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = {},
    distributions: Optional[Sequence[Literal["wheel", "sdist"]]] = ["wheel"],
    config_settings: Optional[Mapping[str, Union[str, Sequence[str]]]] = None,
    isolation: bool = True,
    skip_build_dependency_check: bool = False,
):
    # root_dir = Path(".").resolve()
    pyproject_path = root_dir / "pyproject.toml"

    if not pyproject_path.exists():
        check_dependencies = True
        if name is None:
            raise ValueError(
                f"No pyproject.toml could be found in the root directory {root_dir}, "
                f"you need to create one, or fill the name parameter."
            )

    dependencies = None
    if check_dependencies:
        if isinstance(pipeline, Path):
            pipeline = edsnlp.load(pipeline)
        dependencies = get_deep_dependencies(pipeline)
        for dep in dependencies:
            print("DEPENDENCY", dep[0].ljust(30), dep[1])

    root_dir = root_dir.resolve()

    pyproject = None
    if pyproject_path.exists():
        pyproject = toml.loads((root_dir / "pyproject.toml").read_text())

        if "tool" in pyproject and "poetry" in pyproject["tool"]:
            project_type = "poetry"

    if project_type == "poetry":
        packager = PoetryPackager(
            pyproject=pyproject,
            pipeline=pipeline,
            name=name,
            version=version,
            root_dir=root_dir,
            build_dir=build_dir,
            dist_dir=dist_dir,
            artifacts_name=artifacts_name,
            dependencies=dependencies,
            metadata=metadata,
        )
    else:
        raise Exception(
            "Could not infer project type, only poetry based projects are "
            "supported for now"
        )

    packager.make_src_dir()
    packager.build(
        distributions=distributions,
        config_settings=config_settings,
        isolation=isolation,
        skip_dependency_check=skip_build_dependency_check,
    )
