from typing import Any, Dict, List, Optional
import sys
import io
import traceback
import subprocess
from contextlib import redirect_stdout, redirect_stderr

from .config import WORKSPACE_ROOT
from .canvas import parse_canvas_object


# =============================================================================
# JUPYTER-LIKE CODE EXECUTION TOOLS
# =============================================================================

class NotebookState:
    """
    Maintains persistent state for Jupyter-like code execution.

    This class manages:
    - An ordered list of code cells (the "script")
    - A persistent namespace for variable state across cells
    - Execution history and outputs
    - Canvas items generated during cell execution
    """

    def __init__(self):
        self._cells: List[Dict[str, Any]] = []
        self._namespace: Dict[str, Any] = {}
        self._execution_count: int = 0
        self._ipython_shell = None
        self._canvas_items: List[Dict[str, Any]] = []  # Collected canvas items
        self._initialize_namespace()

    def _initialize_namespace(self):
        """Initialize the namespace with common imports and utilities."""
        # Pre-populate with commonly used modules
        init_code = """
import sys
import os
import json
from pathlib import Path

# Data science essentials (imported if available)
try:
    import pandas as pd
except ImportError:
    pass

try:
    import numpy as np
except ImportError:
    pass

try:
    import matplotlib
    matplotlib.use('Agg')  # Non-interactive backend
    import matplotlib.pyplot as plt
except ImportError:
    pass

try:
    import plotly.express as px
    import plotly.graph_objects as go
except ImportError:
    pass
"""
        # Execute initialization silently
        try:
            exec(init_code, self._namespace)
        except Exception:
            pass  # Ignore import errors for optional packages

        # Inject add_to_canvas function that captures items
        def _add_to_canvas_wrapper(content: Any) -> Dict[str, Any]:
            """Add content to the canvas for visualization.

            Supports: DataFrames, matplotlib figures, plotly figures,
            PIL images, and markdown strings.
            """
            try:
                parsed = parse_canvas_object(content, workspace_root=WORKSPACE_ROOT)
                self._canvas_items.append(parsed)
                return parsed
            except Exception as e:
                error_result = {
                    "type": "error",
                    "data": f"Failed to add to canvas: {str(e)}",
                    "error": str(e)
                }
                self._canvas_items.append(error_result)
                return error_result

        self._namespace["add_to_canvas"] = _add_to_canvas_wrapper

    def _get_ipython(self):
        """Get or create an IPython InteractiveShell for enhanced execution."""
        if self._ipython_shell is None:
            try:
                from IPython.core.interactiveshell import InteractiveShell
                self._ipython_shell = InteractiveShell.instance()
                # Share the namespace
                self._ipython_shell.user_ns = self._namespace
            except ImportError:
                # IPython not available, will use exec() fallback
                pass
        return self._ipython_shell

    @property
    def cells(self) -> List[Dict[str, Any]]:
        """Return a copy of all cells."""
        return [cell.copy() for cell in self._cells]

    @property
    def namespace(self) -> Dict[str, Any]:
        """Return the current namespace (variable state)."""
        return self._namespace

    def get_cell(self, cell_index: int) -> Optional[Dict[str, Any]]:
        """Get a cell by index."""
        if 0 <= cell_index < len(self._cells):
            return self._cells[cell_index].copy()
        return None

    def add_cell(self, code: str, cell_type: str = "code") -> Dict[str, Any]:
        """Add a new cell to the end of the script."""
        cell = {
            "index": len(self._cells),
            "type": cell_type,
            "source": code,
            "execution_count": None,
            "outputs": [],
            "status": "pending"
        }
        self._cells.append(cell)
        return cell.copy()

    def insert_cell(self, index: int, code: str, cell_type: str = "code") -> Dict[str, Any]:
        """Insert a cell at a specific index."""
        if index < 0:
            index = 0
        if index > len(self._cells):
            index = len(self._cells)

        cell = {
            "index": index,
            "type": cell_type,
            "source": code,
            "execution_count": None,
            "outputs": [],
            "status": "pending"
        }
        self._cells.insert(index, cell)

        # Update indices for subsequent cells
        for i in range(index + 1, len(self._cells)):
            self._cells[i]["index"] = i

        return cell.copy()

    def modify_cell(self, cell_index: int, new_code: str) -> Dict[str, Any]:
        """Modify the code in an existing cell."""
        if not (0 <= cell_index < len(self._cells)):
            return {
                "error": f"Cell index {cell_index} out of range. Valid range: 0-{len(self._cells) - 1}"
            }

        self._cells[cell_index]["source"] = new_code
        self._cells[cell_index]["status"] = "modified"
        self._cells[cell_index]["outputs"] = []  # Clear previous outputs

        return self._cells[cell_index].copy()

    def delete_cell(self, cell_index: int) -> Dict[str, Any]:
        """Delete a cell by index."""
        if not (0 <= cell_index < len(self._cells)):
            return {
                "error": f"Cell index {cell_index} out of range. Valid range: 0-{len(self._cells) - 1}"
            }

        deleted_cell = self._cells.pop(cell_index)

        # Update indices for subsequent cells
        for i in range(cell_index, len(self._cells)):
            self._cells[i]["index"] = i

        return {"deleted": deleted_cell, "remaining_cells": len(self._cells)}

    def execute_cell(self, cell_index: int) -> Dict[str, Any]:
        """Execute a single cell and capture its output."""
        if not (0 <= cell_index < len(self._cells)):
            return {
                "error": f"Cell index {cell_index} out of range. Valid range: 0-{len(self._cells) - 1}"
            }

        cell = self._cells[cell_index]

        if cell["type"] != "code":
            return {
                "index": cell_index,
                "type": cell["type"],
                "source": cell["source"],
                "output": "(markdown cell - not executed)",
                "status": "skipped"
            }

        self._execution_count += 1
        cell["execution_count"] = self._execution_count

        # Track canvas items added during this cell's execution
        canvas_count_before = len(self._canvas_items)

        # Capture stdout and stderr
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()

        result = {
            "index": cell_index,
            "execution_count": self._execution_count,
            "source": cell["source"],
            "stdout": "",
            "stderr": "",
            "result": None,
            "error": None,
            "status": "success",
            "canvas_items": []  # Canvas items added during execution
        }

        try:
            # Try IPython first for better execution handling
            ipython = self._get_ipython()

            if ipython is not None:
                # Use IPython's run_cell for magic commands support
                with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                    exec_result = ipython.run_cell(cell["source"], store_history=True)

                result["stdout"] = stdout_capture.getvalue()
                result["stderr"] = stderr_capture.getvalue()

                if exec_result.success:
                    if exec_result.result is not None:
                        result["result"] = repr(exec_result.result)
                else:
                    if exec_result.error_in_exec:
                        result["error"] = str(exec_result.error_in_exec)
                        result["status"] = "error"
                    elif exec_result.error_before_exec:
                        result["error"] = str(exec_result.error_before_exec)
                        result["status"] = "error"
            else:
                # Fallback to exec() if IPython is not available
                with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                    # Compile to check for expression vs statement
                    code = cell["source"].strip()

                    # Try to evaluate as expression first (to get return value)
                    try:
                        # Check if it's a simple expression
                        compiled = compile(code, "<cell>", "eval")
                        exec_result = eval(compiled, self._namespace)
                        if exec_result is not None:
                            result["result"] = repr(exec_result)
                    except SyntaxError:
                        # It's a statement, execute it
                        exec(code, self._namespace)

                result["stdout"] = stdout_capture.getvalue()
                result["stderr"] = stderr_capture.getvalue()

        except Exception:
            result["error"] = traceback.format_exc()
            result["status"] = "error"
            result["stdout"] = stdout_capture.getvalue()
            result["stderr"] = stderr_capture.getvalue()

        # Capture any canvas items added during this cell's execution
        canvas_items_added = self._canvas_items[canvas_count_before:]
        result["canvas_items"] = canvas_items_added

        # Store outputs in cell
        cell["outputs"] = [result]
        cell["status"] = result["status"]

        return result

    def execute_all(self) -> List[Dict[str, Any]]:
        """Execute all cells in order."""
        results = []
        for i in range(len(self._cells)):
            results.append(self.execute_cell(i))
        return results

    def get_script(self) -> str:
        """Get all code cells concatenated as a single script."""
        code_cells = [cell["source"] for cell in self._cells if cell["type"] == "code"]
        return "\n\n".join(code_cells)

    def get_variables(self) -> Dict[str, str]:
        """Get a summary of user-defined variables in the namespace."""
        # Filter out modules, builtins, and private variables
        user_vars = {}
        for name, value in self._namespace.items():
            if name.startswith("_"):
                continue
            if isinstance(value, type(sys)):  # Skip modules
                continue
            if callable(value) and hasattr(value, "__module__"):
                # Skip imported functions
                if value.__module__ != "__main__" and value.__module__ not in [None, "builtins"]:
                    continue
            try:
                # Get a short repr
                value_repr = repr(value)
                if len(value_repr) > 100:
                    value_repr = value_repr[:97] + "..."
                user_vars[name] = f"{type(value).__name__}: {value_repr}"
            except Exception:
                user_vars[name] = f"{type(value).__name__}: <unable to repr>"
        return user_vars

    def get_canvas_items(self) -> List[Dict[str, Any]]:
        """Get all canvas items collected during execution."""
        return self._canvas_items.copy()

    def clear_canvas_items(self) -> Dict[str, Any]:
        """Clear collected canvas items."""
        count = len(self._canvas_items)
        self._canvas_items = []
        return {"cleared": count}

    def reset(self):
        """Reset the notebook state (clear all cells and namespace)."""
        self._cells = []
        self._namespace = {}
        self._execution_count = 0
        self._canvas_items = []
        self._initialize_namespace()
        return {"status": "reset", "message": "Notebook state cleared"}


