"""Apply Black to Jupyter notebooks."""

# Copyright © 2019 Tom Catling <tomcatling@gmail.com>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
# REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
# INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
# OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.

import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple

import black
import click
import nbformat
from attr import dataclass

DEFAULT_LINE_LENGTH = 79
DEFAULT_INCLUDES = r"\.ipynb$"
DEFAULT_EXCLUDES = (
    r"/(\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|_build|buck-out|build|dist|"
    r"\.ipynb_checkpoints)/"
)


@click.command(context_settings=dict(help_option_names=["-h", "--help"]))
@click.option(
    "-l",
    "--line-length",
    type=int,
    default=DEFAULT_LINE_LENGTH,
    help="How many characters per line to allow.",
    show_default=True,
)
@click.option(
    "--check",
    is_flag=True,
    help=(
        "Don't write the files back, just return the status.  Return code 0 "
        "means nothing would change.  Return code 1 means some files would be "
        "reformatted.  Return code 123 means there was an internal error."
    ),
)
@click.option(
    "--include",
    type=str,
    default=DEFAULT_INCLUDES,
    help=(
        "A regular expression that matches files and directories that should "
        "be included on recursive searches.  An empty value means all files "
        "are included regardless of the name.  Use forward slashes for "
        "directories on all platforms (Windows, too).  Exclusions are "
        "calculated first, inclusions later."
    ),
    show_default=True,
)
@click.option(
    "--exclude",
    type=str,
    default=DEFAULT_EXCLUDES,
    help=(
        "A regular expression that matches files and directories that should "
        "be excluded on recursive searches. An empty value means no paths are "
        "excluded. Use forward slashes for directories on all platforms "
        "(Windows, too). Exclusions are calculated first, inclusions later."
    ),
    show_default=True,
)
@click.option(
    "--clear-output",
    is_flag=True,
    help="Clear cell output as part of formatting.",
)
@click.argument(
    "src",
    nargs=-1,
    type=click.Path(
        exists=True,
        file_okay=True,
        dir_okay=True,
        readable=True,
        allow_dash=True,
    ),
    is_eager=True,
)
@click.option(
    "--config",
    type=click.Path(
        exists=False,
        file_okay=True,
        dir_okay=False,
        readable=True,
        allow_dash=False,
    ),
    is_eager=True,
    callback=black.read_pyproject_toml,
    help="Read configuration from PATH.",
)
@click.pass_context
def cli(
    ctx: click.Context,
    line_length: int,
    check: bool,
    include: str,
    exclude: str,
    clear_output: bool,
    src: Tuple[str],
    config: Optional[str],
) -> None:
    """
    The uncompromising code formatter, for Jupyter notebooks.
    """
    write_back = black.WriteBack.from_configuration(check=check, diff=False)
    mode = black.FileMode.from_configuration(
        py36=True,
        pyi=False,
        skip_string_normalization=False,
        skip_numeric_underscore_normalization=False,
    )

    if config:
        black.out(f"Using configuration from {config}.", bold=False, fg="blue")

    try:
        include_regex = black.re_compile_maybe_verbose(include)
    except re.error:
        black.err(f"Invalid regular expression for include given: {include!r}")
        ctx.exit(2)
    try:
        exclude_regex = black.re_compile_maybe_verbose(exclude)
    except re.error:
        black.err(f"Invalid regular expression for exclude given: {exclude!r}")
        ctx.exit(2)

    report = black.Report(check=check, quiet=False, verbose=True)
    root = black.find_project_root(src)
    sources: Set[Path] = set()
    for s in src:
        p = Path(s)
        if p.is_dir():
            sources.update(
                black.gen_python_files_in_dir(
                    p, root, include_regex, exclude_regex, report
                )
            )
        elif p.is_file() or s == "-":
            # if a file was explicitly given, we don't care about its extension
            sources.add(p)
        else:
            black.err(f"invalid path: {s}")
    if len(sources) == 0:
        black.out("No paths given. Nothing to do.")
        ctx.exit(0)

    for source in sources:
        reformat_one(
            src=source,
            line_length=line_length,
            write_back=write_back,
            mode=mode,
            clear_output=clear_output,
            report=report,
        )

    black.out(f"All done!")
    click.secho(str(report), err=True)
    ctx.exit(report.return_code)


