import asyncio
from enum import Enum, auto
from pathlib import Path
from typing import Dict, Generator, List, Optional, Union

from markupsafe import Markup

from .d3 import Script
from .plot import Plot
from .style import CSS, GRID
from .utils import Data, DataInput, arrange

JSCode = Union[Plot, Script]
JSInput = Union[Dict[str, JSCode], List[JSCode]]

class PlotType(Enum):
    SINGLE_PLOT = auto()
    MULTIPLE_PLOTS = auto()
    SINGLE_D3 = auto()
    MULTIPLE_D3 = auto()
    UNIDENTIFIED = auto()

def jupyter_environment():
    """
    Check if the code is executed in a jupyter environment
    """
    try:
        shell = get_ipython().__class__.__name__
        return shell == 'ZMQInteractiveShell'
    except NameError:
        return False

def identify(plot: JSInput) -> PlotType:
    """
    Identify the type of the plot argument

    Parameters
    ----------
    plot : JSInput
        plot javascript code

    Returns
    -------
    PlotType
        type of the plot
    """
    if isinstance(plot, Plot):
        return PlotType.SINGLE_PLOT
    elif isinstance(plot, Script):
        return PlotType.SINGLE_D3
    elif isinstance(plot, dict):
        objs = plot.values()
        if all(map(lambda obj: isinstance(obj, Plot), objs)):
            return PlotType.MULTIPLE_PLOTS
        elif all(map(lambda obj: isinstance(obj, Script), objs)):
            return PlotType.MULTIPLE_D3
    elif isinstance(plot, list):
        if all(map(lambda obj: isinstance(obj, Plot), plot)):
            return PlotType.MULTIPLE_PLOTS
        elif all(map(lambda obj: isinstance(obj, Script), plot)):
            return PlotType.MULTIPLE_D3
    return PlotType.UNIDENTIFIED

async def html(data: dict, plot: JSInput, style:Union[Path, str, dict]=None, svg:bool=False, grid:int=1, event:Optional[str] = None) -> str:
    """
    Return HTML content filled by arguments from detroit templates

    Parameters
    ----------
    data : Union[dict, Generator[dict, None, None]]
        data used for plots
    plot : JSInput
        plot javascript code
    style : Union[Path, str, dict]
        a file or a dictionary defining a CSS file
    svg : bool
        True to get javascript functions to generate a svg object in javascript
    grid : int
        number of columns
    event: Optional[str]
        for dynamic update, this javascript code is inserted into the websocket updates

    Returns
    -------
    str
        HTML content filled by arguments
    """
    from jinja2 import (ChoiceLoader, Environment, PackageLoader,
                        select_autoescape)
    loader = ChoiceLoader([PackageLoader("detroit", "templates"), PackageLoader("detroit", "static")])
    env = Environment(loader=loader, autoescape=select_autoescape(), enable_async=True)
    plot_type = identify(plot)
    style = CSS(style)

    if event: # dynamic updates
        if plot_type != PlotType.SINGLE_D3:
            raise TypeError("Unsupported plot for dynamic updates")
        code = Markup(plot)
        data = Markup(f"const data = {data};" if data else "const data = {};")
        template = env.get_template("websocket.html")
        return await template.render_async(
            code = code,
            data = data,
            style = str(style),
            event = Markup(event),
        )

    if plot_type == PlotType.UNIDENTIFIED:
        raise TypeError("Unsupported type of argument \"plot\"")
    elif plot_type == PlotType.MULTIPLE_D3 or plot_type == PlotType.MULTIPLE_PLOTS:
        single = False
        style.update(GRID(grid))
        if isinstance(plot, dict):
            code = {
                id: {"title": title, "code": Markup(code)}
                for id, (title, code) in enumerate(plot.items())
            }
        else:
            code = {
                id: {"title": "", "code": Markup(code)}
                for id, code in enumerate(plot)
            }
        id = f"plot-{len(plot) - 1}"
        width_code = "boundingRect.width" if grid > 1 else "svg.getBoundingClientRect().width"
        serialize_code = f"mysvg = serialize(makeSVGfromGrid(div, svg, {grid}));" if svg else ""
    else:
        single = True
        code = Markup(plot)
        id = "myplot"
        width_code = "svg.getBoundingClientRect().width"
        serialize_code=(
            "mysvg = serialize(svg);"
            if plot_type == PlotType.SINGLE_D3
            else "mysvg = serialize(makeSVGfromSimple(div, svg));"
        )

    data = Markup(f"const data = {data};")
    if plot_type == PlotType.SINGLE_D3 or plot_type == PlotType.MULTIPLE_D3:
        template = env.get_template("d3.html")
    else:
        template = env.get_template("plot.html")
    return await template.render_async(
        single = single,
        code=code,
        data=data,
        style=str(style),
        serialize=svg,
        serialize_code=serialize_code,
        width_code=width_code,
        id=id,
    )

