"""Code execution engine for Nhandu."""

from __future__ import annotations

import ast
import contextlib
import io
import os
import re
import sys
import traceback
from collections.abc import Callable, Generator
from pathlib import Path
from typing import Any

from nhandu.models import (
    CodeBlock,
    Document,
    ExecutedDocument,
    MarkdownBlock,
)

# Type alias for progress callback
ProgressCallback = Callable[[int, int, CodeBlock], None]


@contextlib.contextmanager
def _script_environment(source_path: Path | None) -> Generator[None, None, None]:
    """
    Context manager to set up Python script environment variables.

    Sets up sys.path and sys.argv to match behavior of running a Python script,
    then restores original values on exit.

    @param source_path: Path to the source document, or None for stdin/in-memory.
    @yield: None
    """
    # Save original state
    original_path = sys.path.copy()
    original_argv = sys.argv.copy()

    try:
        # Set up sys.argv[0] to match script path
        if source_path:
            sys.argv = [str(source_path.absolute()), *sys.argv[1:]]
            # Add script directory to sys.path[0] for relative imports
            script_dir = str(source_path.parent.absolute())
            if script_dir not in sys.path:
                sys.path.insert(0, script_dir)
        else:
            # For stdin/in-memory, use current working directory
            sys.argv = [str(Path.cwd() / "<stdin>"), *sys.argv[1:]]
            cwd = str(Path.cwd().absolute())
            if cwd not in sys.path:
                sys.path.insert(0, cwd)

        yield

    finally:
        # Restore original state
        sys.path[:] = original_path
        sys.argv[:] = original_argv