def reformat_one(
    src: Path,
    line_length: int,
    write_back: black.WriteBack,
    mode: black.FileMode,
    clear_output: bool,
    report: black.Report,
) -> None:
    """Reformat a single file under `src`."""
    try:

        sub_report = SubReport(write_back=write_back)
        changed = black.Changed.NO

        cache: black.Cache = {}
        if write_back is not black.WriteBack.DIFF:
            cache = black.read_cache(line_length, mode)
            res_src = src.resolve()
            if res_src in cache and cache[res_src] == black.get_cache_info(
                res_src
            ):
                changed = black.Changed.CACHED
        if changed is not black.Changed.CACHED:
            sub_report = format_file_in_place(
                src,
                line_length=line_length,
                write_back=write_back,
                mode=mode,
                clear_output=clear_output,
                sub_report=sub_report,
            )
            if sub_report.change_count:
                changed = black.Changed.YES
        if (
            write_back is black.WriteBack.YES
            and changed is not black.Changed.CACHED
        ) or (
            write_back is black.WriteBack.CHECK and changed is black.Changed.NO
        ):
            black.write_cache(cache, [src], line_length, mode)
        report.done(src, changed)
        if changed is not black.Changed.CACHED:
            click.secho(f"    {sub_report}", err=True)
    except Exception as exc:
        report.failed(src, str(exc))


def format_file_in_place(
    src: Path,
    line_length: int,
    write_back: black.WriteBack,
    mode: black.FileMode,
    clear_output: bool,
    sub_report: "SubReport",
) -> "SubReport":
    """
    Format file under `src` path. Return True if changed.
    If `write_back` is YES, write reformatted code to the file.
    """

    with src.open() as fp:
        src_contents = nbformat.read(fp, as_version=nbformat.NO_CONVERT)

    dst_cells: List[Dict[Any, Any]] = []
    for cell in src_contents["cells"]:
        if cell["cell_type"] == "code":
            try:
                cell["source"] = format_cell_source(
                    cell["source"], line_length=line_length, mode=mode
                )
                sub_report.done(black.Changed.YES)
            except black.NothingChanged:
                sub_report.done(black.Changed.NO)
            except black.InvalidInput:
                sub_report.failed()
            if clear_output:
                try:
                    cell["outputs"], cell[
                        "execution_count"
                    ] = clear_cell_outputs(
                        cell["outputs"], cell["execution_count"]
                    )
                    sub_report.done_output(black.Changed.YES)
                except black.NothingChanged:
                    sub_report.done_output(black.Changed.NO)
        dst_cells.append(cell)
    src_contents["cells"] = dst_cells

    # click.echo(src)
    # click.secho(f"    {sub_report}", err=True)

    if write_back is black.WriteBack.YES:
        with src.open("w") as fp:
            nbformat.write(src_contents, fp)

    return sub_report


def clear_cell_outputs(
    src_outputs: List[str], src_execution_count: int
) -> Tuple[List[str], None]:
    if src_outputs == [] and src_execution_count is None:
        raise black.NothingChanged
    return [], None


def format_cell_source(
    src_contents: str, *, line_length: int, mode: black.FileMode
) -> black.FileContent:
    """Reformat contents of cell and return new contents.
    Additionally confirm that the reformatted code is valid by calling
    :func:`assert_equivalent` and :func:`assert_stable` on it.
    `line_length` is passed to :func:`format_str`.
    """

    if src_contents.strip() == "":
        raise black.NothingChanged

    dst_contents = format_str(src_contents, line_length=line_length, mode=mode)

    if src_contents == dst_contents:
        raise black.NothingChanged

    assert_equivalent(src_contents, dst_contents)
    assert_stable(dst_contents, line_length=line_length, mode=mode)

    return dst_contents


