"""
Export pipeline for nblite.

Handles exporting notebooks to different formats and to Python modules.
"""

from __future__ import annotations

import bisect
import re
from dataclasses import dataclass, field
from pathlib import Path

from nblite.config.schema import CellReferenceStyle, ExportMode
from nblite.core.cell import Cell
from nblite.core.notebook import Format, Notebook
from nblite.export.function_export import export_function_notebook, is_function_notebook
from nblite.extensions import HookRegistry, HookType

__all__ = [
    "export_notebook_to_notebook",
    "export_notebook_to_module",
    "export_notebooks_to_module",
    "get_export_targets",
    "ExportResult",
]


# Pattern to extract function/class names for __all__
FUNCTION_PATTERN = re.compile(r"^(?:async\s+)?def\s+(\w+)\s*\(", re.MULTILINE)
CLASS_PATTERN = re.compile(r"^class\s+(\w+)\s*[\(:]", re.MULTILINE)
# Match variable assignments at module level (not starting with _)
VARIABLE_PATTERN = re.compile(r"^([a-zA-Z][a-zA-Z0-9_]*)\s*=", re.MULTILINE)

# Pattern to match "from package.module import ..." statements
# Captures: (1) the package.module part, (2) the rest after "import"
FROM_IMPORT_PATTERN = re.compile(
    r"^(\s*)from\s+([a-zA-Z_][a-zA-Z0-9_.]*)\s+import\s+(.+)$",
    re.MULTILINE,
)


def _get_cell_order(cell: Cell) -> int:
    """
    Get the order value for an exported cell.

    Order values determine the position of cells in the exported module.
    Default values:
    - #|export: 0
    - #|exporti: 0
    - #|export_to: 1

    Args:
        cell: The cell to get order from

    Returns:
        The order value (lower = earlier in output)
    """
    # Check export_to first (has module and order)
    if cell.has_directive("export_to"):
        directive = cell.get_directive("export_to")
        if directive and directive.value_parsed:
            return directive.value_parsed.get("order", 1)

    # Check export/exporti
    for name in ("export", "exporti"):
        if cell.has_directive(name):
            directive = cell.get_directive(name)
            if directive and directive.value_parsed:
                return directive.value_parsed.get("order", 0)
            return 0  # Default if no parsed value

    return 0  # Fallback


def _validate_function_only_directives(notebook: Notebook) -> None:
    """
    Validate that function-only directives are only used in function notebooks.

    Raises:
        ValueError: If #|top_export or #|bottom_export used in non-function notebook.
    """
    if is_function_notebook(notebook):
        return  # Function notebook, all directives allowed

    for cell in notebook.cells:
        if cell.has_directive("top_export"):
            raise ValueError(
                f"#|top_export directive can only be used in function notebooks "
                f"(notebooks with #|export_as_func true). Found in cell {cell.index}."
            )
        if cell.has_directive("bottom_export"):
            raise ValueError(
                f"#|bottom_export directive can only be used in function notebooks "
                f"(notebooks with #|export_as_func true). Found in cell {cell.index}."
            )


@dataclass
class ExportResult:
    """Result of an export operation."""

    success: bool = True
    files_created: list[Path] = field(default_factory=list)
    files_updated: list[Path] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)


def export_notebook_to_notebook(
    notebook: Notebook,
    output_path: Path | str,
    format: str | None = None,
    *,
    no_header: bool = False,
) -> None:
    """
    Export a notebook to another notebook format.

    Args:
        notebook: Source notebook
        output_path: Output path for the notebook
        format: Output format (ipynb, percent). Auto-detected if None.
        no_header: If True, omit YAML frontmatter when serializing to percent format.
    """
    output_path = Path(output_path)

    if format is None:
        format = Format.from_path(output_path)

    content = notebook.to_string(format, no_header=no_header)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(content)