# Global notebook state instance
_notebook_state = NotebookState()


def create_cell(code: str, cell_type: str = "code") -> Dict[str, Any]:
    """
    Create a new code or markdown cell and add it to the end of the script.

    This simulates creating a new cell in a Jupyter notebook. The cell is added
    but not executed - use execute_cell() to run it.

    Args:
        code: The Python code or markdown content for the cell
        cell_type: Either "code" or "markdown" (default: "code")

    Returns:
        Dictionary with cell information including:
        - index: The cell's position in the notebook
        - type: The cell type
        - source: The cell's code/content
        - status: "pending" (not yet executed)

    Examples:
        # Create a code cell
        create_cell("x = 42\\nprint(f'x = {x}')")

        # Create a markdown cell
        create_cell("## Analysis Results", cell_type="markdown")
    """
    return _notebook_state.add_cell(code, cell_type)


def insert_cell(index: int, code: str, cell_type: str = "code") -> Dict[str, Any]:
    """
    Insert a new cell at a specific position in the script.

    This is useful when you need to add code between existing cells,
    such as adding a missing import or intermediate calculation.

    Args:
        index: Position to insert the cell (0-based). Cells after this
               position will be shifted down.
        code: The Python code or markdown content
        cell_type: Either "code" or "markdown" (default: "code")

    Returns:
        Dictionary with cell information including index and status

    Examples:
        # Insert an import at the beginning
        insert_cell(0, "import pandas as pd")

        # Insert a cell between cells 2 and 3
        insert_cell(3, "intermediate_result = process(data)")
    """
    return _notebook_state.insert_cell(index, code, cell_type)


