"""
Notebook Handler

Uses nbformat for reading/writing Jupyter notebooks with full compatibility.
Handles version upgrades (v3→v4), validation, and sanitization.
"""

import asyncio
import uuid
from pathlib import Path
from typing import Any
import weakref

import nbformat
from nbformat import ValidationError
from nbformat.v4 import new_notebook, new_code_cell, new_markdown_cell


# Always use nbformat v4 internally
NBFORMAT_VERSION = 4


class CellSourceConflictError(Exception):
    """Raised when a cell's current source doesn't match the expected source."""

    def __init__(self, current_source: str):
        self.current_source = current_source
        super().__init__(f"Cell source changed (expected_source mismatch)")

# Maximum notebook file size (100MB) - notebooks can be large with embedded outputs
MAX_NOTEBOOK_FILE_SIZE = 100 * 1024 * 1024

# Per-file locks scoped by event loop to prevent cross-loop issues
#
# Structure:
#   _file_locks[loop] -> {file_path: Lock}
#
# WeakKeyDictionary ensures locks are discarded automatically when
# an event loop is garbage-collected (e.g., after a runner restart).
#
# Note: We intentionally do NOT clean up individual file locks after use.
# Eager cleanup creates a race condition where the lock is removed while
# another coroutine is waiting to acquire it. The memory overhead of keeping
# locks for accessed files is negligible, and they're automatically cleaned
# up when the event loop is garbage collected.
_file_locks: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Lock]]" = weakref.WeakKeyDictionary()


def _get_file_lock(path: Path) -> asyncio.Lock:
    """Get or create an asyncio lock for a specific file path in the current event loop."""
    loop = asyncio.get_running_loop()
    loop_locks = _file_locks.setdefault(loop, {})
    file_key = str(path.resolve())
    return loop_locks.setdefault(file_key, asyncio.Lock())


def read_notebook(path: str | Path) -> dict[str, Any]:
    """
    Read a notebook file with validation and version upgrade.

    Tries UTF-8 first, falls back to latin-1 for notebooks with
    non-UTF-8 characters (e.g., created on Windows).

    Args:
        path: Path to .ipynb file

    Returns:
        Notebook as dict (JSON-serializable)

    Raises:
        FileNotFoundError: If file doesn't exist
        ValidationError: If notebook is invalid
        ValueError: If notebook is too large
    """
    path = Path(path)

    # Security: Check file size before reading to prevent OOM
    if path.exists():
        file_size = path.stat().st_size
        if file_size > MAX_NOTEBOOK_FILE_SIZE:
            raise ValueError(
                f"Notebook too large: {file_size / (1024*1024):.1f}MB "
                f"(max {MAX_NOTEBOOK_FILE_SIZE / (1024*1024):.0f}MB)"
            )

    # Try UTF-8 first (most common), fall back to latin-1
    nb = None
    for encoding in ("utf-8", "latin-1"):
        try:
            with open(path, "r", encoding=encoding) as f:
                # nbformat.read handles version detection and upgrades
                nb = nbformat.read(f, as_version=NBFORMAT_VERSION)
            break
        except UnicodeDecodeError:
            if encoding == "latin-1":
                raise  # Should never happen, latin-1 accepts all bytes
            continue

    if nb is None:
        raise ValueError(f"Could not decode notebook: {path}")

    # Validate the notebook
    nbformat.validate(nb)

    # Ensure all cells have IDs (required in nbformat 5.x)
    _ensure_cell_ids(nb)

    # Return as plain dict for JSON serialization
    return dict(nb)


def write_notebook(path: str | Path, notebook: dict[str, Any]) -> None:
    """
    Write a notebook file with validation.

    Args:
        path: Path to .ipynb file
        notebook: Notebook as dict

    Raises:
        ValidationError: If notebook is invalid
    """
    path = Path(path)

    # Convert dict to NotebookNode for validation
    nb = nbformat.from_dict(notebook)

    # Validate before writing
    nbformat.validate(nb)

    # Ensure all cells have IDs
    _ensure_cell_ids(nb)

    # Ensure parent directory exists
    path.parent.mkdir(parents=True, exist_ok=True)

    with open(path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)