def export_notebook_to_module(
    notebook: Notebook,
    output_path: Path | str,
    project_root: Path | str,
    export_mode: ExportMode = ExportMode.PERCENT,
    include_warning: bool = True,
    cell_reference_style: CellReferenceStyle = CellReferenceStyle.RELATIVE,
    package_name: str | None = None,
    target_module: str | None = None,
) -> None:
    """
    Export a notebook to a Python module.

    Args:
        notebook: Source notebook
        output_path: Output path for the module
        project_root: Project root directory for computing relative paths
        export_mode: Export mode (percent with cell markers, or py without)
        include_warning: Include autogenerated warning header
        cell_reference_style: Style for cell references (relative or absolute)
        package_name: Package name for converting absolute imports to relative imports.
            If provided, imports like "from {package_name}.X import Y" will be
            converted to relative imports like "from .X import Y".
        target_module: If specified, only export cells targeting this module.
            Cells with #|export go to default_exp, cells with #|export_to
            go to their specified module.
    """
    output_path = Path(output_path)
    project_root = Path(project_root)

    # Validate function-only directives
    _validate_function_only_directives(notebook)

    # Check if this is a function notebook
    if is_function_notebook(notebook):
        export_function_notebook(
            notebook,
            output_path,
            include_warning=include_warning,
            package_name=package_name,
        )
        return

    source_path = notebook.source_path

    # Calculate relative path for cell references
    if source_path and cell_reference_style == CellReferenceStyle.RELATIVE:
        try:
            source_ref = str(source_path.relative_to(project_root))
        except ValueError:
            source_ref = str(source_path)
    elif cell_reference_style == CellReferenceStyle.ABSOLUTE:
        source_ref = str(source_path)
    elif cell_reference_style == CellReferenceStyle.NONE:
        source_ref = ""
    else:
        raise ValueError(f"Invalid cell reference style: {cell_reference_style}")

    # Calculate module depth for relative imports
    # e.g., my_lib/submodule/utils.py has depth 1 (one subdirectory)
    module_depth = 0
    if package_name:
        try:
            # Get the path relative to project root
            rel_to_project = output_path.relative_to(project_root)
            parts = rel_to_project.parts
            # Find where the package name appears in the path
            # (it may not be at index 0 if package is inside src/ or similar)
            try:
                pkg_index = parts.index(package_name)
                # module_depth = number of directories after the package, minus the filename
                # e.g., src/my_lib/submodule/utils.py with my_lib at index 1:
                #   parts after package = ('submodule', 'utils.py'), depth = 1
                module_depth = len(parts) - pkg_index - 2
            except ValueError:
                # Package name not found in path
                pass
        except ValueError:
            pass

    # Collect exported cells
    exported_content = _collect_exported_content(
        notebook, export_mode, source_ref, package_name, module_depth, target_module
    )

    # Build module content
    lines: list[str] = []

    # Add autogenerated header
    if include_warning:
        lines.append(f"# AUTOGENERATED! DO NOT EDIT! File to edit: {source_ref}")
        lines.append("")

    # Add __all__ list
    all_names = _extract_public_names(exported_content)
    # Add names from #|add_to_all directives
    all_names.extend(_collect_add_to_all_names(notebook))
    # Get names from #|exporti cells to exclude
    exporti_names = _collect_exporti_names(notebook)
    # Deduplicate while preserving order, excluding exporti names
    seen = set()
    unique_names = []
    for name in all_names:
        if name not in seen and name not in exporti_names:
            seen.add(name)
            unique_names.append(name)
    if unique_names:
        all_str = ", ".join(f"'{name}'" for name in sorted(unique_names))
        lines.append(f"__all__ = [{all_str}]")
        lines.append("")

    # Add content
    lines.append(exported_content)

    # Write output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text("\n".join(lines))


