# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any

import click

from marimo._cli.errors import MarimoCLIMissingDependencyError
from marimo._cli.export._common import (
    SandboxVenvPool,
    collect_notebooks,
    is_multi_target,
    run_python_subprocess,
)
from marimo._cli.export.output import STDERR
from marimo._cli.install_hints import get_playwright_chromium_setup_commands
from marimo._cli.parse_args import parse_args
from marimo._cli.print import echo, green, red, yellow
from marimo._dependencies.dependencies import DependencyManager
from marimo._export._html_asset_server import HtmlAssetServer
from marimo._export.file import export_html
from marimo._export.requests import (
    HTMLFileExportRequest,
    NotebookExecutionOptions,
)
from marimo._schemas.export_options import HTMLExportOptions
from marimo._server.utils import asyncio_run
from marimo._utils.marimo_path import MarimoPath
from marimo._utils.paths import marimo_package_path, maybe_make_dirs

if TYPE_CHECKING:
    from marimo._cli.sandbox import SandboxMode


_sandbox_message = (
    "Render notebooks in an isolated environment, with dependencies tracked "
    "via PEP 723 inline metadata. If already declared, dependencies will "
    "install automatically. Requires uv. Only applies when --execute is used."
)
_READINESS_WAIT_TIMEOUT_MS = 30_000
_sandbox_bootstrapped_env = "MARIMO_THUMBNAIL_SANDBOX_BOOTSTRAPPED"
_sandbox_mode_env = "MARIMO_THUMBNAIL_SANDBOX_MODE"
_thumbnail_sandbox_deps = ["playwright"]


def _split_paths_and_args(
    name: str, args: tuple[str, ...]
) -> tuple[list[str], tuple[str, ...]]:
    paths = [name]
    for index, arg in enumerate(args):
        if arg == "--":
            return paths, args[index + 1 :]
        paths.append(arg)
    return paths, ()


def _sandbox_mode_from_env() -> SandboxMode | None:
    from marimo._cli.sandbox import SandboxMode

    if os.environ.get(_sandbox_bootstrapped_env) != "1":
        return None

    mode = os.environ.get(_sandbox_mode_env)
    if mode == SandboxMode.SINGLE.value:
        return SandboxMode.SINGLE
    if mode == SandboxMode.MULTI.value:
        return SandboxMode.MULTI
    return None


def _resolve_thumbnail_sandbox_mode(
    *,
    execute: bool,
    sandbox: bool | None,
    path_targets: list[Path],
    first_target: str,
) -> SandboxMode | None:
    from marimo._cli.sandbox import SandboxMode, resolve_sandbox_mode

    if not execute:
        return None

    env_mode = _sandbox_mode_from_env()
    if env_mode is not None:
        return env_mode

    if is_multi_target(path_targets):
        if sandbox is None:
            return None
        return SandboxMode.MULTI if sandbox else None

    return resolve_sandbox_mode(sandbox=sandbox, name=first_target)


def _bootstrap_thumbnail_sandbox(
    *,
    args: list[str],
    name: str,
    sandbox_mode: SandboxMode,
) -> None:
    from marimo._cli.sandbox import run_in_sandbox

    run_in_sandbox(
        args,
        name=name,
        additional_deps=_thumbnail_sandbox_deps,
        extra_env={
            _sandbox_bootstrapped_env: "1",
            _sandbox_mode_env: sandbox_mode.value,
        },
    )


async def _render_html(
    marimo_path: MarimoPath,
    *,
    execute: bool,
    include_code: bool,
    args: tuple[str, ...],
    asset_url: str | None = None,
    venv_python: str | None = None,
) -> str:
    if not execute:
        result = await export_html(
            HTMLFileExportRequest(
                path=marimo_path,
                options=HTMLExportOptions(
                    files=(),
                    include_code=True,
                    asset_url=asset_url,
                ),
            )
        )
        return result.text

    if venv_python is None:
        cli_args = parse_args(args) if args else {}
        result = await export_html(
            HTMLFileExportRequest(
                path=marimo_path,
                options=HTMLExportOptions(
                    files=(),
                    include_code=include_code,
                    asset_url=asset_url,
                ),
                execution=NotebookExecutionOptions(
                    cli_args=cli_args,
                    argv=list(args),
                    stderr=STDERR,
                ),
            )
        )
        return result.text

    payload = {
        "path": marimo_path.absolute_name,
        "include_code": include_code,
        "args": list(args),
        "asset_url": asset_url,
    }

    # Render in a separate process so we can use a sandboxed venv without polluting the current environment.
    return await asyncio.to_thread(
        _render_html_in_subprocess,
        venv_python,
        payload,
    )


