#!/usr/bin/env python
#
# Copyright (c) Bo Peng and the University of Texas MD Anderson Cancer Center
# Distributed under the terms of the 3-clause BSD License.

import argparse
import base64
import io

import pkg_resources

from sos.utils import dehtml, dot_to_gif, env, linecount_of_file


def get_previewers():
    # Note: data is zest.releaser specific: we want to pass
    # something to the plugin
    group = "sos_previewers"
    result = []
    for entrypoint in pkg_resources.iter_entry_points(group=group):
        # if ':' in entry point name, it should be a function
        try:
            name, priority = entrypoint.name.split(",", 1)
            priority = int(priority)
        except Exception as e:
            env.logger.warning(
                f"Ignore incorrect previewer entry point {entrypoint}: {e}")
            continue
        # If name points to a function in a module. Let us try to import the module
        if ":" in name:
            import importlib

            try:
                mod, func = name.split(":")
                imported = importlib.import_module(mod)
                result.append((getattr(imported, func), entrypoint, priority))
            except ImportError:
                env.logger.warning(f"Failed to load function {mod}:{func}")
        else:
            result.append((name, entrypoint, priority))
    #
    result.sort(key=lambda x: -x[2])
    # we put string infront of functions so that filename matching will be used before
    # content matching. For example, the xlsx file is actually a zip file so it could be
    # previewed as a zip file first.
    return ([x for x in result if isinstance(x[0], str) and x[0] != "*"] +
            [x for x in result if not isinstance(x[0], str)] +
            [x for x in result if x[0] == "*"])


def _preview_img_parser():
    parser = argparse.ArgumentParser(prog="%preview *.pdf")
    parser.add_argument(
        "--width",
        help="Width of the previewed image, can be in any HTML units such as px, em.",
    )
    parser.add_argument(
        "--height",
        help="Height of the previewed image, can be in any HTML units such as px, em.",
    )
    parser.error = lambda msg: env.logger.warning(msg)
    return parser


def preview_img(filename, kernel=None, style=None):
    with open(filename, "rb") as f:
        image = f.read()
    import imghdr

    image_type = imghdr.what(None, image)
    image_data = base64.b64encode(image).decode("ascii")

    args = None
    meta = {}
    if style is not None and "options" in style:
        parser = _preview_img_parser()
        try:
            args = parser.parse_args(style["options"])
            meta.update({
                "image/png":
                    dict(([["width", args.width]] if args.width else []) +
                         ([["height", args.height]] if args.height else []))
            })
        except SystemExit:
            return

    if image_type != "png":
        try:
            if image_type == "gif":
                return {"image/png": image_data}, meta
            else:
                from wand.image import Image

                img = Image(filename=filename)
                return {
                    "image/" + image_type:
                        image_data,
                    "image/png":
                        base64.b64encode(img._repr_png_()).decode("ascii"),
                }, meta
        except Exception:
            return {"image/" + image_type: image_data}, meta
    else:
        return {"image/" + image_type: image_data}, meta


def preview_svg(filename, kernel=None, style=None):
    with open(filename, "r") as f:
        image_data = f.read()
    return {"image/svg+xml": image_data}


def _preview_pdf_parser():
    parser = argparse.ArgumentParser(prog="%preview *.pdf")
    parser.add_argument(
        "--pages", nargs="+", type=int, help="Pages of the PDF to preview.")
    parser.add_argument(
        "--width",
        help="Width of the previewed image, can be in any HTML units such as px, em.",
    )
    parser.add_argument(
        "--height",
        help="Height of the previewed image, can be in any HTML units such as px, em.",
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=150,
        help="resolution of the converted png preview")
    parser.error = lambda msg: env.logger.warning(msg)
    return parser


