"""
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.cell import Cell
    from nblite.core.notebook import Notebook

__all__ = ["export_function_notebook", "is_function_notebook"]


def _get_func_cell_order(cell: Cell, directive_name: str, default: int) -> int:
    """
    Get the order value for a cell with a specific directive.

    Args:
        cell: The cell to get order from
        directive_name: Name of directive to check (e.g., "top_export", "export")
        default: Default order value if not specified

    Returns:
        The order value
    """
    if cell.has_directive(directive_name):
        directive = cell.get_directive(directive_name)
        if directive and directive.value_parsed:
            return directive.value_parsed.get("order", default)
    return default


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 definition (decorators, signature, docstring)
    func_def = _get_function_signature(notebook)
    if func_def is None:
        # Default signature if not provided
        # Use only the last part of the module path (after final dot) as function name
        module_name = notebook.default_exp or "exported_func"
        if "." in module_name:
            module_name = module_name.rsplit(".", 1)[-1]
        func_def = f"def {module_name}():"

    # Add function definition (decorators, signature, docstring - may be multi-line)
    for line in func_def.split("\n"):
        lines.append(line)

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

    # Collect bottom_export cells (code after function)
    bottom_exports = _collect_bottom_exports(notebook)
    if bottom_exports:
        lines.append("")  # Blank line after function
        for code in bottom_exports:
            if package_name:
                code = _transform_imports(code, package_name, module_depth)
            lines.append(code)
            lines.append("")

    # 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, sorted by order."""
    # Collect cells with their order
    cells_with_order: list[tuple[int, int, str]] = []  # (order, index, code)

    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:
                order = _get_func_cell_order(cell, "top_export", -1000)
                cells_with_order.append((order, cell.index, code))

    # Sort by order, then by original index (stable sort)
    cells_with_order.sort(key=lambda x: (x[0], x[1]))

    return [code for _, _, code in cells_with_order]


def _get_function_signature(notebook: Notebook) -> str | None:
    """Extract function definition from #|set_func_signature directive.

    Returns the entire cell content (decorators, signature, docstring) as-is,
    only stripping trailing `...` or `pass` placeholders.

    Returns:
        The function definition string, or None if not found.
    """
    for cell in notebook.cells:
        if not cell.is_code:
            continue
        if cell.has_directive("set_func_signature"):
            code = cell.source_without_directives.rstrip()

            # Remove trailing ellipsis or pass placeholder
            if code.endswith("..."):
                code = code[:-3].rstrip()
            elif code.endswith("pass"):
                code = code[:-4].rstrip()

            return code if code else None

    return None


def _collect_function_body(notebook: Notebook) -> list[str]:
    """Collect code for function body from exported cells, sorted by order."""
    # Collect cells with their order and processed content
    cells_with_order: list[tuple[int, int, list[str]]] = []  # (order, index, lines)

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

        # Skip non-exported cells (must have export, exporti, or func_return)
        has_export = cell.has_directive("export") or cell.has_directive("exporti")
        has_func_return = cell.has_directive("func_return")

        if not has_export and not has_func_return:
            continue

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

        # Get order value (default 0 for export/exporti)
        if cell.has_directive("export"):
            order = _get_func_cell_order(cell, "export", 0)
        elif cell.has_directive("exporti"):
            order = _get_func_cell_order(cell, "exporti", 0)
        else:
            order = 0  # func_return defaults to 0

        cell_lines: list[str] = []

        # Handle func_return - prepend return to all code
        if has_func_return:
            source = cell.source_without_directives.strip()
            if source:
                cell_lines.append(f"return {source}")
        else:
            # Handle func_return_line - inline directive
            # Process the raw source first, then remove other directives
            processed_lines = _process_func_return_lines(cell.source)
            cell_lines.extend(processed_lines)

        if cell_lines:
            cells_with_order.append((order, cell.index, cell_lines))

    # Sort by order, then by original index (stable sort)
    cells_with_order.sort(key=lambda x: (x[0], x[1]))

    # Flatten all lines
    body_lines: list[str] = []
    for _, _, lines in cells_with_order:
        body_lines.extend(lines)

    return body_lines


def _collect_bottom_exports(notebook: Notebook) -> list[str]:
    """Collect code from cells with #|bottom_export directive, sorted by order."""
    # Collect cells with their order
    cells_with_order: list[tuple[int, int, str]] = []  # (order, index, code)

    for cell in notebook.cells:
        if not cell.is_code:
            continue
        if cell.has_directive("bottom_export"):
            code = cell.source_without_directives.strip()
            if code:
                order = _get_func_cell_order(cell, "bottom_export", 1000)
                cells_with_order.append((order, cell.index, code))

    # Sort by order, then by original index (stable sort)
    cells_with_order.sort(key=lambda x: (x[0], x[1]))

    return [code for _, _, code in cells_with_order]


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
            code = re.sub(r"\s*#\|func_return_line\s*$", "", line)

            # Extract leading whitespace to preserve indentation
            leading_ws = ""
            for char in code:
                if char in " \t":
                    leading_ws += char
                else:
                    break

            code_stripped = code.strip()
            if code_stripped:
                # Inline with code: preserve indentation and prepend return
                lines.append(f"{leading_ws}return {code_stripped}")
            else:
                # Directive on its own line: add bare return with same indentation
                # Get indentation from the original line
                orig_leading_ws = ""
                for char in line:
                    if char in " \t":
                        orig_leading_ws += char
                    else:
                        break
                lines.append(f"{orig_leading_ws}return")
        else:
            lines.append(line)
    return lines