class CodeExecutor:
    """Executes code blocks and captures output."""

    def __init__(self, working_dir: str | None = None) -> None:
        self.working_dir = working_dir
        self.namespace: dict[str, Any] = {}
        self.figure_counter = 0
        self.plot_dpi = 100  # Default DPI for figures
        self.number_format = ".4f"  # Default format for inline floats
        self._plt: Any | None = None
        self._matplotlib_initialized = False
        # Reuse parser for inline code extraction
        from nhandu.parser import PythonLiterateParser

        self._inline_parser = PythonLiterateParser()

    def _init_matplotlib(self) -> None:
        """Initialize matplotlib once (lazy, on first use)."""
        if self._matplotlib_initialized:
            return
        self._matplotlib_initialized = True
        try:
            import matplotlib

            matplotlib.use("Agg")  # Non-interactive backend
            import matplotlib.pyplot as plt

            self._plt = plt
        except ImportError:
            self._plt = None

    def _resolve_working_dir(self, doc: Document) -> str | None:
        """Resolve execution working directory with document-relative paths."""
        candidate = doc.metadata.working_dir or self.working_dir
        if candidate is None:
            return None

        path = Path(candidate)
        if path.is_absolute() or doc.source_path is None:
            return str(path)

        return str((doc.source_path.parent / path).resolve())

    def execute_document(
        self,
        doc: Document,
        on_progress: ProgressCallback | None = None,
    ) -> ExecutedDocument:
        """Execute all code blocks in a document."""
        original_dir = os.getcwd()

        try:
            # Set working directory if specified
            resolved_working_dir = self._resolve_working_dir(doc)
            if resolved_working_dir:
                os.chdir(resolved_working_dir)

            # Determine figure output directory (created lazily on first figure)
            if doc.source_path:
                output_dir = doc.source_path.parent / "figures"
            else:
                output_dir = Path("figures")

            # Reset namespace for fresh execution with script environment
            self.namespace = self._create_initial_namespace(doc.source_path)
            self.figure_counter = 0
            self.plot_dpi = doc.metadata.plot_dpi or 100
            self.number_format = doc.metadata.number_format or ".4f"

            # Initialize matplotlib (once per executor lifetime)
            self._init_matplotlib()

            executed_doc = ExecutedDocument(
                blocks=[],
                metadata=doc.metadata,
                source_path=doc.source_path,
                namespace=self.namespace,
            )

            # Count code blocks for progress reporting
            total_code_blocks = sum(1 for b in doc.blocks if isinstance(b, CodeBlock))
            code_block_index = 0

            # Set up script environment (sys.path, sys.argv)
            with _script_environment(doc.source_path):
                for block in doc.blocks:
                    if isinstance(block, CodeBlock):
                        # Create a copy to avoid mutating the input document
                        executed_block = CodeBlock(
                            content=block.content,
                            hidden=block.hidden,
                            line_number=block.line_number,
                        )
                        self._execute_code_block(
                            executed_block, output_dir, doc.source_path
                        )
                        executed_doc.blocks.append(executed_block)

                        code_block_index += 1
                        if on_progress:
                            on_progress(
                                code_block_index,
                                total_code_blocks,
                                executed_block,
                            )
                    elif isinstance(block, MarkdownBlock):
                        executed_block_md = self._process_markdown_block(block)
                        executed_doc.blocks.append(executed_block_md)
                    else:
                        executed_doc.blocks.append(block)

        finally:
            # Always restore original directory
            os.chdir(original_dir)

        return executed_doc

    @staticmethod
    def _create_initial_namespace(source_path: Path | None) -> dict[str, Any]:
        """
        Create initial namespace with Python script environment.

        Provides standard special variables that would be present when
        running a Python script: __name__, __file__, __doc__, __package__, __builtins__.

        @param source_path: Path to the source document, or None for stdin/in-memory.
        @return: Dictionary with initial namespace.
        """
        # Determine __file__ value
        if source_path:
            file_path = str(source_path.absolute())
        else:
            # For stdin/in-memory, use a sentinel path in current directory
            file_path = str(Path.cwd() / "<stdin>")

        namespace = {
            "__name__": "__main__",
            "__file__": file_path,
            "__doc__": None,
            "__package__": None,
            "__builtins__": __builtins__,
        }
        return namespace

    def _execute_code_block(
        self, block: CodeBlock, output_dir: Path, source_path: Path | None
    ) -> None:
        """Execute a single code block."""
        stdout_buffer = io.StringIO()
        stderr_buffer = io.StringIO()

        plt = self._plt

        # Clear any existing plots
        if plt:
            plt.close("all")

        try:
            with contextlib.redirect_stdout(stdout_buffer):
                with contextlib.redirect_stderr(stderr_buffer):
                    code = block.content.strip()
                    self._execute_code(code)

        except Exception as e:
            # Format error message with better location info
            block.error = self._format_error(e, block, source_path)

        # Always capture stdout (even if there was an error, partial output
        # printed before the exception is still useful)
        output = stdout_buffer.getvalue()
        if output:
            block.output = output.rstrip()

        # Capture any matplotlib figures (only when no error)
        if plt and not block.error:
            fig_nums = plt.get_fignums()
            if fig_nums:
                # Lazily create figures directory on first figure
                output_dir.mkdir(exist_ok=True)
                figures = []
                for fig_num in fig_nums:
                    fig = plt.figure(fig_num)
                    figure_path = output_dir / f"figure_{self.figure_counter}.png"
                    fig.savefig(
                        figure_path,
                        dpi=self.plot_dpi,
                        bbox_inches="tight",
                    )
                    figures.append(figure_path)
                    self.figure_counter += 1
                    plt.close(fig)
                block.figures = figures

    def _execute_code(self, code: str) -> None:
        """
        Execute code with proper expression detection using AST.

        If the last statement is an expression, its result is printed
        (Jupyter-style behavior). Uses ast.parse for robust detection
        instead of fragile heuristics.
        """
        try:
            tree = ast.parse(code)
        except SyntaxError:
            # If AST parsing fails, try exec directly (handles edge cases
            # like IPython magic-like syntax, though unlikely)
            exec(code, self.namespace)
            return

        if tree.body and isinstance(tree.body[-1], ast.Expr):
            # Last statement is an expression — execute preceding
            # statements, then eval the expression for its result
            if len(tree.body) > 1:
                mod = ast.Module(body=tree.body[:-1], type_ignores=[])
                ast.fix_missing_locations(mod)
                exec(compile(mod, "<string>", "exec"), self.namespace)

            # Evaluate last expression
            expr = ast.Expression(body=tree.body[-1].value)
            ast.fix_missing_locations(expr)
            result = eval(compile(expr, "<string>", "eval"), self.namespace)
            if result is not None:
                print(repr(result))
        else:
            # Execute entire block as statements
            exec(compile(tree, "<string>", "exec"), self.namespace)

    @staticmethod
    def _format_error(
        exc: Exception, block: CodeBlock, source_path: Path | None
    ) -> str:
        """Format an exception with source location information.

        Extracts the line number from the traceback and calculates the actual
        source file line number when possible.
        """
        error_type = type(exc).__name__
        error_message = str(exc)

        # Try to extract line number from traceback
        tb = traceback.extract_tb(exc.__traceback__)
        code_line_number = None

        # Look for the line number in the code block execution
        # The traceback will show "<string>" for exec/eval code
        for frame in reversed(tb):
            if frame.filename == "<string>":
                code_line_number = frame.lineno
                break

        # Build the error message
        parts = [f"{error_type}: {error_message}"]

        if code_line_number is not None and block.line_number:
            # Calculate actual source line number
            # block.line_number is where the code block starts in source
            actual_line = block.line_number + code_line_number - 1

            if source_path:
                parts.append(f'  File "{source_path}", line {actual_line}')
            else:
                parts.append(f"  Line {actual_line}")

            # Try to show the offending line
            lines = block.content.strip().split("\n")
            if 0 < code_line_number <= len(lines):
                offending_line = lines[code_line_number - 1].strip()
                parts.append(f"    {offending_line}")
        elif block.line_number:
            # Fallback if we couldn't extract line from traceback
            if source_path:
                parts.append(f'  File "{source_path}", near line {block.line_number}')
            else:
                parts.append(f"  Near line {block.line_number}")

        return "\n".join(parts)

    def _process_markdown_block(self, block: MarkdownBlock) -> MarkdownBlock:
        """Process inline code in markdown blocks via single-pass re.sub."""
        # Pattern matches both <%= expr %> and <% stmt %>
        pattern = self._inline_parser.inline_code_pattern

        if not pattern.search(block.content):
            return block

        def replace_inline(match: re.Match[str]) -> str:
            full = match.group(0)
            code = match.group(1)
            is_expression = full.startswith("<%=")

            try:
                if is_expression:
                    result = eval(code.strip(), self.namespace)
                    if result is None:
                        return ""
                    if isinstance(result, float):
                        return format(result, self.number_format)
                    return str(result)
                else:
                    exec(code.strip(), self.namespace)
                    return ""
            except Exception as e:
                return f"[ERROR: {type(e).__name__}: {e}]"

        text = pattern.sub(replace_inline, block.content)
        return MarkdownBlock(text, block.line_number)


def execute(
    doc: Document,
    working_dir: str | None = None,
    on_progress: ProgressCallback | None = None,
) -> ExecutedDocument:
    """Execute code blocks in a document."""
    executor = CodeExecutor(working_dir)
    return executor.execute_document(doc, on_progress=on_progress)
