"""jlabx — Launch JupyterLab with extensions."""

from __future__ import annotations

import ast
import difflib
import json
import os
import signal
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

from jlabx._version import __version__  # noqa: E402

CORE_EXTENSIONS = [
    "jupyterlab",
    "jupyter-collaboration",
    "jupyter-lsp",
    "python-lsp-server",
]

DEFAULT_USER_EXTENSIONS = [
    "jupyterlab-vim",
    "jupyterlab-myst",
    "jupyterlab-git",
]


# ---------------------------------------------------------------------------
# Config file management
# ---------------------------------------------------------------------------


def _config_dir() -> Path:
    xdg = os.environ.get("XDG_CONFIG_HOME")
    if xdg:
        return Path(xdg) / "jlabx"
    return Path.home() / ".config" / "jlabx"


def _config_file() -> Path:
    return _config_dir() / "config.toml"


def _write_config(extensions: list[str]) -> None:
    """Write config file with given extensions."""
    config_dir = _config_dir()
    config_dir.mkdir(parents=True, exist_ok=True)

    lines = [
        "# jlabx configuration",
        "#",
        "# Extra packages to include when launching JupyterLab.",
        "# Core extensions (jupyter-collaboration, jupyter-lsp, python-lsp-server)",
        "# are always included and don't need to be listed here.",
        "#",
        "# To skip these extras temporarily, run: jlabx --no-extras",
        "",
        "extensions = [",
    ]
    for ext in extensions:
        lines.append(f'    "{ext}",')
    lines.append("]")

    _config_file().write_text("\n".join(lines) + "\n")


def _parse_extensions() -> list[str]:
    """Parse extensions from the TOML config file."""
    cf = _config_file()
    if not cf.exists():
        return []

    extensions: list[str] = []
    in_array = False

    for raw_line in cf.read_text().splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue

        # Detect start of extensions array
        if line.startswith("extensions") and "=" in line and "[" in line:
            in_array = True
            after = line.split("[", 1)[1]
            if "]" in after:
                after = after.split("]")[0]
                in_array = False
            for part in after.split(","):
                part = part.strip().strip('"').strip("'").strip()
                if part:
                    extensions.append(part)
            continue

        if in_array:
            if "]" in line:
                line = line.split("]")[0]
                in_array = False
            # Strip inline comments
            if "#" in line:
                line = line[: line.index("#")]
            line = line.strip().rstrip(",").strip('"').strip("'").strip()
            if line:
                extensions.append(line)

    return extensions


def _ensure_config() -> None:
    """Create default config on first run."""
    cf = _config_file()
    if not cf.exists():
        _write_config(DEFAULT_USER_EXTENSIONS)
        print(f"Created config file: {cf}")
        print("  Edit it to customize which extensions are loaded.")
        print()


# ---------------------------------------------------------------------------
# Subcommands
# ---------------------------------------------------------------------------


def _cmd_list() -> None:
    print("Core extensions (always included):")
    for pkg in CORE_EXTENSIONS:
        print(f"  - {pkg}")
    print()

    cf = _config_file()
    exts = _parse_extensions()
    print(f"User extensions (from {cf}):")
    if not exts:
        print("  (none)")
    else:
        for pkg in exts:
            print(f"  + {pkg}")


def _cmd_add(packages: list[str]) -> None:
    if not packages:
        print("Usage: jlabx add <package> [package...]")
        sys.exit(1)

    current = _parse_extensions()
    for pkg in packages:
        if pkg in current:
            print(f"Already configured: {pkg}")
        else:
            current.append(pkg)
            print(f"Added: {pkg}")
    _write_config(current)


def _cmd_remove(packages: list[str]) -> None:
    if not packages:
        print("Usage: jlabx remove <package> [package...]")
        sys.exit(1)

    current = _parse_extensions()
    new_list: list[str] = []
    removed: set[str] = set()
    for ext in current:
        if ext in packages:
            print(f"Removed: {ext}")
            removed.add(ext)
        else:
            new_list.append(ext)

    for pkg in packages:
        if pkg not in removed:
            print(f"Not found: {pkg}")

    _write_config(new_list)


def _check_uv() -> None:
    """Ensure uv is available."""
    if shutil.which("uv"):
        return

    print("Error: 'uv' not found in PATH.")
    print()
    print("Install uv:")
    print("  curl -LsSf https://astral.sh/uv/install.sh | sh")
    print()
    print("Or see: https://docs.astral.sh/uv/getting-started/installation/")
    sys.exit(1)