def modify_cell(cell_index: int, new_code: str) -> Dict[str, Any]:
    """
    Modify the code in an existing cell.

    Use this to fix errors, update logic, or refine code in a cell.
    The cell's outputs are cleared and status set to "modified".
    You'll need to re-execute the cell to see the new results.

    Args:
        cell_index: The index of the cell to modify (0-based)
        new_code: The new code to replace the existing code

    Returns:
        Dictionary with updated cell information, or error if index invalid

    Examples:
        # Fix a typo in cell 2
        modify_cell(2, "result = data.groupby('category').mean()")

        # Update a calculation
        modify_cell(0, "threshold = 0.95  # Updated from 0.9")
    """
    return _notebook_state.modify_cell(cell_index, new_code)


def delete_cell(cell_index: int) -> Dict[str, Any]:
    """
    Delete a cell from the script.

    Removes the cell at the specified index. Subsequent cells will have
    their indices updated. Note: This does NOT undo any side effects
    from executing the deleted cell (variables remain in namespace).

    Args:
        cell_index: The index of the cell to delete (0-based)

    Returns:
        Dictionary with deleted cell info and remaining cell count

    Examples:
        # Remove cell 3
        delete_cell(3)
    """
    return _notebook_state.delete_cell(cell_index)


def execute_cell(cell_index: int) -> Dict[str, Any]:
    """
    Execute a single cell and return its output.

    Runs the code in the specified cell within the persistent namespace.
    Variables created or modified will be available to subsequent cells.
    Captures stdout, stderr, and the cell's return value (if any).

    Args:
        cell_index: The index of the cell to execute (0-based)

    Returns:
        Dictionary containing:
        - index: Cell index
        - execution_count: Global execution counter
        - source: The executed code
        - stdout: Captured print() output
        - stderr: Captured error output
        - result: Return value of the last expression (if any)
        - error: Error traceback (if execution failed)
        - status: "success" or "error"

    Examples:
        # Execute the first cell
        execute_cell(0)

        # Execute and check for errors
        result = execute_cell(2)
        if result["status"] == "error":
            print(result["error"])
    """
    return _notebook_state.execute_cell(cell_index)


def execute_all_cells() -> List[Dict[str, Any]]:
    """
    Execute all cells in the script in order.

    Runs each cell sequentially from the beginning. Useful for
    re-running the entire notebook after modifications.

    Returns:
        List of execution results for each cell

    Examples:
        # Run entire notebook
        results = execute_all_cells()
        errors = [r for r in results if r.get("status") == "error"]
    """
    return _notebook_state.execute_all()


