from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import nbformat


@dataclass(frozen=True)
class Cell:
    """Represents a notebook cell with type and content."""

    cell_type: str  # "code" or "markdown"
    content: str
    index: int  # 1-indexed position in original notebook


def read_ipynb_cells(
    path: Path,
    *,
    include_markdown: bool = False,
    cell_indices: list[int] | None = None,
) -> list[Cell]:
    """
    Read cells from a notebook.

    Args:
        path: Path to .ipynb file
        include_markdown: If True, include markdown cells
        cell_indices: If provided, only include cells at these 1-indexed positions
    """
    nb = nbformat.read(path, as_version=4)
    cells: list[Cell] = []
    idx = 0

    for cell in nb.cells:
        cell_type = cell.get("cell_type", "")
        src = str(cell.get("source") or "").rstrip()

        if not src:
            continue

        if cell_type == "code":
            idx += 1
            if cell_indices is None or idx in cell_indices:
                cells.append(Cell(cell_type="code", content=src, index=idx))
        elif cell_type == "markdown" and include_markdown:
            idx += 1
            if cell_indices is None or idx in cell_indices:
                cells.append(Cell(cell_type="markdown", content=src, index=idx))

    return cells


# Backward-compatible wrapper
def read_ipynb_code_cells(
    path: Path,
    *,
    max_cells: int | None = None,
) -> list[str]:
    """Legacy function for backward compatibility."""
    cells = read_ipynb_cells(path, include_markdown=False)
    code_only = [c.content for c in cells if c.cell_type == "code"]
    if max_cells is not None:
        return code_only[:max_cells]
    return code_only
