"""
Function notebook export.

Exports notebooks as Python functions where the notebook body
becomes the function body and directives control structure.
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from nblite.core.notebook import Notebook

__all__ = ["export_function_notebook", "is_function_notebook"]


def is_function_notebook(notebook: Notebook) -> bool:
    """
    Check if a notebook is a function notebook.

    A function notebook has the `#|export_as_func` directive set to true.

    Args:
        notebook: The notebook to check.

    Returns:
        True if the notebook is a function notebook.
    """
    for cell in notebook.cells:
        if cell.has_directive("export_as_func"):
            directive = cell.get_directive("export_as_func")
            if directive and directive.value.strip().lower() == "true":
                return True
    return False


def export_function_notebook(
    notebook: Notebook,
    output_path: Path,
    *,
    include_warning: bool = True,
    package_name: str | None = None,
) -> None:
    """
    Export a function notebook to a Python module.

    The notebook is structured as:
    - `#|top_export`: Code that appears before the function (imports)
    - `#|set_func_signature`: Defines the function signature
    - `#|export`: Code that becomes the function body
    - `#|func_return`: Code that becomes a return statement
    - `#|func_return_line`: Inline directive to make a line a return statement

    Args:
        notebook: The notebook to export.
        output_path: Path to write the module.
        include_warning: Whether to include autogenerated warning.
        package_name: Package name for converting absolute imports to relative imports.

    Raises:
        ValueError: If required directives are missing.
    """
    # Import here to avoid circular imports
    from nblite.export.pipeline import _transform_imports

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Calculate module depth for relative imports (same logic as pipeline.py)
    module_depth = 0
    if package_name:
        # Count parent directories from output_path that are within the package
        parts = output_path.parts
        try:
            pkg_idx = parts.index(package_name)
            # Depth is number of directories between package and file (excluding file)
            module_depth = len(parts) - pkg_idx - 2
        except ValueError:
            pass

    lines: list[str] = []

    # Add autogenerated warning
    if include_warning:
        lines.append("# AUTOGENERATED! DO NOT EDIT!")
        lines.append("")

    # Collect top_export cells (imports, etc.)
    top_exports = _collect_top_exports(notebook)
    for code in top_exports:
        if package_name:
            code = _transform_imports(code, package_name, module_depth)
        lines.append(code)
        lines.append("")

    # Get function signature
    signature = _get_function_signature(notebook)
    if signature is None:
        # Default signature if not provided
        module_name = notebook.default_exp or "exported_func"
        signature = f"def {module_name}():"

    # Add function definition
    lines.append(signature)

    # Collect function body
    body_lines = _collect_function_body(notebook)
    if body_lines:
        for line in body_lines:
            lines.append(f"    {line}")
    else:
        lines.append("    pass")

    # Write output
    content = "\n".join(lines)
    # Clean up extra blank lines
    while "\n\n\n" in content:
        content = content.replace("\n\n\n", "\n\n")

    output_path.write_text(content.strip() + "\n")


def _collect_top_exports(notebook: Notebook) -> list[str]:
    """Collect code from cells with #|top_export directive."""
    exports = []

    for cell in notebook.cells:
        if not cell.is_code:
            continue
        if cell.has_directive("top_export"):
            code = cell.source_without_directives.strip()
            if code:
                exports.append(code)

    return exports


def _get_function_signature(notebook: Notebook) -> str | None:
    """Extract function signature from #|set_func_signature directive."""
    for cell in notebook.cells:
        if not cell.is_code:
            continue
        if cell.has_directive("set_func_signature"):
            code = cell.source_without_directives.strip()
            # Extract the def line
            for line in code.split("\n"):
                line = line.strip()
                if line.startswith("def "):
                    # Remove any ellipsis or pass
                    if line.endswith("..."):
                        line = line[:-3].rstrip()
                    # Ensure it ends with colon
                    if not line.endswith(":"):
                        line += ":"
                    return line
    return None


def _collect_function_body(notebook: Notebook) -> list[str]:
    """Collect code for function body from exported cells."""
    body_lines: list[str] = []

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

        # Skip non-exported cells
        if not cell.has_directive("export") and not cell.has_directive("func_return"):
            continue

        # Skip cells with set_func_signature or top_export
        if cell.has_directive("set_func_signature") or cell.has_directive("top_export"):
            continue

        # Handle func_return - prepend return to all code
        if cell.has_directive("func_return"):
            source = cell.source_without_directives.strip()
            if source:
                body_lines.append(f"return {source}")
            continue

        # Handle func_return_line - inline directive
        # Process the raw source first, then remove other directives
        processed_lines = _process_func_return_lines(cell.source)
        for line in processed_lines:
            body_lines.append(line)

    return body_lines


def _process_func_return_lines(source: str) -> list[str]:
    """Process source code, handling #|func_return_line inline directives."""
    lines = []
    for line in source.split("\n"):
        stripped = line.strip()
        # Skip directive-only lines (lines that start with #| and have no code before)
        if stripped.startswith("#|") and "#|func_return_line" not in stripped:
            continue

        # Check for inline func_return_line directive
        if "#|func_return_line" in line:
            # Remove the directive and prepend return
            code = re.sub(r"\s*#\|func_return_line\s*$", "", line)
            code = code.rstrip()
            if code:
                lines.append(f"return {code}")
        else:
            lines.append(line)
    return lines