def preview_pdf(filename, kernel=None, style=None):
    use_png = False
    warn = kernel.warn if kernel is not None else env.logger.warning
    if style is not None and "style" in style:
        if style["style"] != "png":
            if style["style"] is not None:
                warn(
                    f'Option --style of PDF preview only accept parameter png: {style["style"]} provided'
                )
        else:
            use_png = True
    args = None
    if style is not None and "options" in style:
        parser = _preview_pdf_parser()
        try:
            args = parser.parse_args(style["options"])
        except SystemExit:
            return
    meta = {}
    embed_options = ""
    if args and (args.width or args.height):
        meta.update({
            "image/png":
                dict(([["width", args.width]] if args.width else []) +
                     ([["height", args.height]] if args.height else []))
        })
        embed_options += (f'width="{args.width}" ' if args.width else " ") + (
            f'height="{args.height}" ' if args.height else " ")
    if use_png:
        try:
            # this import will fail even if wand is installed
            # if imagemagick is not installed properly.
            from wand.image import Image

            img = Image(filename=filename, resolution=args.dpi)

            if img.width == 0 or img.height == 0:
                raise ValueError("Image appears to have zero width or height")
            nPages = len(img.sequence)
            pages = list(range(nPages))

            if args and args.pages is not None:
                pages = [x - 1 for x in args.pages]
                for p in pages:
                    if p >= nPages:
                        warn(
                            f"Page {p} out of range of the pdf file ({nPages} pages)"
                        )
                        pages = list(range(nPages))
                        break
            # single page PDF
            if len(pages) == 1 and pages[0] == 0:
                return {
                    "image/png":
                        base64.b64encode(img._repr_png_()).decode("ascii")
                }, meta
            elif len(pages) == 1:
                # if only one page
                return {
                    "image/png":
                        base64.b64encode(
                            Image(img.sequence[pages[0]])._repr_png_()
                        ).decode("ascii")
                }, meta
            else:
                widths = [img.sequence[p].width for p in pages]
                heights = [img.sequence[p].height for p in pages]
                image = Image(width=max(widths), height=sum(heights))
                for i, p in enumerate(pages):
                    image.composite(
                        img.sequence[p], top=sum(heights[:i]), left=0)
                image.format = "png"
                with io.BytesIO() as out:
                    image.save(file=out)
                    img_data = out.getvalue()
                return {
                    "image/png": base64.b64encode(img_data).decode("ascii")
                }, meta
        except Exception as e:
            warn(e)
            return {
                "text/html":
                    f'<embed src="{filename}" {embed_options} type="application/pdf" />'
            }
    else:
        # by default use iframe, because PDF figure can have multiple pages (#693)
        return {
            "text/html":
                f'<embed src="{filename}" {embed_options} type="application/pdf" />'
        }


def preview_html(filename, kernel=None, style=None):
    with open(filename) as html:
        content = html.read()
    return {"text/html": content, "text/plain": dehtml(content)}


def _get_txt_parser():
    parser = argparse.ArgumentParser(prog="%preview *.txt")
    parser.add_argument(
        "-l",
        "--limit",
        type=int,
        default=5,
        help="""Maximum number
        of lines to preview.""",
    )
    parser.error = lambda msg: env.logger.warning(msg)
    return parser


def preview_txt(filename, kernel=None, style=None):
    if style is not None and "options" in style:
        parser = _get_txt_parser()
        try:
            args = parser.parse_args(style["options"])
        except SystemExit:
            return
        limit = args.limit
    else:
        limit = 5

    nlines = linecount_of_file(filename)
    content = "HINT: {} line{}{}\n".format(
        nlines,
        "s" if nlines > 1 else "",
        f" ({limit} displayed, see --limit)" if nlines > limit else "",
    )

    with open(filename, "r") as fin:
        if limit < 0:
            content += fin.read()
        else:
            for _ in range(limit):
                content += fin.readline()
    return content


def preview_csv(filename, kernel=None, style=None):
    import pandas

    from .visualize import Visualizer

    data = pandas.read_csv(filename)
    return Visualizer(kernel, style).preview(data)


def preview_xls(filename, kernel=None, style=None):
    import pandas

    from .visualize import Visualizer

    data = pandas.read_excel(filename)
    return Visualizer(kernel, style).preview(data)


def preview_zip(filename, kernel=None, style=None):
    import zipfile

    names = zipfile.ZipFile(filename).namelist()
    return (f"{len(names)} files\n" + "\n".join(names[:5]) +
            ("\n..." if len(names) > 5 else ""))


def preview_tar(filename, kernel=None, style=None):
    import tarfile

    with tarfile.open(filename, "r:*") as tar:
        # only extract files
        names = [x.name for x in tar.getmembers() if x.isfile()]
    return (f"{len(names)} files\n" + "\n".join(names[:5]) +
            ("\n..." if len(names) > 5 else ""))


def preview_gz(filename, kernel=None, style=None):
    import tarfile

    # .tar.gz
    if tarfile.is_tarfile(filename):
        return preview_tar(filename, kernel, style)
    import gzip

    content = b""
    with gzip.open(filename, "rb") as fin:
        for _ in range(5):
            content += fin.readline()
    try:
        return content.decode()[:2000]
    except Exception:
        return "binary data"


def preview_md(filename, kernel=None, style=None):
    import markdown

    with open(filename) as fin:
        text = fin.read()
    html = markdown.markdown(text)
    return {"text/html": html, "text/plain": text}


def preview_dot(filename, kernel=None, style=None):
    data = dot_to_gif(
        filename, warn=kernel.warn if kernel else env.logger.warning)
    # according to https://github.com/ipython/ipython/issues/10045
    # I have to use 'image/png' instead of 'image/gif' to get the gif displayed.
    return {"image/png": data}