def export_notebooks_to_module(
    notebooks: list[tuple[Notebook, str]],
    output_path: Path | str,
    project_root: Path | str,
    export_mode: ExportMode = ExportMode.PERCENT,
    include_warning: bool = True,
    cell_reference_style: CellReferenceStyle = CellReferenceStyle.RELATIVE,
    package_name: str | None = None,
    target_module: str | None = None,
) -> None:
    """
    Export multiple notebooks to a single Python module.

    This aggregates cells from all notebooks that target the same module,
    sorting them by (order, notebook_path, cell_index) for deterministic output.

    Args:
        notebooks: List of (notebook, source_ref) tuples. source_ref is the
            relative path to use for cell references.
        output_path: Output path for the module
        project_root: Project root directory for computing relative paths
        export_mode: Export mode (percent with cell markers, or py without)
        include_warning: Include autogenerated warning header
        cell_reference_style: Style for cell references (relative or absolute)
        package_name: Package name for converting absolute imports to relative imports.
        target_module: If specified, only export cells targeting this module.
    """
    output_path = Path(output_path)
    project_root = Path(project_root)

    # Validate function-only directives in all notebooks
    for notebook, _source_ref in notebooks:
        _validate_function_only_directives(notebook)

    # Check if any notebook is a function notebook - these should not be aggregated
    for notebook, _source_ref in notebooks:
        if is_function_notebook(notebook):
            raise ValueError(
                f"Function notebooks cannot be aggregated with other notebooks. "
                f"Notebook {notebook.source_path} has #|export_as_func true."
            )

    # Calculate module depth for relative imports
    module_depth = 0
    if package_name:
        try:
            rel_to_project = output_path.relative_to(project_root)
            parts = rel_to_project.parts
            try:
                pkg_index = parts.index(package_name)
                module_depth = len(parts) - pkg_index - 2
            except ValueError:
                pass
        except ValueError:
            pass

    # Collect exported content from all notebooks
    exported_content, source_refs = _collect_exported_content_multi(
        notebooks, export_mode, package_name, module_depth, target_module
    )

    # Build module content
    lines: list[str] = []

    # Add autogenerated header with all source files
    if include_warning:
        if len(source_refs) == 1:
            lines.append(f"# AUTOGENERATED! DO NOT EDIT! File to edit: {source_refs[0]}")
        else:
            files_str = ", ".join(sorted(source_refs))
            lines.append(f"# AUTOGENERATED! DO NOT EDIT! Files to edit: {files_str}")
        lines.append("")

    # Add __all__ list (aggregated from all notebooks)
    all_names = _extract_public_names(exported_content)
    # Add names from #|add_to_all directives from all notebooks
    # Also collect exporti names to exclude
    exporti_names: set[str] = set()
    for nb, _source_ref in notebooks:
        all_names.extend(_collect_add_to_all_names(nb))
        exporti_names.update(_collect_exporti_names(nb))
    # Deduplicate while preserving order, excluding exporti names
    seen = set()
    unique_names = []
    for name in all_names:
        if name not in seen and name not in exporti_names:
            seen.add(name)
            unique_names.append(name)
    if unique_names:
        all_str = ", ".join(f"'{name}'" for name in sorted(unique_names))
        lines.append(f"__all__ = [{all_str}]")
        lines.append("")

    # Add content
    lines.append(exported_content)

    # Write output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text("\n".join(lines))


def _collect_exported_content_multi(
    notebooks: list[tuple[Notebook, str]],
    export_mode: ExportMode,
    package_name: str | None = None,
    module_depth: int = 0,
    target_module: str | None = None,
) -> tuple[str, list[str]]:
    """
    Collect content from multiple notebooks for a single module.

    Cells are sorted by (order, notebook_path, cell_index) for deterministic
    and stable ordering across notebooks.

    Args:
        notebooks: List of (notebook, source_ref) tuples
        export_mode: Export mode (percent or py)
        package_name: Package name for import transformation
        module_depth: Depth of module within package (for relative imports)
        target_module: If specified, only include cells targeting this module.

    Returns:
        Tuple of (aggregated content, list of source file paths that contributed)

    Hooks triggered:
        PRE_CELL_EXPORT: Before each cell export
        POST_CELL_EXPORT: After each cell export
    """
    # Collect all cells with their metadata for sorting
    # Each entry: (order, notebook_path_str, cell_index, cell, notebook, source_ref)
    cells_to_export: list[tuple[int, str, int, Cell, Notebook, str]] = []
    source_refs_set: set[str] = set()

    for notebook, source_ref in notebooks:
        default_exp = notebook.default_exp
        notebook_path_str = str(notebook.source_path) if notebook.source_path else ""

        for cell in notebook.cells:
            if not cell.is_code:
                continue

            # Check for export directives
            has_export = cell.has_directive("export") or cell.has_directive("exporti")
            has_export_to = cell.has_directive("export_to")

            if not (has_export or has_export_to):
                continue

            # Filter by target module if specified
            if target_module is not None:
                cell_target = None
                if has_export_to:
                    export_to_directive = cell.get_directive("export_to")
                    if export_to_directive and export_to_directive.value_parsed:
                        cell_target = export_to_directive.value_parsed.get("module", "")
                elif has_export:
                    cell_target = default_exp

                if cell_target != target_module:
                    continue

            order = _get_cell_order(cell)
            cells_to_export.append(
                (order, notebook_path_str, cell.index, cell, notebook, source_ref)
            )
            source_refs_set.add(source_ref)

    # Sort by (order, notebook_path, cell_index) for deterministic output
    cells_to_export.sort(key=lambda x: (x[0], x[1], x[2]))

    # Output cells in sorted order
    parts: list[str] = []

    for _order, _nb_path, _cell_idx, cell, notebook, source_ref in cells_to_export:
        # Get source without directives
        source = cell.source_without_directives.strip()
        if not source:
            continue

        # Transform absolute imports to relative imports
        if package_name:
            source = _transform_imports(source, package_name, module_depth)

        # Trigger PRE_CELL_EXPORT hook
        HookRegistry.trigger(
            HookType.PRE_CELL_EXPORT,
            cell=cell,
            notebook=notebook,
        )

        if export_mode == ExportMode.PERCENT:
            # Add cell marker with notebook-specific source reference
            if source_ref:
                parts.append(f"# %% {source_ref} {cell.index}")
            else:
                parts.append(f"# %% {cell.index}")
            parts.append(source)
            parts.append("")
        else:
            # Plain Python mode
            parts.append(source)
            parts.append("")

        # Trigger POST_CELL_EXPORT hook
        HookRegistry.trigger(
            HookType.POST_CELL_EXPORT,
            cell=cell,
            notebook=notebook,
            source=source,
        )

    return "\n".join(parts), list(source_refs_set)


