import pyarrow as pa
from pyarrow import parquet as pq
import nbconvert as nbc
import nbformat as nbf
import subprocess

from dsi.backends.filesystem import Filesystem


class Parquet(Filesystem):
    """
    Support for a Parquet back-end.

    Parquet is a convenient format when metadata are larger than SQLite supports.
    """

    def __init__(self, filename, **kwargs):
        super().__init__(filename=filename)
        self.filename = filename
        try:
            self.compression = kwargs['compression']
        except KeyError:
            self.compression = None

    def get_artifacts(self):
        """Get Parquet data from filename."""
        table = pq.read_table(self.filename)
        resout = table.to_pydict()
        return resout

    def put_artifacts(self, collection):
        """Put artifacts into file at filename path."""
        table = pa.table(collection)
        pq.write_table(table, self.filename, compression=self.compression)

    @staticmethod
    def get_cmd_output(cmd: list) -> str:
        """
        Runs a given command and returns the stdout if successful.

        If stderr is not empty, an exception is raised with the stderr text.
        """
        proc = subprocess.run(cmd, capture_output=True, shell=True)
        if proc.stderr != b"":
            raise Exception(proc.stderr)
        return proc.stdout.strip().decode("utf-8")

    def inspect_artifacts(self, collection, interactive=False):
        """Populate a Jupyter notebook with tools required to look at Parquet data."""
        nb = nbf.v4.new_notebook()
        text = """\
        # This notebook was auto-generated by a DSI Backend for Parquet.
        # Execute the Jupyter notebook cells below and interact with "df"
        # to explore your data.
        """
        code1 = """\
        import pandas as pd
        df = pd.read_parquet('{}')
        df.head()
        """.format(self.filename)

        code2 = """\
        df.info()
        """

        code3 = """\
        df.describe()
        """

        nb['cells'] = [nbf.v4.new_markdown_cell(text),
                       nbf.v4.new_code_cell(code1),
                       nbf.v4.new_code_cell(code2),
                       nbf.v4.new_code_cell(code3)]

        fname = 'dsi_parquet_backend_output.ipynb'

        print('Writing Jupyter notebook...')
        with open(fname, 'w') as fh:
            nbf.write(nb, fh)

        # open the jupyter notebook for static page generation
        with open(fname, 'r', encoding='utf-8') as fh:
            nb_content = nbf.read(fh, as_version=4)
        # Init executor for notebook
        run_nb = nbc.preprocessors.ExecutePreprocessor(timeout=-1) # No timeout
        # Execute the notebook
        run_nb.preprocess(nb_content, {'metadata':{'path':'.'}})

        if interactive:
            print('Opening Jupyter notebook...')
            self.get_cmd_output(cmd=['jupyter-lab ./dsi_parquet_backend_output.ipynb'])
        else:
#            self.get_cmd_output(cmd=['jupyter nbconvert --to html {}'.format(fname)])
            # Init HTML exporter
            html_exporter = nbc.HTMLExporter()
            html_content,_ = html_exporter.from_notebook_node(nb_content)
            # Save HTML file
            html_filename = 'dsi_parquet_backend_output.html'
            with open(html_filename, 'w', encoding='utf-8') as fh:
                fh.write(html_content)