def _print_extensions(no_extras: bool, user_extensions: list[str]) -> None:
    """Print what extensions are being loaded."""
    print("Core extensions:")
    for pkg in CORE_EXTENSIONS:
        print(f"  - {pkg}")
    if no_extras:
        print("User extensions: skipped (--no-extras)")
    elif user_extensions:
        print(f"User extensions (from {_config_file()}):")
        for pkg in user_extensions:
            print(f"  + {pkg}")
    else:
        print("User extensions: none configured")
    print()


def _run_with_signals(cmd: list[str]) -> None:
    """Launch a subprocess with signal forwarding and wait for exit."""
    proc = subprocess.Popen(cmd)

    def handler(sig: int, frame: object) -> None:
        print("\nShutting down JupyterLab...")
        proc.terminate()
        try:
            proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            print("Force killing hung process...")
            proc.kill()
        sys.exit(0)

    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)

    sys.exit(proc.wait())


_EMPTY_NOTEBOOK: dict = {
    "cells": [],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3",
        },
        "language_info": {"name": "python", "version": "3.12.0"},
    },
    "nbformat": 4,
    "nbformat_minor": 5,
}


def _create_notebook(path: str) -> dict:
    """Create an empty notebook file and return its contents."""
    nb = _EMPTY_NOTEBOOK.copy()
    with open(path, "w") as f:
        json.dump(nb, f, indent=1)
        f.write("\n")
    return nb


def _find_similar_notebooks(notebook: str) -> list[str]:
    """Find .ipynb files in the same directory with similar names."""
    path = Path(notebook)
    parent = path.parent or Path(".")
    name = path.name
    candidates = [p.name for p in parent.glob("*.ipynb")]
    if not candidates:
        return []
    return difflib.get_close_matches(name, candidates, n=3, cutoff=0.6)


def _resolve_notebook(notebook: str) -> str:
    """Resolve a notebook path, prompting for similar names or creating new.

    Returns the notebook path to use (may differ from input if user picks
    a suggestion).
    """
    if Path(notebook).exists():
        return notebook

    similar = _find_similar_notebooks(notebook)
    if not similar:
        print(f"Creating new notebook: {notebook}")
        _create_notebook(notebook)
        return notebook

    print(f"{notebook} does not exist.")
    print()
    for i, name in enumerate(similar, 1):
        print(f"  {i}) {name}")
    print(f"  c) Create {Path(notebook).name}")
    print()

    try:
        choice = input("Choice [1]: ").strip()
    except (EOFError, KeyboardInterrupt):
        print()
        sys.exit(1)

    if choice == "c":
        print(f"Creating new notebook: {notebook}")
        _create_notebook(notebook)
        return notebook

    # Default to first option
    if not choice:
        choice = "1"

    try:
        idx = int(choice) - 1
        if 0 <= idx < len(similar):
            picked = str(Path(notebook).parent / similar[idx])
            print(f"Launching {picked}")
            return picked
    except ValueError:
        pass

    print(f"Invalid choice. Creating {notebook}")
    _create_notebook(notebook)
    return notebook


def _read_notebook(notebook: str) -> dict:
    """Read a notebook as JSON. Returns empty dict on error."""
    try:
        with open(notebook) as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}


def _has_juv_metadata(nb: dict) -> bool:
    """Check if a notebook has PEP 723 inline script metadata."""
    for cell in nb.get("cells", []):
        if cell.get("cell_type") != "code":
            continue
        source = "".join(cell.get("source", []))
        if "# /// script" in source:
            return True
    return False


def _extract_imports(nb: dict) -> list[str]:
    """Extract top-level import names from notebook code cells."""
    modules: set[str] = set()
    for cell in nb.get("cells", []):
        if cell.get("cell_type") != "code":
            continue
        source = "".join(cell.get("source", []))
        try:
            tree = ast.parse(source)
        except SyntaxError:
            continue
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    modules.add(alias.name.split(".")[0])
            elif isinstance(node, ast.ImportFrom):
                if node.module and node.level == 0:
                    modules.add(node.module.split(".")[0])
    return sorted(modules)


def _filter_third_party(modules: list[str]) -> list[str]:
    """Filter out stdlib and obviously internal modules."""
    # sys.stdlib_module_names is available on 3.10+ (our minimum)
    stdlib: frozenset[str] = sys.stdlib_module_names  # type: ignore[attr-defined]
    return [m for m in modules if m not in stdlib and not m.startswith("_")]