async def _save(data: dict, plot: JSInput, output: Union[Path, str], style: Union[str, dict], grid: int, scale_factor: float):
    """
    Save the plot into a `output` file given arguments.
    It supports :

    * :code:`.svg` files
    * :code:`.png` files
    * :code:`.pdf` files

    Parameters
    ----------
    data : DataInput
        data used for plots
    plot : JSInput
        plot javascript code
    output : Union[Path, str]
        output file name
    style : Union[Path, str, dict]
        a file or a dictionary defining a CSS file
    grid : int
        number of columns
    scale_factor : float
        only for :code:`.png` file; the more the number is higher, the more the quality of image will be
    """
    from playwright.async_api import async_playwright
    if isinstance(output, str):
        output = Path(output)
    input = Path("~detroit-tmp.html")
    input.write_text(await html(data, plot, style=style, grid=grid, svg=output.suffix == ".svg"))
    async with async_playwright() as p:
        browser = await p.chromium.launch()
        if output.suffix == ".png":
            context = await browser.new_context(device_scale_factor=scale_factor)
            page = await context.new_page()
            await page.goto(f'file://{input.absolute()}')
            width = await page.evaluate_handle("width");
            height = await page.evaluate_handle("height");
            await page.set_viewport_size({'width': int(float(str(width))), 'height': int(float(str(height)))})
            await page.goto(f'file://{input.absolute()}')
            await page.screenshot(path=output, type='png')
        elif output.suffix == ".pdf":
            page = await browser.new_page()
            await page.goto(f'file://{input.absolute()}')
            await page.pdf(path=output)
        elif output.suffix == ".svg":
            page = await browser.new_page()
            await page.set_viewport_size({'width': 2560, 'height': 1440})
            await page.goto(f'file://{input.absolute()}')
            jshandle = await page.evaluate_handle("mysvg")
            output.write_text(str(jshandle))
        else:
            await browser.close()
            input.unlink()
            raise ValueError(f"Unsupported \"{output.suffix}\" file")
        await browser.close()
        input.unlink()

def save(data: DataInput, plot: JSInput, output:Union[Path, str], style:Union[Path, str, dict]=None, grid:int=1, scale_factor:float=1) -> str:
    """
    Save the plot into a `output` file given arguments
    It supports :

    * :code:`.svg` files
    * :code:`.png` files
    * :code:`.pdf` files

    Parameters
    ----------
    data : DataInput
        data used for plots
    plot : JSInput
        plot javascript code
    output : Union[Path, str]
        output file name
    style : Union[Path, str, dict]
        a file or a dictionary defining a CSS file
    grid : int
        number of columns
    scale_factor : float
        only for :code:`.png` file; the more the number is higher, the more the quality of image will be

    Returns
    -------
    str
        Message which confirms the save of the file

    Examples
    --------

    >>> df = pl.read_csv("barley.csv")
    >>> data = Data(arrange(df))
    >>> color = {"legend": js("true"), "label": "Elevation (m)"}
    >>> contour = Plot.contour(
    ...     data.values,
    ...     {
    ...         "width": data.width,
    ...         "height": data.height,
    ...         "fill": js("Plot.identity"),
    ...         "stroke": "black"
    ...     },
    ... )
    >>> plot = Plot.plot({"color": , "marks": [contour]})
    >>> save(df, plot, "figure.png", scale_factor=2)
    "figure.png" saved.
    """
    asyncio.run(_save(arrange(data), plot, output, style, grid, scale_factor))
    return f"{output} saved."

def websocket_save(
    generator: Generator[dict, None, None],
    script: Script,
    event: str,
    style:Union[Path, str, dict]=None,
    init_data: Optional[DataInput] = None,
    record_video_dir:str="videos/",
    time:int = 2000,
    width:int = 660,
    height:int = 420,
    verbose=False,
) -> str:
    """
    Record the rendering into a headless browser and
    save the video as :code:`.webm` in :code:`record_video_dir`.

    Parameters
    ----------
    generator : Generator[dict, None, None]
        generator function which return new updates of data
    script : Script
        plot javascript code (only d3 supported)
    event : str
        javascript code called during updates
    style : Union[Path, str, dict]
        a file or a dictionary defining a CSS file
    init_data : Optional[DataInput]
        initial data before updates
    record_video_dir : str
        directory where videos are saved
    time : int
        timeout to wait
    width : int
        width of the video
    height : int
        height of the video
    verbose :
        :code:`True` to enable logging

    Returns
    -------
    str
        Message which confirms the save of the file

    Examples
    --------

    Check :ref:`the full example<Update Guide>` for more information

    >>> websocket_save(
    ...     generator,
    ...     script,
    ...     event=js("update(received_data.values, received_data.xmax);"),
    ...     init_data=data,
    ... )
    """
    from quart import Quart, websocket
    from hypercorn.asyncio import serve
    from hypercorn import Config
    from playwright.async_api import async_playwright
    import logging

    if not verbose:
        logging.disable(logging.CRITICAL)

    app = Quart("detroit")

    @app.websocket("/ws")
    async def ws():
        for data in generator():
            await websocket.send_json(data)

    @app.route("/")
    async def main():
        return await html(data=arrange(init_data), plot=script, style=style, event=event)

    shutdown_event = asyncio.Event()
    async def local_save(url: str):
        async with async_playwright() as p:
            browser = await p.chromium.launch()
            context = await browser.new_context(
                record_video_dir=record_video_dir,
                viewport={"width": width, "height": height}
            )
            page = await context.new_page()
            await page.goto(url)
            await page.wait_for_timeout(time)
            await page.close()
            await context.close()
            await browser.close()
        shutdown_event.set()

    config = Config()
    config.bind = ["localhost:5000"]

    loop = asyncio.get_event_loop()
    loop.create_task(local_save("http://localhost:5000"))
    loop.run_until_complete(
        serve(app, config, shutdown_trigger=shutdown_event.wait)
    )
    return f"Video saved."

