"""
Export pipeline for nblite.

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

from __future__ import annotations

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).

    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
    """
    # 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

    def replace_import(match: re.Match) -> str:
        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