# Import name → PyPI package name (from github.com/ianhi/script-bisect)
_IMPORT_TO_PACKAGE: dict[str, str] = {
    # Image and computer vision
    "cv2": "opencv-python",
    "PIL": "Pillow",
    "skimage": "scikit-image",
    # Machine learning and data science
    "sklearn": "scikit-learn",
    # Data formats and parsing
    "yaml": "PyYAML",
    "bs4": "beautifulsoup4",
    # Authentication and security
    "jwt": "PyJWT",
    # Date and time
    "dateutil": "python-dateutil",
    # Environment and configuration
    "dotenv": "python-dotenv",
    # File type detection
    "magic": "python-magic",
    # Hardware
    "serial": "pyserial",
    "usb": "pyusb",
    # Database
    "psycopg2": "psycopg2-binary",
    # Attrs
    "attr": "attrs",
    # System
    "gi": "pygobject",
}


def _imports_to_packages(modules: list[str]) -> list[str]:
    """Map import names to likely PyPI package names."""
    return [_IMPORT_TO_PACKAGE.get(m, m) for m in modules]


def _cmd_notebook(notebook: str, args: list[str]) -> None:
    """Launch a single notebook via juv with jlabx extensions."""
    _check_uv()

    no_extras = "--no-extras" in args
    init_deps = "--init-deps" in args
    passthrough = [a for a in args if a not in ("--no-extras", "--init-deps")]

    user_extensions = [] if no_extras else _parse_extensions()

    # juv handles jupyterlab itself, so only pass the other core extensions
    extras = [pkg for pkg in CORE_EXTENSIONS if pkg != "jupyterlab"] + user_extensions

    notebook = _resolve_notebook(notebook)

    nb = _read_notebook(notebook)
    if not nb:
        print(f"Error: could not read {notebook}")
        sys.exit(1)

    # If no PEP 723 metadata, detect imports from notebook cells
    detected_packages: list[str] = []
    if not _has_juv_metadata(nb):
        imports = _filter_third_party(_extract_imports(nb))
        detected_packages = _imports_to_packages(imports)

        if init_deps:
            # Write PEP 723 metadata into the notebook via juv
            print(f"Initializing PEP 723 metadata in {notebook}...")
            subprocess.run(["uvx", "juv", "init", notebook], check=True)
            if detected_packages:
                print(f"  Adding detected deps: {', '.join(detected_packages)}")
                subprocess.run(
                    ["uvx", "juv", "add", notebook] + detected_packages, check=True
                )
            print()
            # Clear detected_packages since juv now manages them
            detected_packages = []
        else:
            if detected_packages:
                print(
                    f"No PEP 723 metadata — detected imports: {', '.join(detected_packages)}"
                )
            else:
                print(f"No PEP 723 metadata or third-party imports in {notebook}")
            print("  Tip: use --init-deps to persist deps in the notebook via juv")
            print()

    _print_extensions(no_extras, user_extensions)
    print(f"Notebook mode — launching {notebook} via juv")
    print()

    # Build --with args: jlabx extensions + any auto-detected packages
    with_args: list[str] = []
    for pkg in extras + detected_packages:
        with_args.extend(["--with", pkg])

    # Port handling
    port = os.environ.get("JUPYTER_PORT", "8888")
    has_port = any("--port" in a for a in passthrough)
    port_args = [] if has_port else ["--port", port]

    # juv passthrough args go after --
    juv_passthrough = (
        ["--"] + port_args + passthrough if (port_args or passthrough) else []
    )

    cmd = ["uvx", "juv", "run"] + with_args + [notebook] + juv_passthrough
    _run_with_signals(cmd)


