from typing import Any, Generator
from types import TracebackType
from IPython.core.inputtransformer2 import TransformerManager
from IPython import get_ipython
from IPython.core.ultratb import AutoFormattedTB
from IPython.core.interactiveshell import InteractiveShell
from nbformat.v4.nbbase import new_notebook, new_markdown_cell, new_code_cell, new_output
from nbformat import NotebookNode
from contextlib import redirect_stdout, redirect_stderr
import nbformat
import io
from markdown import Markdown
from bs4 import BeautifulSoup
import os
import sys
from .editor import EditorDialog
from ..key_handler import KeyHandler
from ..form_elements import TextElement
from ..dialog import Dialog, HTMLElement
from ..subdialogs import AppletDialog
from ..framework import backup, beep, load
from .. import debugging
import traceback as tb


def zzz_notebook_exc_writer(self: Any, etype: type[BaseException], value: BaseException, tb: TracebackType, tb_offset: int | None = None) -> None:
    """Legacy notebook exception writer (not currently used)."""
    #this produces all sorts of incomprehensible strings of numbers (possibly timing values) and does not usefully reference cell numbers
    aotb = AutoFormattedTB(mode='Plain', color_scheme='NoColor', tb_offset=tb_offset)
    formatted_tb = aotb.structured_traceback(etype, value, tb)
    print("\n".join(formatted_tb))


def notebook_exc_writer(self: Any, exc_type: type[BaseException], exc: BaseException, exc_tb: TracebackType, tb_offset: int | None = None) -> None: 
    print(f'*** Notebook Traceback for applet = {self.dialog} , most recent first. ***',file=sys.stderr)
    try:
        debugging.exc_writer(exc_type,exc,exc_tb,exc_applet=self.dialog)
    except Exception as e:
        tb.print_exc()