def render(data: DataInput, plot: JSInput, style:Union[Path, str, dict]=None, grid:int=1):
    """
    Launch a web application to render plot. In a jupyter environment, display directly the plot.

    Parameters
    ----------
    data : DataInput
        data used for plots
    plot : JSInput
        plot javascript code
    style : Union[Path, str, dict]
        a file or a dictionary defining a CSS file
    grid : int
        number of columns

    Examples
    --------

    >>> df = pl.read_csv("barley.csv")
    >>> data = Data(arrange(df))
    >>> color = {"legend": js("true"), "label": "Elevation (m)"}
    >>> contour = Plot.contour(
    ...     data.values,
    ...     {
    ...         "width": data.width,
    ...         "height": data.height,
    ...         "fill": js("Plot.identity"),
    ...         "stroke": "black"
    ...     },
    ... )
    >>> plot = Plot.plot({"color": , "marks": [contour]})
    >>> render(df, plot)
    """
    from quart import Quart
    try:
        import nest_asyncio
        nest_asyncio.apply()
        from IPython import get_ipython
        from IPython.display import HTML, display
        JUPYTER_INSTALLED = True
    except:
        JUPYTER_INSTALLED = False
    data = arrange(data)
    if JUPYTER_INSTALLED and jupyter_environment():
        display(
            HTML(asyncio.run(html(data, plot, style=style, grid=grid))),
            metadata={"isolated": True}
        )
    else:
        app = Quart("detroit")

        @app.route("/data")
        def get_data():
            return data

        @app.route("/")
        async def main():
            return await html(data, plot, style=style, grid=grid)

        app.run()

def websocket_render(generator: Generator[dict, None, None], script: Script, event: str, style:Union[Path, str, dict]=None, init_data: Optional[DataInput] = None):
    """
    Launch a web application to render plot with dynamic updates of data
    Jupyter environment not supported

    Parameters
    ----------
    generator : Generator[dict, None, None]
        Generator function which return new updates of data
    script : Script
        Plot javascript code (only d3 supported)
    event : str
        Javascript code called during updates
    style : Union[Path, str, dict]
        A file or a dictionary defining a CSS file
    init_data : Optional[DataInput]
        Initial data before updates

    Examples
    --------

    Check :ref:`the full example<Update Guide>` for more information

    >>> update = function("data", "xmax", name="update")
    >>> update(d3(content="x").domain([0, js("xmax")]).nice())
    >>> update(svg.selectAll("g.xaxis").call(d3.axisBottom(x)))
    >>> update(
    ...     svg(content="line").datum(data)
    ...       .attr("d", d3.line()
    ...       .x(function("d").inline("x(d.x)"))
    ...       .y(function("d").inline("y(d.y)"))
    ...     )
    ... )
    >>> script(update)
    >>> def generator():
    ...     s = 1000
    ...     istep = 4 * pi / 1000
    ...     xv = [istep * i for i in range(1000)]
    ...     yv = list(map(sin, xv))
    ...     for i in range(1000):
    ...         xmax = istep * (s + i)
    ...         xv.append(xmax)
    ...         yv.append(sin(istep * (s + i)))
    ...         yield {"values": arrange([xv, yv]), "xmax": xmax}
    >>> websocket_render(
    ...     generator,
    ...     script,
    ...     event=js("update(received_data.values, received_data.xmax);"),
    ...     init_data=data
    ... )
    """
    from quart import Quart, websocket
    try:
        import nest_asyncio
        nest_asyncio.apply()
        from IPython import get_ipython
        from IPython.display import HTML, display
        JUPYTER_INSTALLED = True
    except:
        JUPYTER_INSTALLED = False
    if JUPYTER_INSTALLED and jupyter_environment():
        raise EnvironmentError("Dynamic updates in Jupyter environment unsupported")
    else:
        app = Quart("detroit")

        @app.websocket("/ws")
        async def ws():
            for data in generator():
                await websocket.send_json(data)

        @app.route("/")
        async def main():
            return await html(data=arrange(init_data), plot=script, style=style, event=event)

        app.run()