def _cmd_launch(args: list[str]) -> None:
    _check_uv()

    no_extras = "--no-extras" in args
    force_uv = "--uv" in args

    # Parse --python flag
    python_version: str | None = None
    passthrough: list[str] = []
    skip_next = False
    for i, a in enumerate(args):
        if skip_next:
            skip_next = False
            continue
        if a in ("--no-extras", "--uv"):
            continue
        if a == "--python" and i + 1 < len(args):
            python_version = args[i + 1]
            skip_next = True
            continue
        if a.startswith("--python="):
            python_version = a.split("=", 1)[1]
            continue
        passthrough.append(a)

    # Build extension list
    user_extensions = [] if no_extras else _parse_extensions()
    all_extensions = CORE_EXTENSIONS + user_extensions

    _print_extensions(no_extras, user_extensions)

    # Port handling
    port = os.environ.get("JUPYTER_PORT", "8888")
    has_port = any("--port" in a for a in passthrough)
    port_args = [] if has_port else ["--port", port]

    # Build --with args (for uv modes)
    with_args: list[str] = []
    for pkg in all_extensions:
        with_args.extend(["--with", pkg])

    python_args = ["--python", python_version] if python_version else []

    # Detect environment and build command
    cwd = Path.cwd()
    is_pixi = not force_uv and (
        (cwd / "pixi.toml").exists() or (cwd / "pixi.lock").exists()
    )
    if is_pixi:
        print("Pixi project detected — ephemeral venv for extensions")
        pixi_python = subprocess.run(
            ["pixi", "run", "which", "python"],
            capture_output=True, text=True, check=True,
        ).stdout.strip()

        tmpdir = tempfile.mkdtemp(prefix="jlabx-")
        venv_path = os.path.join(tmpdir, ".venv")
        try:
            subprocess.run(
                [pixi_python, "-m", "venv", "--system-site-packages", venv_path],
                check=True,
            )
            venv_python = os.path.join(venv_path, "bin", "python")
            subprocess.run(
                ["uv", "pip", "install", "--quiet", "--python", venv_python]
                + all_extensions,
                check=True,
            )
            jupyter_lab = os.path.join(venv_path, "bin", "jupyter-lab")
            cmd = [jupyter_lab] + port_args + passthrough
            _run_with_signals(cmd)
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)
        return
    elif (cwd / "pyproject.toml").exists():
        print("Python project detected — installing local package + extras...")
        cmd = (
            ["uv", "run"]
            + python_args
            + with_args
            + ["jupyter-lab"]
            + port_args
            + passthrough
        )
    else:
        print("Standalone mode — launching JupyterLab with extras...")
        cmd = (
            ["uvx"]
            + python_args
            + ["--from", "jupyterlab"]
            + with_args
            + ["jupyter-lab"]
            + port_args
            + passthrough
        )

    _run_with_signals(cmd)


def _cmd_help() -> None:
    print("jlabx — Launch JupyterLab with extensions")
    print()
    print("Usage:")
    print("  jlabx                        Launch JupyterLab")
    print("  jlabx notebook.ipynb         Launch a notebook via juv (creates if new)")
    print("  jlabx notebook.ipynb --init-deps  Detect imports & persist as juv deps")
    print("  jlabx --uv                   Force uv mode in a pixi project")
    print("  jlabx --python 3.11          Use a specific Python version (uv mode)")
    print("  jlabx --no-extras [args]     Launch without user extensions")
    print("  jlabx list                   Show configured extensions")
    print("  jlabx add <pkg> [pkg...]     Add user extensions")
    print("  jlabx remove <pkg> [pkg...]  Remove user extensions")
    print()
    print(f"Config: {_config_file()}")
    print()
    print("Install: uv tool install jlabx")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main() -> None:
    _ensure_config()

    args = sys.argv[1:]

    if not args:
        _cmd_launch([])
        return

    # Check if any arg is a .ipynb file → notebook mode
    # Also catch common extension typos (.ipnyb, .ipnb, .inpyb, etc.)
    _NB_EXTENSIONS = {".ipynb", ".ipnyb", ".ipnb", ".inpyb", ".ipybn"}

    def _is_notebook_arg(arg: str) -> bool:
        return any(arg.endswith(ext) for ext in _NB_EXTENSIONS)

    notebooks = [a for a in args if _is_notebook_arg(a)]
    if notebooks:
        nb_arg = notebooks[0]
        # Fix the extension if mistyped
        for ext in _NB_EXTENSIONS - {".ipynb"}:
            if nb_arg.endswith(ext):
                corrected = nb_arg[: -len(ext)] + ".ipynb"
                print(f"Correcting extension: {nb_arg} → {corrected}")
                nb_arg = corrected
                break
        rest = [a for a in args if not _is_notebook_arg(a)]
        _cmd_notebook(nb_arg, rest)
        return

    cmd = args[0]
    if cmd == "list":
        _cmd_list()
    elif cmd == "add":
        _cmd_add(args[1:])
    elif cmd == "remove":
        _cmd_remove(args[1:])
    elif cmd in ("--help", "-h"):
        _cmd_help()
    elif cmd == "--version":
        print(f"jlabx {__version__}")
    else:
        _cmd_launch(args)


if __name__ == "__main__":
    main()