class IShell(InteractiveShell):
    """Custom IPython InteractiveShell for notebook execution.

    Attributes:
        dialog: Reference to parent NotebookDialog.
        cell: Currently executing notebook cell.
    """

    def __init__(self, dialog: 'NotebookDialog', *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.dialog: NotebookDialog = dialog
        self.cell: NotebookNode | None = None

    def set_cell(self, cell: NotebookNode) -> None:
        """Mark cell as currently executing.

        Args:
            cell: Notebook cell to mark.
        """
        #mark this as the cell that is currently executing which will be tested in convert_file_to_cell
        cell.execution_count = -1
        self.cell = cell

    def set_next_input(self, source: str, replace: bool) -> None:
        #print(f'set_next_input with source {source} and replace {replace}')
        if replace:
            self.cell.source = source
        return super().set_next_input(source, replace)


class OutputsDialog(Dialog):
    """Dialog for viewing notebook cell outputs."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        if len(self):
            self.cursor = len(self) - 1

    @KeyHandler.hotkey("Control+Rmenu")
    def return_to_cell(self, modified_name: str) -> bool | None:
        """Alias for escape - return to parent cell.

        Args:
            modified_name: Modified key name.

        Returns:
            Result of cancel().
        """
        return self.cancel(modified_name)


class NotebookDialog(EditorDialog):
    """Jupyter notebook editor dialog.

    Provides full notebook functionality including:
    - Cell navigation (Control+Up/Down)
    - Cell execution (Control+Return)
    - Cell creation (F7 for code, Shift+F7 for markdown)
    - Output viewing (Control+Rmenu)
    - Integrated IPython shell with magic commands

    Attributes:
        notebook: NotebookNode containing all cells and metadata.
        global_context: Global execution context dictionary.
        current_id: ID of currently displayed cell.
        shell: IShell instance for executing code cells.
    """

    def __init__(self, parent: Dialog, name_seed: str, element_type: type[TextElement] = TextElement, filler: Any = [], revert_to: Any = None) -> None:
        self.notebook: NotebookNode | None = None
        self.global_context: dict[str, Any] = {}
        self.current_id: str | None = None
        self.shell: IShell = IShell(self)
        self.shell.set_custom_exc((Exception,),notebook_exc_writer)          
        if filler:
            filename = filler if isinstance(filler, str) else filler.filename 
            self.load_notebook(filename)
        else: 
            #self.notebook = nbformat.reads(sample_content, as_version=4)
            self.notebook = new_notebook()
            #add one markdown cell and one code cell to the new notebook 
            self.add_cell(markdown=True,init=True)
            self.add_cell(markdown=False,init=True)            
            self.set_new_file(template='new_notebook.ipynb')
            #the above sets self.filler and self.name so use them in the super to avoid them getting overridden 
            filler = self.filler
            name_seed=self.name 
        self.current_id = self.notebook.cells[0].id
        super().__init__(parent, name_seed=name_seed, element_type=element_type, filler=filler, revert_to=revert_to)

    def load_notebook(self, filename: str) -> None:
        with open(filename, 'r', encoding='utf-8') as f:
            self.notebook = nbformat.read(f, as_version=4)

    def generate_content(self, filler: Any, element_type: type[TextElement] = TextElement) -> Generator[TextElement | HTMLElement, None, None]:
        if hasattr(self.notebook,'cells') and len(self.notebook.cells):
            if not self.current_cell().cell_type=='markdown': 
                self.say_indents=True
                for l in self.current_cell().source.split('\n'):
                    if l.strip():
                        yield TextElement(self,'blank line',l) 
            else:
                self.say_indents=False
                soup = self.markdown_as_soup(self.current_cell().source)
                #descendents doesn't work but not sure children is right as will not pick up grandchildren - need to check out markdown possibilities and why descendents is not working 
                for crouton  in soup.children :  
                    #navigable strings (rather than elements) have a NoneType name so replace that with an empty string  
                    name = crouton.name
                    if (name is None) or (name.lower()=='p'):
                        name = ""
                    txt = crouton.get_text()
                    if isinstance(txt, str) and txt.strip():
                        yield HTMLElement(self, name, txt, html=crouton)

    def empty_message(self) -> str:
        if not isinstance(self.notebook, nbformat.NotebookNode):
            return 'No notebook loaded'
        elif not hasattr(self.notebook, 'cells'):
            return 'notebook has no cells list'
        else:
            return f'cell {self.cell_index(self.current_cell())} is an empty {self.current_cell().cell_type} cell'

    def convert_pseudo_filename(self, pseudo_file: str) -> str:
        if pseudo_file in self.shell.compile._filename_map:  
            execution_count = self.shell.compile._filename_map[pseudo_file]
            max_count=0
            for i  in range(len(self.notebook.cells)): 
                if hasattr(self.notebook.cells[i],'execution_count') and not (self.notebook.cells[i].execution_count is None): 
                    this_count = self.notebook.cells[i].execution_count 
                    max_count = max(this_count,execution_count)
                    if this_count== execution_count:  
                        return f'cell {i+1}' 
                    elif this_count == -1:
                        #temporarily store the cell number for the running cell which had it's execution count set to -1 in shell.set_cell  
                        running_cell = i+1
            #no match found in any cell for the execution count found in the filename map 
            if execution_count == max_count:
                #no match found for execution_count which is the highest value in the map so must be the running cell for which execution_count only gets set after run_cell returns, so return the running cell number
                return f'cell {running_cell}'
        return f'{pseudo_file} not in shell.compile._filename_map'

    @KeyHandler.hotkey("Control+Up")
    def previous_cell(self, modified_name: str) -> None:
        self.move_cell(-1)

    @KeyHandler.hotkey("Control+Down")
    def next_cell(self, modified_name: str) -> None:
        self.move_cell(1)

    def move_cell(self, increment: int) -> None:
        if 0 <= (self.cell_index(self.current_cell()) + increment) <= (len(self.notebook.cells) - 1):
            #before refilling dialog, copy current contents to current cell as it may have been edited 
            self.contents_to_notebook()
            self.current_id =self.notebook.cells[self.cell_index(self.current_cell())+increment].id 
            self.refill(self.filler)
            self.say_strings(self.cell_name(self.current_cell()))
        else:
            beep()

    def cell_index(self, cell: NotebookNode) -> int:
        return self.notebook.cells.index(cell)

    def is_current_cell(self, cell: NotebookNode) -> bool:
        return self.current_id == cell.id

    def current_cell(self) -> NotebookNode:
        """Return the current cell."""
        #enhance the below to do similar to self.current() which returns an empty element if dialog is empty,
        return [cell for cell in self.notebook.cells if cell.id == self.current_id][0]

    @KeyHandler.hotkey("Control+Rmenu")
    def command_outputs(self, modified_name: str) -> None:
        self.show_outputs(modified_name.startswith('Shift'))

    def show_outputs(self, all: bool = False) -> None:
        outputs=[]
        targets = [self.current_cell()] if not all else self.notebook.cells  
        cells = [c for c in targets if c.cell_type=='code']    
        for cell in cells:
            for o in cell.outputs:  
                outputs.append(f'cell {self.notebook.cells.index(cell)+1} {o.name} {o.output_type}')
                outputs.extend(self.output_content(o))
        OutputsDialog(self, name_seed='output streams', filler=outputs).run(speak='%current')
        self.say_current()

    def output_content(self, o: Any) -> list[str]:
        if hasattr(o, 'text'):
            content = [l for l in o.text.split('\n') if len(l.strip())]
        else:
            content = [f'cell {self.cell_index(self.current_cell())+1} output {o.name} {o.output_type} not supported']
        return content

    @KeyHandler.hotkey("F7")
    def command_add_cell(self, modified_name: str) -> None:
        self.add_cell(markdown=modified_name.startswith('Shift'))

    def add_cell(self, markdown: bool = False, init: bool = False) -> None:
        if (not init) and len(self.notebook.cells):
            current_index = self.cell_index(self.current_cell())
        else:
            init = True 
        if markdown:
            self.notebook.cells.append(new_markdown_cell(f'#markdown cell {len(self.notebook.cells)+1}\n##this is a subheading\nand this is text'))
        else:
            self.notebook.cells.append(new_code_cell(f'#code  cell {len(self.notebook.cells)+1}'))
        if not init:
            self.move_cell(len(self.notebook.cells) - 1 - current_index)

    def contents_to_notebook(self) -> None:
        """Copy dialog contents back to current notebook cell."""
        #need to convert edited HTML back to markdown for markdown cels
        if self.current_cell().cell_type == 'code':
            self.current_cell().source = '\n'.join([e.value for e in self])

    @KeyHandler.hotkey("Control+S")
    def command_save(self, modified_name: str) -> None:
        """Save notebook to file."""
        #copy contents of dialog to current cell as it may have been edited
        self.contents_to_notebook()
        super().command_save(modified_name)

    def write(self, filename: str) -> None:
        with open(filename, 'w', encoding='utf-8') as f:
            nbformat.write(self.notebook, f)

    @KeyHandler.hotkey("Control+T")
    def command_cell_names(self, modified_name: str) -> None:
        self.say_strings(self.cell_name(self.current_cell(), extended=modified_name.startswith('Shift')))

    def cell_name(self, cell: NotebookNode, extended: bool = False) -> str:
        title = f'cell {self.notebook.cells.index(cell)+1} {cell.cell_type}'  
        if extended:
            title += f'  include details of outputs, id, etc'
        return title

    def markdown_as_soup(self, markdown_source: str) -> BeautifulSoup:
        md = Markdown()
        html_content = md.convert(markdown_source)
        soup = BeautifulSoup(html_content, 'html.parser')
        return soup

    def execute_cell(self, cell: NotebookNode, global_context: dict[str, Any]) -> None:
        # Function to execute a code cell and capture its output    
        stdout_io = io.StringIO()
        stderr_io = io.StringIO()
        with redirect_stdout(stdout_io), redirect_stderr(stderr_io):
            try:
                self.shell.set_cell(cell)
                result = self.shell.run_cell(cell.source,store_history=True,cell_id=cell.id)  
                cell.execution_count = result.execution_count 
                error = False  
                if self.is_current_cell(cell):
                    self.refill(self.filler) #reload in case %load magic has changed the cell contents 
            except Exception as e:
                stderr_io.write(f'unexpectedly reached except block after run_cell, not custom exc writer, with exception {e}')            
                exc = e
                error = True
        stdout_output = stdout_io.getvalue()
        stderr_output = stderr_io.getvalue()
        # Clear existing outputs
        cell.outputs = []
        # Update the cell outputs
        if stdout_output:
            cell.outputs.append(new_output(output_type="stream", name="stdout", text=stdout_output))
        if stderr_output or error:
            cell.outputs.append(new_output(output_type="stream", name="stderr", text=stderr_output))
            if error:
                debugging.exc_writer(exc.__class__, exc, exc.__traceback__, exc_applet=self)

    def exception_frame_stack_cell_and_source(self, context: Any) -> Generator[tuple[str, str], None, None]:  # type: ignore[misc]  # Dead code - undefined obj, pseudo_file
        """Pre-pend the exception, frame and stack and line no on to the below yielding filler (not currently used)."""
        # This appears to be incomplete legacy code - obj and pseudo_file are undefined
        obj = self  # Assume obj refers to self
        pseudo_file = ""  # Undefined - would need to be passed as parameter
        if not hasattr(obj.shell, 'compile') or not hasattr(obj.shell.compile, '_filename_map'):
            yield ("Error: shell.compile._filename_map not available", '')
            return
        if pseudo_file not in obj.shell.compile._filename_map:
            yield (f"Error: {pseudo_file} not in filename map", '')
            return
        execution_count = obj.shell.compile._filename_map[pseudo_file]
        for cell in obj.notebook.cells:
            if cell.execution_count == execution_count:
                for l in cell.source.split('\n'):
                    yield (l, '')
                return
        yield (f'could not find cell with execution count = {execution_count}', '')

    @KeyHandler.hotkey("Control+Return")
    def execute(self, modified_name: str) -> None:
        self.contents_to_notebook()
        all =  'Shift' in modified_name
        if all:
            # Execute the entire notebook
            i=0
            for cell in self.notebook.cells:
                if cell.cell_type=='code':
                    self.execute_cell(cell, self.global_context)
                    i+=1
            self.say_strings(f'executed {i} code cells')                    
        else:
            cell = self.current_cell()
            if cell.cell_type!='code':
                self.say_strings(f'cannot execute cell {self.notebook.cells.index(cell)} which is of type {cell.cell_type}')
            else:
                self.execute_cell(cell, self.global_context)
                if  (hasattr(cell,'outputs') and cell.outputs and (len(cell.outputs)==1) and (cell.outputs[0].name=='stdout')):            
                    self.say_strings(self.output_content(cell.outputs[0]))            
                else:
                    self.say_strings('executed current cell')
#end of file