def _render_html_in_subprocess(
    venv_python: str, payload: dict[str, Any]
) -> str:
    """Render a notebook to HTML in a separate Python process."""
    script = r"""
import asyncio
import json
import sys

from marimo._cli.parse_args import parse_args
from marimo._export.file import export_html
from marimo._export.requests import (
    HTMLFileExportRequest,
    NotebookExecutionOptions,
)
from marimo._schemas.export_options import HTMLExportOptions
from marimo._utils.marimo_path import MarimoPath

payload = json.loads(sys.argv[1])
path = MarimoPath(payload["path"])
include_code = bool(payload.get("include_code", False))
args = payload.get("args") or []
asset_url = payload.get("asset_url")

cli_args = parse_args(tuple(args)) if args else {}
result = asyncio.run(
    export_html(
        HTMLFileExportRequest(
            path=path,
            options=HTMLExportOptions(
                files=(),
                include_code=include_code,
                asset_url=asset_url,
            ),
            execution=NotebookExecutionOptions(
                cli_args=cli_args,
                argv=list(args),
            ),
        )
    )
)
sys.stdout.write(result.text)
"""

    return run_python_subprocess(
        venv_python=venv_python,
        script=script,
        payload=payload,
        action="render notebook",
    )


async def _generate_thumbnails(
    *,
    notebooks: list[MarimoPath],
    width: int,
    height: int,
    scale: int,
    timeout_ms: int,
    output: Path | None,
    overwrite: bool,
    include_code: bool,
    execute: bool,
    notebook_args: tuple[str, ...],
    continue_on_error: bool,
    sandbox_mode: SandboxMode | None,
) -> None:
    from marimo._cli.sandbox import SandboxMode
    from marimo._metadata.opengraph import default_opengraph_image_abs

    failures: list[tuple[MarimoPath, Exception]] = []

    try:
        from playwright.async_api import (  # type: ignore[import-not-found]
            TimeoutError as PlaywrightTimeoutError,
            async_playwright,
        )
    except ModuleNotFoundError as e:
        if getattr(e, "name", None) == "playwright":
            raise MarimoCLIMissingDependencyError(
                "Playwright is required to generate thumbnails.",
                "nbconvert[webpdf]",
                followup_commands=get_playwright_chromium_setup_commands(),
            ) from None
        raise

    use_per_notebook_sandbox = sandbox_mode is SandboxMode.MULTI

    if use_per_notebook_sandbox and not DependencyManager.which("uv"):
        raise MarimoCLIMissingDependencyError(
            "uv is required for --sandbox thumbnail generation.",
            "uv",
            additional_tip="Install uv from https://github.com/astral-sh/uv",
        )

    static_dir = marimo_package_path() / "_static"

    sandbox_pool: SandboxVenvPool | None = (
        SandboxVenvPool() if use_per_notebook_sandbox else None
    )
    try:
        with HtmlAssetServer(
            directory=static_dir, route="/__marimo_thumbnail__.html"
        ) as server:
            async with async_playwright() as playwright:
                browser = await playwright.chromium.launch()
                context = await browser.new_context(
                    viewport={"width": width, "height": height},
                    device_scale_factor=scale,
                )
                page = await context.new_page()
                await page.emulate_media(reduced_motion="reduce")

                for index, notebook in enumerate(notebooks):
                    try:
                        out_path = (
                            output
                            if output is not None
                            else default_opengraph_image_abs(
                                str(notebook.path)
                            )
                        )
                        if out_path.exists() and not overwrite:
                            echo(
                                red("skip")
                                + f": {notebook.short_name} (exists, use --overwrite)"
                            )
                            continue

                        maybe_make_dirs(out_path)

                        echo(f"Rendering {notebook.short_name}...")
                        venv_python = (
                            sandbox_pool.get_python(str(notebook.path))
                            if sandbox_pool is not None
                            else None
                        )
                        html = await _render_html(
                            notebook,
                            execute=execute,
                            include_code=include_code,
                            args=notebook_args,
                            asset_url=server.base_url,
                            venv_python=venv_python,
                        )
                        server.set_html(html)

                        echo(f"Screenshotting -> {out_path}...")
                        page_url = f"{server.page_url}?v={index}"
                        await page.goto(page_url, wait_until="load")
                        # Hide chrome and watermarks marked for print so thumbnails stay focused on notebook content.
                        await page.add_style_tag(
                            content=(
                                '.print\\:hidden,[data-testid="watermark"]{display:none !important;}'
                            )
                        )
                        # Nb renderer starts cell contents as invisible for a short period to avoid flicker
                        # --> we wait for the first cell container to be visible before snapshotting.
                        try:
                            await page.wait_for_function(
                                r"""
() => {
  const root = document.getElementById("root");
  if (!root) return false;

  const cells = Array.from(document.querySelectorAll('div[id^="cell-"]'));
  if (cells.length > 0) {
    const hasVisibleCell = cells.some((cell) => {
      const style = window.getComputedStyle(cell);
      return style.visibility !== "hidden" && style.display !== "none";
    });
    if (hasVisibleCell) return true;
  }

  return root.childElementCount > 0;
}
""",
                                timeout=_READINESS_WAIT_TIMEOUT_MS,
                            )
                        except PlaywrightTimeoutError:
                            echo(
                                yellow("warning")
                                + ": readiness check timed out; capturing screenshot anyway."
                            )
                        await page.wait_for_timeout(timeout_ms)
                        await page.screenshot(
                            path=str(out_path), full_page=False
                        )

                        echo(green("ok") + f": {out_path}")
                    except Exception as e:
                        failures.append((notebook, e))
                        echo(red("error") + f": {notebook.short_name}: {e}")
                        if not continue_on_error:
                            raise

                await context.close()
                await browser.close()
    finally:
        if sandbox_pool is not None:
            sandbox_pool.close()

    if failures:
        raise click.ClickException(
            f"Failed to generate thumbnails for {len(failures)} notebooks."
        )