async def write_notebook_safe(path: str | Path, notebook: dict[str, Any]) -> None:
    """Lock-aware async wrapper around write_notebook.

    Acquires the per-file lock so full-notebook writes don't race with
    granular cell updates (update_cell, update_cell_outputs, etc.).
    Also offloads blocking I/O to a thread.
    """
    path = Path(path)
    lock = _get_file_lock(path)
    async with lock:
        await asyncio.to_thread(write_notebook, path, notebook)


def validate_notebook(notebook: dict[str, Any]) -> list[str]:
    """
    Validate a notebook and return any errors.

    Args:
        notebook: Notebook as dict

    Returns:
        List of validation error messages (empty if valid)
    """
    errors = []

    try:
        nb = nbformat.from_dict(notebook)
        nbformat.validate(nb)
    except ValidationError as e:
        errors.append(str(e))
    except Exception as e:
        errors.append(f"Invalid notebook structure: {e}")

    return errors


def create_empty_notebook() -> dict[str, Any]:
    """
    Create a new empty notebook with Python 3 kernel.

    Returns:
        Empty notebook as dict
    """
    nb = new_notebook()
    nb.metadata.kernelspec = {
        "name": "python3",
        "display_name": "Python 3",
        "language": "python",
    }
    nb.metadata.language_info = {
        "name": "python",
        "version": "3.11",
        "file_extension": ".py",
        "mimetype": "text/x-python",
    }
    return dict(nb)


async def update_cell(
    path: str | Path,
    cell_id: str,
    source: str,
    expected_source: str | None = None,
) -> str:
    """
    Update a specific cell's source code in a notebook.

    This is safer than patching raw JSON because it:
    1. Acquires a per-file lock to prevent concurrent modifications
    2. Reads the notebook with validation
    3. Updates only the specified cell
    4. Writes back with validation

    Args:
        path: Path to .ipynb file
        cell_id: ID of the cell to update
        source: New source code for the cell
        expected_source: If provided, the cell's current source must match
            this value or a CellSourceConflictError is raised.

    Returns:
        The cell_type of the updated cell ("code" or "markdown").

    Raises:
        FileNotFoundError: If file doesn't exist
        ValueError: If cell_id not found
        CellSourceConflictError: If expected_source doesn't match current source
        ValidationError: If notebook is invalid
    """
    path = Path(path)

    lock = _get_file_lock(path)
    async with lock:
        return await asyncio.to_thread(
            _update_cell_sync, path, cell_id, source, expected_source
        )


def _update_cell_sync(
    path: Path,
    cell_id: str,
    source: str,
    expected_source: str | None = None,
) -> str:
    """Perform the blocking notebook read/modify/write flow in a thread.

    Returns the cell_type of the updated cell.
    """
    # Read and validate
    with open(path, "r", encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=NBFORMAT_VERSION)

    nbformat.validate(nb)
    _ensure_cell_ids(nb)

    # Find and update the cell
    cell_type: str | None = None
    for cell in nb.cells:
        if cell.get("id") == cell_id:
            # Optimistic concurrency check
            if expected_source is not None and cell["source"] != expected_source:
                raise CellSourceConflictError(cell["source"])
            cell["source"] = source
            cell_type = cell["cell_type"]  # guaranteed by nbformat
            break

    if cell_type is None:
        raise ValueError(f"Cell not found: {cell_id}")

    # Validate after update
    nbformat.validate(nb)

    # Write back
    with open(path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)

    return cell_type


async def update_cell_outputs(
    path: str | Path,
    cell_id: str,
    outputs: list[dict[str, Any]],
    execution_count: int | None = None,
) -> bool:
    """
    Update a code cell's outputs and execution_count on disk.

    Acquires the per-file lock, reads the notebook, finds the cell,
    updates its outputs and execution_count, validates, and writes back.

    Returns:
        True if the cell was found and updated, False if cell not found
        or not a code cell (cell may have been deleted between execution
        and save).
    """
    path = Path(path)
    lock = _get_file_lock(path)
    async with lock:
        return await asyncio.to_thread(
            _update_cell_outputs_sync, path, cell_id, outputs, execution_count
        )