def get_export_targets(notebook: Notebook) -> dict[str, list[int]]:
    """
    Get all export target modules from a notebook.

    Returns a dict mapping module paths to lists of cell indices.
    - Cells with #|export or #|exporti use the notebook's default_exp (key: "" if no default_exp)
    - Cells with #|export_to use the specified module path

    Args:
        notebook: Source notebook

    Returns:
        Dict mapping module paths to cell indices. Key "" means default_exp target.
    """
    targets: dict[str, list[int]] = {}
    default_exp = notebook.default_exp

    for cell in notebook.cells:
        if not cell.is_code:
            continue

        has_export = cell.has_directive("export") or cell.has_directive("exporti")
        has_export_to = cell.has_directive("export_to")

        if not (has_export or has_export_to):
            continue

        if has_export_to:
            # Use the module specified in export_to
            export_to_directive = cell.get_directive("export_to")
            if export_to_directive and export_to_directive.value_parsed:
                target_module = export_to_directive.value_parsed.get("module", "")
                if target_module:
                    if target_module not in targets:
                        targets[target_module] = []
                    targets[target_module].append(cell.index)
        elif has_export:
            # Use default_exp target (empty string as key if no default_exp)
            target_key = default_exp or ""
            if target_key not in targets:
                targets[target_key] = []
            targets[target_key].append(cell.index)

    return targets


def _collect_exported_content(
    notebook: Notebook,
    export_mode: ExportMode,
    source_ref: str | None,
    package_name: str | None = None,
    module_depth: int = 0,
    target_module: str | None = None,
) -> str:
    """
    Collect content from exported cells.

    Cells are sorted by their order value (from #|export, #|exporti, or #|export_to).
    Cells with the same order value maintain their original notebook order (stable sort).

    Args:
        notebook: Source notebook
        export_mode: Export mode (percent or py)
        source_ref: Source reference string for cell markers
        package_name: Package name for import transformation
        module_depth: Depth of module within package (for relative imports)
        target_module: If specified, only include cells targeting this module.
                       Cells with #|export go to default_exp, cells with #|export_to
                       go to their specified module.

    Hooks triggered:
        PRE_CELL_EXPORT: Before each cell export (cell=cell, notebook=notebook)
        POST_CELL_EXPORT: After each cell export (cell=cell, notebook=notebook, source=str)
    """
    parts: list[str] = []
    default_exp = notebook.default_exp

    # First pass: collect cells that should be exported
    cells_to_export: list[Cell] = []

    for cell in notebook.cells:
        if not cell.is_code:
            continue

        # Check for export directives
        has_export = cell.has_directive("export") or cell.has_directive("exporti")
        has_export_to = cell.has_directive("export_to")

        if not (has_export or has_export_to):
            continue

        # Filter by target module if specified
        if target_module is not None:
            cell_target = None
            if has_export_to:
                export_to_directive = cell.get_directive("export_to")
                if export_to_directive and export_to_directive.value_parsed:
                    cell_target = export_to_directive.value_parsed.get("module", "")
            elif has_export:
                cell_target = default_exp

            if cell_target != target_module:
                continue

        cells_to_export.append(cell)

    # Sort cells by order value (stable sort maintains original order for same values)
    cells_to_export.sort(key=lambda c: (_get_cell_order(c), c.index))

    # Second pass: output cells in sorted order
    for cell in cells_to_export:
        # Get source without directives
        source = cell.source_without_directives.strip()
        if not source:
            continue

        # Transform absolute imports to relative imports
        if package_name:
            source = _transform_imports(source, package_name, module_depth)

        # Trigger PRE_CELL_EXPORT hook
        HookRegistry.trigger(
            HookType.PRE_CELL_EXPORT,
            cell=cell,
            notebook=notebook,
        )

        if export_mode == ExportMode.PERCENT:
            # Add cell marker
            if source_ref is not None:
                parts.append(f"# %% {source_ref} {cell.index}")
            else:
                parts.append(f"# %% {cell.index}")
            parts.append(source)
            parts.append("")
        else:
            # Plain Python mode
            parts.append(source)
            parts.append("")

        # Trigger POST_CELL_EXPORT hook
        HookRegistry.trigger(
            HookType.POST_CELL_EXPORT,
            cell=cell,
            notebook=notebook,
            source=source,
        )

    return "\n".join(parts)