@click.command(
    "thumbnail", help="Generate OpenGraph thumbnails for notebooks."
)
@click.argument(
    "name",
    type=click.Path(
        exists=True, file_okay=True, dir_okay=True, path_type=Path
    ),
)
@click.option(
    "--width",
    type=int,
    default=1200,
    help="Viewport width for the screenshot.",
)
@click.option(
    "--height",
    type=int,
    default=630,
    help="Viewport height for the screenshot.",
)
@click.option(
    "--scale",
    type=click.IntRange(min=1, max=4),
    default=2,
    help=(
        "Device scale factor for screenshots. Output resolution will be "
        "`width*scale` x `height*scale`."
    ),
)
@click.option(
    "--timeout-ms",
    type=int,
    default=1500,
    help="Additional time to wait after page load before screenshot.",
)
@click.option(
    "--output",
    type=click.Path(path_type=Path),
    default=None,
    help=(
        "Output filename. If omitted, writes to "
        "`<notebook_dir>/__marimo__/assets/<notebook_stem>/opengraph.png`."
    ),
)
@click.option(
    "--overwrite/--no-overwrite",
    default=False,
    help="Overwrite existing thumbnails.",
)
@click.option(
    "--include-code/--no-include-code",
    default=False,
    help="Whether to include code in the rendered HTML before screenshot.",
)
@click.option(
    "--execute/--no-execute",
    default=False,
    help=(
        "Execute notebooks and include their outputs in thumbnails. "
        "In --no-execute mode (default), thumbnails are generated from notebook "
        "structure without running code (and will not include outputs)."
    ),
)
@click.option(
    "--sandbox/--no-sandbox",
    is_flag=True,
    default=None,
    type=bool,
    help=_sandbox_message,
)
@click.option(
    "--continue-on-error/--fail-fast",
    default=True,
    help="Continue processing other notebooks if one notebook fails.",
)
@click.argument("args", nargs=-1, type=click.UNPROCESSED)
def thumbnail(
    name: Path,
    width: int,
    height: int,
    scale: int,
    timeout_ms: int,
    output: Path | None,
    overwrite: bool,
    include_code: bool,
    execute: bool,
    sandbox: bool | None,
    continue_on_error: bool,
    args: tuple[str, ...],
) -> None:
    """Generate thumbnails for one or more notebooks (or directories)."""
    paths, notebook_args = _split_paths_and_args(str(name), args)
    path_targets = [Path(p) for p in paths]
    notebooks = collect_notebooks(path_targets)
    if not notebooks:
        raise click.ClickException("No marimo notebooks found.")
    if output is not None and len(notebooks) > 1:
        raise click.UsageError(
            "--output can only be used when generating thumbnail for a single notebook."
        )

    if not execute and sandbox:
        raise click.UsageError("--sandbox requires --execute.")

    sandbox_mode = _resolve_thumbnail_sandbox_mode(
        execute=execute,
        sandbox=sandbox,
        path_targets=path_targets,
        first_target=str(name),
    )

    if (
        execute
        and sandbox_mode is not None
        and _sandbox_mode_from_env() is None
    ):
        _bootstrap_thumbnail_sandbox(
            args=sys.argv[1:],
            name=str(name),
            sandbox_mode=sandbox_mode,
        )
        return

    try:
        DependencyManager.playwright.require("for thumbnail generation")
    except ModuleNotFoundError as e:
        if getattr(e, "name", None) == "playwright":
            raise MarimoCLIMissingDependencyError(
                "Playwright is required for thumbnail generation.",
                "playwright",
                followup_commands=get_playwright_chromium_setup_commands(),
            ) from None
        raise

    asyncio_run(
        _generate_thumbnails(
            notebooks=notebooks,
            width=width,
            height=height,
            scale=scale,
            timeout_ms=timeout_ms,
            output=output,
            overwrite=overwrite,
            include_code=include_code,
            execute=execute,
            notebook_args=notebook_args,
            continue_on_error=continue_on_error,
            sandbox_mode=sandbox_mode,
        )
    )