def format_str(
    src_contents: str,
    line_length: int,
    *,
    mode: black.FileMode = black.FileMode.AUTO_DETECT,
) -> black.FileContent:
    trailing_semi_colon = src_contents.rstrip()[-1] == ";"
    src_contents = hide_magic(src_contents)
    dst_contents = black.format_str(
        src_contents, line_length=line_length, mode=mode
    )
    dst_contents = dst_contents.rstrip()
    if trailing_semi_colon:
        dst_contents = f"{dst_contents};"
    dst_contents = reveal_magic(dst_contents)
    return dst_contents


def assert_equivalent(src: str, dst: str) -> None:
    black.assert_equivalent(hide_magic(src), hide_magic(dst))


def assert_stable(
    dst: str,
    line_length: int,
    mode: black.FileMode = black.FileMode.AUTO_DETECT,
) -> None:
    new_dst = format_str(dst, line_length=line_length, mode=mode)
    if dst != new_dst:
        raise AssertionError(
            "INTERNAL ERROR: Black produced different code on the second pass "
            "of the formatter."
        ) from None


def contains_magic(line: str) -> bool:
    if len(line) == 0:
        return False
    else:
        return line[0] == "%" or line[0] == "!"


def hide_magic(source: str) -> str:
    """
    Black can't deal with cell or line magic, so we
    disguise it as a comment. This keeps it in the same
    place in the reformatted code.
    """

    def _hide_magic_line(line: str) -> str:
        return f"###MAGIC###{line}" if contains_magic(line) else line

    return "\n".join(_hide_magic_line(l) for l in source.split("\n"))


def reveal_magic(source: str) -> str:
    """
    Reveal any notebook magic hidden by hide_magic().
    """
    return source.replace("###MAGIC###", "")


@dataclass
class SubReport:
    """
    Provides a reformatting counter for notebook cells.
    Can be rendered with `str(report)`.
    """

    write_back: black.WriteBack
    change_count: int = 0
    same_count: int = 0
    failure_count: int = 0
    output_change_count: int = 0
    output_same_count: int = 0

    def done(self, changed: black.Changed) -> None:
        """
        Increment the counter for successful reformatting.
        """
        if changed is black.Changed.YES:
            self.change_count += 1
        else:
            self.same_count += 1

    def done_output(self, changed: black.Changed) -> None:
        """
        Increment the counter for successful clear output.
        """
        if changed is black.Changed.YES:
            self.output_change_count += 1
        else:
            self.output_same_count += 1

    def failed(self) -> None:
        """
        Increment the counter for failed reformatting.
        """
        self.failure_count += 1

    def __str__(self) -> str:
        """
        Render a report of the current state.
        """
        if self.write_back is black.WriteBack.CHECK:
            reformatted = "would be reformatted"
            unchanged = "would be left unchanged"
            failed = "would fail to reformat"
            cleared = "would be cleared"
        else:
            reformatted = "reformatted"
            unchanged = "left unchanged"
            failed = "failed to reformat"
            cleared = "cleared"
        report = []
        if self.change_count:
            s = "s" if self.change_count > 1 else ""
            report.append(
                click.style(
                    f"{self.change_count} cell{s} {reformatted}", bold=True
                )
            )
        if self.same_count:
            s = "s" if self.same_count > 1 else ""
            report.append(f"{self.same_count} cell{s} {unchanged}")
        if self.failure_count:
            s = "s" if self.failure_count > 1 else ""
            report.append(
                click.style(f"{self.failure_count} cell{s} {failed}", fg="red")
            )
        if self.output_change_count:
            s = "s" if self.change_count > 1 else ""
            report.append(
                click.style(
                    f"{self.change_count} output{s} {cleared}", bold=True
                )
            )
        if self.output_same_count:
            s = "s" if self.same_count > 1 else ""
            report.append(f"{self.same_count} output{s} {unchanged}")
        return ", ".join(report) + "."


if __name__ == "__main__":
    cli()