def _update_cell_outputs_sync(
    path: Path,
    cell_id: str,
    outputs: list[dict[str, Any]],
    execution_count: int | None,
) -> bool:
    """Perform the blocking outputs update in a thread."""
    with open(path, "r", encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=NBFORMAT_VERSION)

    nbformat.validate(nb)
    _ensure_cell_ids(nb)

    for cell in nb.cells:
        if cell.get("id") == cell_id:
            if cell.get("cell_type") != "code":
                return False
            cell["outputs"] = [nbformat.from_dict(o) for o in outputs]
            if execution_count is not None:
                cell["execution_count"] = execution_count
            break
    else:
        return False

    nbformat.validate(nb)

    with open(path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)

    return True


async def add_cell(
    path: str | Path,
    cell_type: str,
    source: str = "",
    after_cell_id: str | None = None,
) -> str:
    """
    Add a new cell to a notebook.

    Args:
        path: Path to .ipynb file
        cell_type: Type of cell ("code" or "markdown")
        source: Initial source code for the cell
        after_cell_id: Insert after this cell ID. If None, appends to end.

    Returns:
        The ID of the newly created cell

    Raises:
        FileNotFoundError: If file doesn't exist
        ValueError: If cell_type is invalid or after_cell_id not found
        ValidationError: If notebook is invalid
    """
    path = Path(path)

    lock = _get_file_lock(path)
    async with lock:
        cell_id = await asyncio.to_thread(
            _add_cell_sync, path, cell_type, source, after_cell_id
        )

    return cell_id


def _add_cell_sync(
    path: Path,
    cell_type: str,
    source: str,
    after_cell_id: str | None,
) -> str:
    """Perform the blocking add cell operation in a thread."""
    # Read and validate
    with open(path, "r", encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=NBFORMAT_VERSION)

    nbformat.validate(nb)
    _ensure_cell_ids(nb)

    # Create the new cell
    cell_id = str(uuid.uuid4())
    if cell_type == "code":
        new_cell = new_code_cell(source=source, id=cell_id)
    elif cell_type == "markdown":
        new_cell = new_markdown_cell(source=source, id=cell_id)
    else:
        raise ValueError(f"Invalid cell_type: {cell_type}. Must be 'code' or 'markdown'")

    # Find insertion position
    if after_cell_id is None:
        # Append to end
        nb.cells.append(new_cell)
    elif after_cell_id == "__top__":
        # Insert at the beginning
        nb.cells.insert(0, new_cell)
    else:
        # Find the cell and insert after it
        insert_idx = None
        for i, cell in enumerate(nb.cells):
            if cell.get("id") == after_cell_id:
                insert_idx = i + 1
                break

        if insert_idx is None:
            raise ValueError(f"Cell not found: {after_cell_id}")

        nb.cells.insert(insert_idx, new_cell)

    # Validate after modification
    nbformat.validate(nb)

    # Write back
    with open(path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)

    return cell_id


async def delete_cell(path: str | Path, cell_id: str) -> None:
    """
    Delete a cell from a notebook.

    Args:
        path: Path to .ipynb file
        cell_id: ID of the cell to delete

    Raises:
        FileNotFoundError: If file doesn't exist
        ValueError: If cell_id not found
        ValidationError: If notebook is invalid
    """
    path = Path(path)

    lock = _get_file_lock(path)
    async with lock:
        await asyncio.to_thread(_delete_cell_sync, path, cell_id)


def _delete_cell_sync(path: Path, cell_id: str) -> None:
    """Perform the blocking delete cell operation in a thread."""
    # Read and validate
    with open(path, "r", encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=NBFORMAT_VERSION)

    nbformat.validate(nb)
    _ensure_cell_ids(nb)

    # Find and remove the cell
    cell_found = False
    for i, cell in enumerate(nb.cells):
        if cell.get("id") == cell_id:
            nb.cells.pop(i)
            cell_found = True
            break

    if not cell_found:
        raise ValueError(f"Cell not found: {cell_id}")

    # Validate after modification
    nbformat.validate(nb)

    # Write back
    with open(path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)


def _ensure_cell_ids(nb: nbformat.NotebookNode) -> None:
    """Ensure all cells have unique IDs."""
    existing_ids: set[str] = set()

    for cell in nb.cells:
        if not cell.get("id") or cell["id"] in existing_ids:
            cell["id"] = str(uuid.uuid4())
        existing_ids.add(cell["id"])