def get_script() -> Dict[str, Any]:
    """
    Get the complete script and current state.

    Returns all cells, the concatenated code, and current variable state.
    Useful for reviewing the notebook or exporting the code.

    Returns:
        Dictionary containing:
        - cells: List of all cells with their content and outputs
        - script: All code cells concatenated as a single script
        - variables: Summary of user-defined variables
        - cell_count: Total number of cells

    Examples:
        # Review current state
        state = get_script()
        print(f"Notebook has {state['cell_count']} cells")
        print(state['script'])
    """
    return {
        "cells": _notebook_state.cells,
        "script": _notebook_state.get_script(),
        "variables": _notebook_state.get_variables(),
        "cell_count": len(_notebook_state.cells)
    }


def get_variables() -> Dict[str, str]:
    """
    Get a summary of all user-defined variables in the namespace.

    Returns variable names with their types and values (truncated if long).
    Useful for understanding what data is available for use in new cells.

    Returns:
        Dictionary mapping variable names to "type: value" strings

    Examples:
        # Check available variables
        vars = get_variables()
        for name, info in vars.items():
            print(f"{name}: {info}")
    """
    return _notebook_state.get_variables()


def reset_notebook() -> Dict[str, Any]:
    """
    Reset the notebook state completely.

    Clears all cells and resets the namespace to its initial state.
    Use with caution - this cannot be undone.

    Returns:
        Dictionary confirming the reset

    Examples:
        # Start fresh
        reset_notebook()
    """
    return _notebook_state.reset()


def get_notebook_canvas_items() -> List[Dict[str, Any]]:
    """
    Get all canvas items generated during notebook cell execution.

    When code in cells calls add_to_canvas(), the items are collected here.
    Use this to retrieve visualizations generated by executed code.

    Returns:
        List of canvas item dictionaries with type and data

    Examples:
        # After executing cells that created charts
        items = get_notebook_canvas_items()
        for item in items:
            print(f"Type: {item['type']}")
    """
    return _notebook_state.get_canvas_items()


def clear_notebook_canvas_items() -> Dict[str, Any]:
    """
    Clear all canvas items collected from notebook execution.

    Returns:
        Dictionary with count of cleared items
    """
    return _notebook_state.clear_canvas_items()


# =============================================================================
# CANVAS TOOLS
# =============================================================================

def add_to_canvas(content: Any) -> Dict[str, Any]:
    """Add an item to the canvas for visualization. Canvas is like a note-taking tool where
    you can store charts, dataframes, images, and markdown text for the user to see.

    Args:
        content: Can be a pandas DataFrame, matplotlib Figure, plotly Figure,
                PIL Image, dictionary (for Plotly JSON), or string (for Markdown)
        workspace_root: Path to the workspace root directory

    Returns:
        Dictionary with the parsed canvas object

    Examples:
        # Add a DataFrame
        import pandas as pd
        df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
        add_to_canvas(df)

        # Add a Matplotlib chart
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots()
        ax.plot([1, 2, 3], [1, 4, 9])
        add_to_canvas(fig)

        # Add Markdown text
        add_to_canvas("## Key Findings\\n- Point 1\\n- Point 2")
    """
    try:
        # Parse the content into canvas format
        parsed = parse_canvas_object(content, workspace_root=WORKSPACE_ROOT)
        # Return the parsed object (deepagents will handle the JSON serialization)
        return parsed
    except Exception as e:
        return {
            "type": "error",
            "data": f"Failed to add to canvas: {str(e)}",
            "error": str(e)
        }


# =============================================================================
# BASH TOOL
# =============================================================================

def bash(command: str, timeout: int = 60) -> Dict[str, Any]:
    """Execute a bash command and return the output.

    Runs the command in the workspace directory. Use this for file operations,
    git commands, installing packages, or any shell operations.

    Args:
        command: The bash command to execute
        timeout: Maximum time in seconds to wait for the command (default: 60)

    Returns:
        Dictionary containing:
        - stdout: Standard output from the command
        - stderr: Standard error output
        - return_code: Exit code (0 typically means success)
        - status: "success" or "error"

    Examples:
        # List files
        bash("ls -la")

        # Check git status
        bash("git status")

        # Install a package
        bash("pip install pandas")

        # Run a script
        bash("python script.py")
    """
    try:
        result = subprocess.run(
            command,
            shell=True,
            cwd=str(WORKSPACE_ROOT),
            capture_output=True,
            text=True,
            timeout=timeout
        )

        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "return_code": result.returncode,
            "status": "success" if result.returncode == 0 else "error"
        }

    except subprocess.TimeoutExpired:
        return {
            "stdout": "",
            "stderr": f"Command timed out after {timeout} seconds",
            "return_code": -1,
            "status": "error"
        }
    except Exception as e:
        return {
            "stdout": "",
            "stderr": str(e),
            "return_code": -1,
            "status": "error"
        }