def _transform_imports(source: str, package_name: str, module_depth: int) -> str:
    """
    Transform absolute imports to relative imports.

    Converts imports like "from {package_name}.X import Y" to relative imports
    like "from .X import Y" (with appropriate number of dots based on depth).

    Imports inside string literals (f-strings, docstrings, etc.) are left unchanged.

    Args:
        source: Source code to transform
        package_name: The package name to convert (e.g., "my_lib")
        module_depth: Depth of module within package (0 = directly in package)
            - my_lib/core.py has depth 0 -> needs 1 dot
            - my_lib/submodule/utils.py has depth 1 -> needs 2 dots

    Returns:
        Transformed source code
    """
    from nblite.core.directive import _get_string_ranges, _is_position_in_string

    # Number of dots needed: depth + 1
    # depth 0 (my_lib/core.py) -> 1 dot (.)
    # depth 1 (my_lib/submodule/utils.py) -> 2 dots (..)
    num_dots = module_depth + 1
    dots = "." * num_dots

    # Compute string ranges to avoid rewriting imports inside string literals
    string_ranges = _get_string_ranges(source)

    # Precompute line start offsets for offset-to-position conversion
    line_starts = [0]
    for i, ch in enumerate(source):
        if ch == "\n":
            line_starts.append(i + 1)

    def replace_import(match: re.Match) -> str:
        # Skip matches inside string literals (f-strings, docstrings, etc.)
        if string_ranges:
            offset = match.start()
            line_idx = bisect.bisect_right(line_starts, offset) - 1
            col = offset - line_starts[line_idx]
            if _is_position_in_string(line_idx, col, string_ranges):
                return match.group(0)

        indent = match.group(1)
        module_path = match.group(2)
        imports = match.group(3)

        # Check if this import is from our package
        if module_path == package_name:
            # from my_lib import X -> from .. import X (with appropriate dots based on depth)
            return f"{indent}from {dots} import {imports}"
        elif module_path.startswith(f"{package_name}."):
            # from my_lib.core import X -> from .core import X (with appropriate dots)
            rest = module_path[len(package_name) + 1 :]  # Remove "my_lib."
            return f"{indent}from {dots}{rest} import {imports}"
        else:
            # Not our package, leave unchanged
            return match.group(0)

    return FROM_IMPORT_PATTERN.sub(replace_import, source)


def _extract_public_names(content: str) -> list[str]:
    """Extract public names (functions, classes, constants) from content."""
    names: list[str] = []

    # Find functions
    for match in FUNCTION_PATTERN.finditer(content):
        name = match.group(1)
        if not name.startswith("_"):
            names.append(name)

    # Find classes
    for match in CLASS_PATTERN.finditer(content):
        name = match.group(1)
        if not name.startswith("_"):
            names.append(name)

    # Find module-level variables (not starting with _)
    for match in VARIABLE_PATTERN.finditer(content):
        name = match.group(1)
        if not name.startswith("_"):
            names.append(name)

    return names


def _collect_add_to_all_names(notebook: Notebook) -> list[str]:
    """
    Collect names from #|add_to_all directives in a notebook.

    Args:
        notebook: Source notebook

    Returns:
        List of names to add to __all__
    """
    names: list[str] = []
    for cell in notebook.cells:
        if cell.has_directive("add_to_all"):
            directive = cell.get_directive("add_to_all")
            if directive and directive.value_parsed:
                names.extend(directive.value_parsed)
    return names


def _collect_exporti_names(notebook: Notebook) -> set[str]:
    """
    Collect names from #|exporti cells to exclude from __all__.

    Args:
        notebook: Source notebook

    Returns:
        Set of names to exclude from __all__
    """
    names: set[str] = set()
    for cell in notebook.cells:
        if not cell.is_code:
            continue
        if cell.has_directive("exporti"):
            source = cell.source_without_directives
            # Extract names from this cell
            for match in FUNCTION_PATTERN.finditer(source):
                names.add(match.group(1))
            for match in CLASS_PATTERN.finditer(source):
                names.add(match.group(1))
            for match in VARIABLE_PATTERN.finditer(source):
                names.add(match.group(1))
    return names
