from __future__ import annotations

import asyncio
import io
import traceback as tb

from abc import abstractmethod
from collections.abc import Iterable, Iterator
from contextlib import nullcontext
from functools import partial
from types import FunctionType
from typing import Any, Self, final

import panel as pn
import param

from panel.io import hold
from panel.io.cache import is_equal
from panel.layout.base import (
    Column, ListLike, NamedListLike, Row,
)
from panel.pane import Markdown
from panel.viewable import Viewable, Viewer
from panel_material_ui import (
    Accordion, BreakpointSwitcher, Button, Card, ChatFeed, ChatMessage,
    Container, Dialog, Divider, FileDownload, IconButton, Progress, Select,
    SpeedDial, TextAreaInput, TextInput, Typography,
)

from ..views.base import Panel, View
from .actor import (
    Actor, ContextProvider, NullStep, TContext,
)
from .config import MissingContextError
from .context import (
    LWW, ContextError, ValidationIssue, collect_task_outputs,
    input_dependency_keys, merge_contexts, validate_task_inputs,
    validate_taskgroup_exclusions,
)
from .editors import LumenEditor
from .export import (
    format_output, make_md_cell, make_preamble, write_notebook,
)
from .llm import Llm, Message
from .tools import FunctionTool, Tool
from .utils import (
    extract_block_source, get_block_names, wrap_logfire_on_method,
)


class Task(Viewer):
    """
    A `Task` defines a single unit of work that can be executed and rendered.
    """

    abort_on_error = param.Boolean(default=True, doc="""
        If True, the report will abort if an error occurs.""")

    context = param.Dict(default={}, allow_refs=True, doc="""
        The context to pass to the task.""")

    history = param.List(doc="""
        Conversation history to include as context for the task.""")

    instruction = param.String(default="", doc="""
        The instruction to give to the task.""")

    interface = param.ClassSelector(class_=ChatFeed, doc="""
        The chat interface to use for the task.""")

    llm = param.ClassSelector(class_=Llm, doc="""
        The LLM to use for the task.""")

    out_context = param.Dict(default={}, doc="""
        The context generated by the task.""")

    parent = param.Parameter(doc="""
        The parent task group.""")

    running = param.Boolean(doc="""
        Whether the task is currently running.""")

    status = param.Selector(objects=["idle", "running", "success", "error"], default="idle", doc="""
        The current status of the task.""")

    steps_layout = param.ClassSelector(default=None, class_=(ListLike, NamedListLike), allow_None=True, doc="""
        The layout progress updates will be streamed to.""")

    title = param.String(doc="""
        The title of the task.""")

    views = param.List(allow_refs=True, doc="""
        The generated viewable outputs of the task.""")

    _header = param.List()

    def __init_subclass__(cls, **kwargs):
        """
        Apply wrap_logfire to all the subclasses' execute automatically
        """
        super().__init_subclass__(**kwargs)
        wrap_logfire_on_method(cls, "execute")

    def __init__(self, **params):
        super().__init__(**params)
        self._prepared = False
        self._init_view()
        self._null_step = NullStep()
        self._title = Typography(f"{'#'*self.level} {self.title}", margin=(10, 10, 0, 10))

    @param.depends("title", watch=True)
    def _update_title(self):
        self._title = f"{'#'*self.level} {self.title}"

    def __repr__(self):
        params = []
        if self.instruction:
            params.append(f"instruction='{self.instruction}'")
        if self.title:
            params.append(f"title='{self.title}'")
        return f"{self.__class__.__name__}({', '.join(params)})"

    def _init_view(self):
        self._view = self._container = Column(
            sizing_mode='stretch_width', styles={'min-height': 'unset'}, height_policy='fit'
        )

    def _add_step(self, title: str = "", **kwargs):
        """Private contextmanager for adding steps to the interface.

        If self.interface is None, returns a nullcontext that captures calls.
        Otherwise, returns the interface's add_step contextmanager.
        """
        if self.steps_layout is not None and 'steps_layout' not in kwargs:
            kwargs['steps_layout'] = self.steps_layout
        return nullcontext(self._null_step) if self.interface is None else self.interface.add_step(title=title, **kwargs)

    def _populate_view(self):
        self._view[:] = self._header

    def _render_controls(self):
        return [
            TextInput.from_param(
                self.param.title, sizing_mode="stretch_width", margin=(10, 0)
            ),
            TextAreaInput.from_param(
                self.param.instruction, sizing_mode="stretch_width", margin=(10, 0)
            ),
        ]

    def _render_message_history(self, context: TContext) -> list[Message]:
        messages = list(self.history)
        if not self.instruction:
            return messages

        user_msg = None
        for msg in messages[::-1]:
            if msg.get("role") == "user":
                user_msg = msg
                break
        if not user_msg:
            messages.append({"role": "user", "content": self.instruction})
        elif self.instruction not in user_msg.get("content"):
            user_msg["content"] = f'{user_msg["content"]}\n\nInstruction: {self.instruction}'
        return messages

    def _render_output(self, out):
        if isinstance(out, str):
            return Typography(out, margin=(20, 10))
        elif isinstance(out, ChatMessage):
            return Typography(out.object, margin=(20, 10))
        elif isinstance(out, (Viewable, View, LumenEditor)):
            return out

    def __panel__(self):
        return self._container

    @abstractmethod
    async def _execute(self, context: TContext, **kwargs) -> tuple[list[Any], TContext]:
        raise NotImplementedError(f"{self.__class__.__name__} does not implement the execute.")

    def cleanup(self):
        """
        Cleanup when the task is not longer needed.

        Distinct from reset which only resets the outputs.
        """
        self.reset()

    def editor(self, show_title: bool = True) -> Viewable:
        """
        Returns the editor for the task.

        Arguments
        ----------
        show_title: bool
            Whether to show the title of the task.

        Returns
        -------
        The editor for the task.
        """
        return Column(*self._render_controls())

    async def execute(self, context: TContext | None = None, **kwargs) -> tuple[list[Any], TContext]:
        """
        Executes the task.

        Arguments
        ----------
        **kwargs: dict
            Additional keyword arguments to pass to the task.

        Returns
        -------
        The outputs of the task.
        """
        context = dict(self.context or {}, **(context or {}))
        self.param.update(status="running", running=True)
        await asyncio.sleep(0.1)  # Allow _running watcher to insert progress bar before execution starts
        try:
            if not self._prepared:
                await self.prepare(context)
            views, out_context = await self._execute(context, **kwargs)
        except Exception:
            self.status = "error"
            raise
        finally:
            if self.status != "error":
                self.status = "success"
            self.running = False
        self.out_context = out_context
        return views, out_context

    async def prepare(self, context: TContext):
        self._prepared = True

    def reset(self):
        """Resets the view, removing generated outputs."""
        self.status = "idle"
        self._view[:] = []
        self.out_context.clear()


class TaskGroup(Task):

    _current = param.Integer(default=0)

    _tasks = param.List(item_type=Task)

    level = 3

    __abstract = True

    def __init__(self, *tasks, **params):
        if not tasks:
            tasks = params.pop('tasks', [])
        else:
            tasks = list(tasks)
        self._task_watchers = {}
        _tasks = []
        current = 0
        for i, task in enumerate(tasks):
            if isinstance(task, FunctionType):
                task = FunctionTool(task)
            if isinstance(task, Actor) and not issubclass(Actor, self.param._tasks.item_type):
                task = ActorTask(task)
            if isinstance(task, Task):
                task.parent = self
                self._task_watchers[task] = task.param.watch(
                    self._sync_context, 'out_context'
                )
            if current == i and task.status == "success":
                current = i+1
            _tasks.append(task)
        if "status" not in params and current == len(tasks):
            params["status"] = "success"
        super().__init__(_tasks=_tasks, _current=current, **params)
        self._watchers = {}
        self._init_view()
        self._init_views()
        self._populate_view()

    def _init_views(self):
        views = self.param._header.rx()
        for task in self:
            views += task.param.views.rx()
        self.views = views

    async def _sync_context(self, event: param.parameterized.Event):
        task = event.obj
        if task not in self:
            if task in self._task_watchers:
                watcher = self._task_watchers.pop(task)
                task.param.unwatch(watcher)
            return
        index = self._tasks.index(task)+1
        if index > self._current or self.running:
            # No need to invalidate if the task hasn't been executed
            return
        changed = {k for k, v in event.new.items() if k in event.old and not is_equal(event.old.get(k), v)}
        if changed:
            self.invalidate(changed, start=index)
            root = self
            while root.parent is not None:
                root = root.parent
            if not (root is self and index >= len(self)):
                await root.execute()

    def __len__(self) -> int:
        return len(self._tasks)

    def __iter__(self) -> Iterator[Task]:
        return iter(self._tasks)

    def __getitem__(self, index) -> Task:
        return self._tasks[index]

    def __contains__(self, value: Any):
        return value in self._tasks

    def __repr__(self):
        params = []
        if self.title:
            params.append(f"title='{self.title}'")
        tasks = [f"\n    {task!r}" for task in self]
        return f"{self.__class__.__name__}({', '.join(params)}{''.join(tasks)})"

    def _populate_view(self):
        self._view[:] = self._header + self._tasks

    async def _run_task(
        self, i: int, task: Self | Actor, context: TContext, **kwargs
    ) -> tuple[list[Any], TContext]:
        if task.status == "success":
            return task.views, task.out_context
        await asyncio.sleep(0.01)
        messages = self._render_message_history(context)
        subcontext = self._get_context(i, context, task)
        task_history = messages + task.history
        with task.param.update(
            interface=self.interface,
            llm=task.llm or self.llm,
            steps_layout=self.steps_layout,
            running=True,
            history=task_history
        ):
            outputs, out_context = await task.execute(subcontext, **kwargs)
        if isinstance(task, ContextProvider):
            unprovided = [
                p for p in task.output_schema.__required_keys__
                if p not in out_context
            ]
            if unprovided:
                raise RuntimeError(f"{task.__class.__name__} failed to provide declared context.")
        return outputs, out_context

    def _render_tasks(self) -> list[tuple[str, Viewable]]:
        tasks = []
        for task in self:
            task_type = type(task).__name__
            tasks.append((
                f'{task_type}: {task.title}',
                task.editor(show_title=False)
            ))
        return tasks

    def _get_context(
        self, i: int, context: TContext | None, task: Task | ContextProvider
    ) -> TContext | None:
        contexts = ([self.context] if self.context else []) + ([context] if context else [])
        contexts += [task.out_context for task in self[:i]]
        if isinstance(task, ContextProvider):
            return merge_contexts(task.input_schema, contexts)
        elif isinstance(task, TaskGroup):
            # Accumulate contexts requested by subtasks and merge based
            subcontexts = []
            for subtask in task:
                subtask_context = self._get_context(i, context, subtask)
                if subtask_context is not None:
                    subcontexts.append(subtask_context)
            return merge_contexts(LWW, subcontexts)
        else:
            raise TypeError("Abstract Task does not implement _get_context.")

    async def _execute(self, context: TContext, **kwargs):
        """
        Executes the tasks.

        Arguments
        ---------
        context: TContext
            The context given to the task
        **kwargs: dict
            Additional keyword arguments to pass to the tasks.
        """
        if self.title and not self._header:
            self._header = views = [self._title]

        views = list(self._header)
        self.status = "running"
        for i, task in enumerate(self):
            new = []
            try:
                new, new_context = await self._run_task(i, task, context, **kwargs)
            except MissingContextError:
                # Re-raise MissingContextError to allow retry logic at Plan level
                raise
            except Exception as e:
                tb.print_exception(e)
                self.status = "error"
                new_context = {"__error__": str(e), "__error_type__": type(e)}
                task.out_context = new_context
                if self.abort_on_error:
                    if self.parent is not None:
                        raise e
                    break
                else:
                    continue
            else:
                if task.status == "error":
                    self.status = "error"
                    if self.abort_on_error:
                        break
                views += new
            finally:
                self._current = i + (0 if task.status == "error" else 1)
        if self.status != "error":
            self.status = "success"
        contexts = [self.context] if self.context else []
        contexts += [task.out_context for task in self]
        return views, merge_contexts(LWW, contexts)

    def append(self, task: Task | Actor):
        """
        Appends a task to the collection.

        Arguments
        ----------
        task: Task | Actor
            The task to append.
        """
        if not isinstance(task, Task):
            task = ActorTask(task)
        task.parent = self
        self._tasks.append(task)
        self._populate_view()
        self._init_views()

    def remove(self, task: Task | list[Task]):
        tasks = task if isinstance(task, list) else [task]
        for t in tasks:
            self._tasks.remove(t)
        self._populate_view()
        self._init_views()
        self._current = min(self._current, len(self))

    def cleanup(self):
        for task, watcher in self._task_watchers.items():
            task.param.unwatch(watcher)
        super().cleanup()

    def editor(self, show_title=True):
        """
        Returns the editor for the tasks.

        Arguments
        ----------
        show_title: bool
            Whether to show the title of the tasks.

        Returns
        -------
        The editor for the tasks.
        """
        return Column(
            *self._render_controls(),
            *((
                Divider(sizing_mode="stretch_width", margin=(10, 0)),
                Accordion(
                    *self._render_tasks(), margin=(10, 0, 10, 0), title_variant="h4",
                    toggle=True, sizing_mode="stretch_width"
                ),
              ) if self._tasks else ()
            )
        )

    def insert(self, index, task: Task | Actor):
        """
        Inserts a task at the given index.

        Arguments
        ----------
        index: int
            The index to insert the task at.
        task: Task | Actor
            The task to insert.
        """
        if not isinstance(task, Task):
            task = ActorTask(task)
        task.parent = self
        self._tasks.insert(index, task)
        self._populate_view()
        self._init_views()

    def invalidate(
        self, keys: Iterable[str], start: int = 0, propagate: bool = True
    ) -> tuple[list[Task], set[str]]:
        """
        Invalidates tasks and propagates context dependencies within a TaskGroup.

        Parameters
        ----------
        keys : Iterable[str]
            A set of context keys that have been modified or invalidated by the user.
            These represent outputs whose dependent tasks should be re-evaluated.
        start : int
            Index of the task to start invalidating.
        propagate: bool
            Whether to propagate the invalidation to the parent task group.

        Returns
        -------
        tuple[bool, set[str]]
            A tuple ``(invalidated, keys)`` where:
            - ``invalidated`` is ``True`` if any tasks in this group (or nested groups)
                were marked invalid due to overlapping input dependencies.
            - ``keys`` is the full, cumulative set of invalidated context keys after
                propagating dependencies through all affected tasks.

        Notes
        -----
        - If a task's input keys intersect with the provided invalidation keys,
        that task is considered stale and will be rerun.
        - Nested TaskGroups are traversed recursively, and their invalidations
        propagate upward to the parent group.
        """
        if propagate and self.parent is not None and self in self.parent:
            parent_idx = self.parent._tasks.index(self)
            if start >= len(self):
                # If the last task is being invalidated,
                # only subsequent tasks on the parent
                # have to be invalidated
                parent_idx += 1
            self.parent.invalidate(keys, start=parent_idx)
            return

        if start == 0:
            self._header = []

        keys = set(keys)
        invalidated = []
        for i, task in enumerate(self):
            if i < start:
                continue
            subtask_invalidated, subtask_keys = task.invalidate(keys, propagate=False)
            if subtask_invalidated:
                invalidated.append(i)
                keys |= subtask_keys

        if invalidated:
            self.status = "idle"
            self._current = max(min(invalidated), 0)
        return [self[i] for i in invalidated], keys

    def merge(self, other: TaskGroup):
        """
        Merges another task group into the current task group.

        Parameters
        ---------
        other: TaskGroup
            The other task group to merge with.

        Returns
        -------
        self: TaskGroup
            The current task group.
        """
        for task in other:
            other.parent = self
            self._tasks.append(task)
        self._view[:] = list(self._view) + list(other._view)
        self.context = merge_contexts(LWW, [self.context, other.context])
        self.out_context = merge_contexts(LWW, [self.context, self.out_context])
        self.history += [h for h in other.history if h not in self.history]
        self._init_views()
        return self

    async def prepare(self, context: TContext):
        for task in self:
            await task.prepare(context)
        self._prepared = True

    def reset(self, start: int = 0):
        """
        Resets the view, removing generated outputs.
        """
        self.status = "idle"
        self._current = start
        with hold():
            self._header = []
            self._view.clear()
            for task in self[start:]:
                if isinstance(task, Task):
                    task.reset()
            self._populate_view()

    def to_notebook(self):
        """
        Returns the notebook representation of the tasks.
        """
        if len(self) and not self.status == "success":
            raise RuntimeError(
                "Report has not been executed, run report before exporting to_notebook."
            )
        cells, extensions = [], ['tabulator']
        for out in self.views:
            ext = None
            if isinstance(out, Typography):
                level = int(out.variant[1:]) if out.variant and out.variant.startswith('h') else 0
                prefix = f"{'#'*level} " if level else ''
                cell = make_md_cell(f"{prefix}{out.object}")
            elif isinstance(out, Markdown):
                cell = make_md_cell(out.object)
            elif isinstance(out, LumenEditor):
                cell, ext = format_output(out)
            elif isinstance(out, Viewable):
                cell, ext = format_output(Panel(object=out))
            cells.append(cell)
            if ext and ext not in extensions:
                extensions.append(ext)
        cells = make_preamble("", extensions=extensions) + cells
        return write_notebook(cells)

    def validate(
        self,
        context: TContext | None = None,
        available_types: dict[str, Any] | None = None,
        path: str | None = None,
        raise_on_error: bool = True
    ):
        """
        Validate the task group and its subtasks.

        Parameters
        ----------
        context : TContext | None, optional
            The context to validate against, by default None
        available_types : dict[str, Any] | None, optional
            Dictionary of available types for validation, by default None
        path : str | None, optional
            Path to the current task group for error reporting, by default None
        raise_on_error : bool, optional
            Whether to raise an error if validation issues are found, by default True

        Returns
        -------
        tuple[list[ValidationIssue], dict[str, Any]]
            A tuple containing:
            - List of validation issues found
            - Dictionary of output types from all tasks
        """
        cur_path = path or self.name
        issues: list[ValidationIssue] = validate_taskgroup_exclusions(self, path=cur_path)
        value_ctx = dict((self.context or {}), **(context or {}))
        types_out: dict[str, Any] = dict(available_types or {})
        for idx, t in enumerate(self):
            subpath = f"{cur_path}[{idx}] -> {t.name}"
            sub_issues, sub_types = t.validate(
                value_ctx,
                available_types=types_out,
                path=subpath,
            )
            issues.extend(sub_issues)
            types_out.update(sub_types)
        if path is None and issues and raise_on_error:
            raise ContextError(issues)
        return issues, types_out


class Section(TaskGroup):
    """
    A `Section` is a `TaskGroup` representing a sequence of related tasks.
    """

    level = 2

    def __init__(self, *tasks, **params):
        self._watchers = {}
        self._placeholder = None
        super().__init__(*tasks, **params)
        self._update_placeholder()

    def __repr__(self) -> str:
        params = []
        if self.title:
            params.append(f"title='{self.title}'")
        tasks = [f"\n    {task!r}" for task in self]
        return f"{self.__class__.__name__}({', '.join(params)}{''.join(tasks)})"

    def _render_controls(self):
        return [
            TextInput.from_param(
                self.param.title, sizing_mode="stretch_width", margin=(10, 0)
            )
        ]

    def _init_view(self):
        self._dialog = Dialog(
            self.editor(),
            show_close_button=True,
            close_on_click=True,
            title=f'Section Configuration: {self.title}',
        )
        self._settings = IconButton(
            icon='settings', on_click=self._open_settings, size="small", color="default"
        )
        self._placeholder = Typography(
            "", variant="body2", margin=(10, 10),
            sx={"color": "text.secondary"}
        )
        self._view = Column(sizing_mode='stretch_width', styles={'min-height': 'unset'}, height_policy='fit')
        self._container = Column(
            Row(
                self._settings,
                styles={'position': 'absolute', 'top': '-57.5px', 'right': '20px'}
            ),
            self._dialog,
            self._view,
            sizing_mode='stretch_width',
            styles={'min-height': 'unset'},
            height_policy='fit'
        )

    def _open_settings(self, event):
        self._dialog.open = True

    @param.depends('status', '_tasks', 'views', watch=True)
    def _update_placeholder(self):
        """Update placeholder text based on execution status."""
        if self._placeholder is None:
            return
        has_outputs = self.status != "idle" or bool(self.views)
        if not has_outputs:
            n = len(self._tasks)
            task_word = "task" if n == 1 else "tasks"
            self._placeholder.object = f"{n} {task_word} ready"
            self._placeholder.visible = True
        else:
            self._placeholder.visible = False

    def _populate_view(self):
        self._view[:] = self._header + [self._placeholder] + list(self._tasks)

    async def _run_task(self, i: int, task: Task | Actor, context: TContext | None, **kwargs) -> list[Any]:
        if context is not None:
            instructions = "\n".join(
                f"{i+1}. {task.instruction}" if hasattr(task, 'instruction') else f"{i+1}. <no instruction>"
                for i, task in enumerate(self)
            )
            context['reasoning'] = f"{self.title}\n\n{instructions}"
        return await super()._run_task(i, task, context, **kwargs)

    @param.depends('running', watch=True)
    async def _running(self):
        await asyncio.sleep(0.05)
        if not self.running:
            return
        # Start with indeterminate, switch to determinate after first task
        loader = Progress(
            variant='indeterminate',
            sizing_mode='stretch_width',
            margin=(0, 0, 10, 0)
        )
        self._view.insert(0, loader)

        # Switch to determinate mode once first task completes
        while self._current == 0 and self.running:
            await asyncio.sleep(0.05)

        if self.running and self._current > 0:
            loader.variant = 'determinate'
            progress_value = self.param._current.rx.pipe(lambda x: (x / len(self) * 100) if len(self) > 0 else 0)
            loader.value = progress_value

        while self.running:
            await asyncio.sleep(0.05)
        self._view.remove(loader)


class Report(TaskGroup):
    """
    A `Report` is a `TaskGroup` consisting of a sequence of `Section` objects.

    The `Report` UI renders buttons to execute, clear, export and configure the
    Report.
    """

    auto_execute = param.Boolean(default=False, doc="""
        If True, automatically execute the report on initialization.""")

    _tasks = param.List(item_type=Section)

    level = 1

    def __init__(self, *tasks, **params):
        if not tasks:
            tasks = params.pop('tasks', [])
        else:
            tasks = list(tasks)
        super().__init__(*tasks, **params)
        pn.state.execute(partial(self.prepare, self.context))
        if self.auto_execute:
            pn.state.execute(self.execute)

    def _init_view(self):
        self._header_title = Typography(
            self.param.title, variant="h1", margin=(0, 0, 0, 10)
        )
        self._view = Accordion(
            sizing_mode="stretch_width", min_height=0, margin=(0, 5, 15, 5), toggle=True,
            sx={
                "& .MuiAccordionDetails-root": {
                    "p": "0 calc(2 * var(--mui-spacing)) 1em !important",
                    "overflowX": "auto"
                }
            }
        )
        self._run = IconButton(
            icon="play_arrow", on_click=self._execute_event, margin=0, size="large",
            description="Execute Report", loading=self.param.running,
        )
        self._clear = IconButton(
            icon="clear", on_click=lambda _: self.reset(), margin=0, size="large",
            description="Clear outputs", visible=False
        )
        self._collapse = IconButton(
            styles={"margin-left": "auto"}, on_click=self._expand_all, icon="unfold_less",
            size="large", color="default", margin=(0, 0, 10, 0),
            description="Collapse/Expand Sections", visible=False
        )
        self._settings = IconButton(
            icon="settings", on_click=self._open_settings, size="large", color="default",
            margin=0, description="Configure Report", visible=False
        )
        self._export = IconButton(
            icon="get_app", on_click=self._trigger_download, size="large", color="default",
            margin=0, description="Export Report to .ipynb", visible=False
        )
        self._download = FileDownload(
            callback=self._notebook_export, filename=f"{self.title or 'Report'}.ipynb", visible=False
        )
        self._dialog = Dialog(
            TextInput.from_param(self.param.title, margin=(10, 0, 0, 0), sizing_mode="stretch_width"),
            show_close_button=True,
            close_on_click=True,
            title=f"Report Settings: {self.title}",
        )
        self._menu = Row(
            self._header_title,
            self._run,
            self._clear,
            self._collapse,
            self._export,
            self._settings,
            self._download,
            sizing_mode="stretch_width"
        )
        self._dial = SpeedDial(
            items=[
                {"label": "Execute Report", "icon": "play_arrow"},
                {"label": "Clear Report", "icon": "clear"},
                {"label": "Export Report to Notebook", "icon": "get_app"},
                {"label": "Configure Report", "icon": "settings"}
            ],
            color="default",
            direction="down",
            icon="more_vert",
            size="small",
            styles={"margin": "10px 10px 10px auto"},
            on_click=self._trigger_event,
            sx={"& .MuiFab-root": {"boxShadow": "var(--mui-shadows-2)"}}
        )
        self._collapsed_menu = Row(
            self._header_title,
            self._dial,
            sizing_mode="stretch_width"
        )
        self._switcher = BreakpointSwitcher(
            media_query='(min-width: 800px)',
            small=self._collapsed_menu,
            large=self._menu,
            sizing_mode="stretch_width"
        )
        self._container = Column(
            self._view,
            self._dialog,
            margin=(0, 0, 0, 5),
            sizing_mode="stretch_both",
            height_policy='fit',
        )
        self._update_run_state()
        self._update_icon_visibility()

    async def _trigger_event(self, item: dict):
        icon = item["icon"]
        if icon == "play_arrow":
            await self._execute_event()
        elif icon == "clear":
            self.reset()
        elif icon == "get_app":
            self._trigger_download()
        elif icon == "settings":
            self._open_settings()

    def _trigger_download(self, event=None):
        self._download.filename = f"{self.title or 'Report'}.ipynb"
        self._download.transfer()

    @param.depends('_current', '_tasks', watch=True)
    def _update_run_state(self):
        # Play button is always enabled - users can re-run after changing settings
        pass

    @param.depends('status', 'views', watch=True)
    def _update_icon_visibility(self):
        """Show/hide icons based on whether report has outputs."""
        has_outputs = self.status in ("success", "error") or bool(self.views)
        self._clear.visible = has_outputs
        self._collapse.visible = has_outputs
        self._export.visible = has_outputs
        self._settings.visible = has_outputs
        # Only animate play button when no outputs
        if has_outputs:
            self._run.sx = {}
        else:
            self._run.sx = {
                "animation": "pulse 1s ease-in-out 2",
                "@keyframes pulse": {
                    "0%, 100%": {"transform": "translateX(0)"},
                    "50%": {"transform": "translateX(7px)"},
                },
            }

    async def _execute_event(self, event=None):
        await asyncio.sleep(0.01)  # yield the event loop to allow button loading state to update
        await self.execute()

    async def _notebook_export(self):
        if len(self) and self.status != "success":
            await self.execute()
        return io.StringIO(self.to_notebook())

    def _expand_all(self, event=None):
        if self._collapse.icon == "unfold_less":
            self._view.active = []
            self._collapse.icon = "expand"
        else:
            self._view.active = list(range(len(self._view)))
            self._collapse.icon = "unfold_less"

    def _open_settings(self, event=None):
        self._dialog.open = True

    def _populate_view(self):
        self._view[:] = objects = [(task.title, task) for task in self]
        has_outputs = self.status in ("success", "error") or bool(self.views)
        if has_outputs:
            self._view.active = list(range(len(objects)))
        else:
            self._view.active = list(range(len(objects))) if objects else []

    async def _run_task(self, i: int, task: Section, context: TContext, **kwargs):
        self._view.active = self._view.active + [i]
        outputs = await super()._run_task(i, task, context, **kwargs)
        return outputs

    def __panel__(self):
        return Container(
            self._switcher,
            self._container,
            align="center",
            sizing_mode="stretch_both",
            stylesheets=[":host > div { overflow-y: auto; }"],
            sx={
                "minWidth": "320px",
                "width": "100%",
                "height": "auto",
                ".mui-light &": { "background-color": "var(--mui-palette-grey-100)"},
                ".mui-dark &": { "background-color": "var(--mui-palette-grey-900)"}
            }
        )


class ExecutableTask(Task, ContextProvider):

    level = 3

    def reset(self):
        super().reset()
        self.views = []

    def invalidate(
        self, keys: Iterable[str], start: int = 0, propagate: bool = True
    ) -> tuple[list[Task], set[str]]:
        """
        Invalidates tasks and propagates context dependencies within a TaskGroup.

        Parameters
        ----------
        keys : Iterable[str]
            A set of context keys that have been modified or invalidated by the user.
            These represent outputs whose dependent tasks should be re-evaluated.
        start : int
            Index of the task to start invalidating.
        propagate: bool
            Whether to propagate the invalidation to the parent task group.

        Returns
        -------
        tuple[bool, set[str]]
            A tuple ``(invalidated, keys)`` where:
            - ``invalidated`` is ``True`` if any tasks in this group (or nested groups)
                were marked invalid due to overlapping input dependencies.
            - ``keys`` is the full, cumulative set of invalidated context keys after
                propagating dependencies through all affected tasks.

        Notes
        -----
        - If a task's input keys intersect with the provided invalidation keys,
        that task is considered stale and will be rerun.
        - Nested TaskGroups are traversed recursively, and their invalidations
        propagate upward to the parent group.
        """
        if propagate and self.parent is not None and self in self.parent:
            parent_idx = self.parent._tasks.index(self)
            if start >= len(self):
                # If the last task is being invalidated,
                # only subsequent tasks on the parent
                # have to be invalidated
                parent_idx += 1
            self.parent.invalidate(keys, start=parent_idx)
            return

        keys = set(keys)
        deps = input_dependency_keys(self.input_schema)
        # If no declared dependencies (empty input_schema), consider it as
        # dependent on all context keys - always invalidate when anything changes.
        # This ensures agents like ChatAgent that have no explicit input_schema
        # are re-run when upstream context changes.
        should_invalidate = (not deps and keys) or (deps & keys)
        if should_invalidate:
            if start == 0:
                # Do no reset if invalidation was called
                # solely to notify parent
                self.reset()
            keys |= set(self.output_schema.__annotations__)
            return [self], keys
        return [], keys

    def validate(
        self,
        context: TContext | None = None,
        available_types: dict[str, Any] | None = None,
        path: str | None = None,
        raise_on_error: bool = True
    ):
        """
        Validate the task group and its subtasks.

        Parameters
        ----------
        context : TContext | None, optional
            The context to validate against, by default None
        available_types : dict[str, Any] | None, optional
            Dictionary of available types for validation, by default None
        path : str | None, optional
            Path to the current task group for error reporting, by default None
        raise_on_error : bool, optional
            Whether to raise an error if validation issues are found, by default True

        Returns
        -------
        tuple[list[ValidationIssue], dict[str, Any]]
            A tuple containing:
            - List of validation issues found
            - Dictionary of output types from all tasks
        """
        subpath = f"{path} -> {self.name}" if path else self.name
        value_ctx = dict((self.context or {}), **(context or {}))
        issues = validate_task_inputs(
            self, value_ctx, available_types, subpath
        )
        types_out = collect_task_outputs(self)
        if path is None and issues and raise_on_error:
            raise ContextError(issues)
        return issues, types_out


class ActorTask(ExecutableTask):
    """
    An `ActorTask` wraps an `Actor` in a way that allows persisting
    the message history, instructions, and context as well as making
    it possible to track its outputs.
    """

    actor = param.ClassSelector(class_=Actor)

    level = 3

    def __init__(self, actor: Actor, **params):
        views = params.get("views", None)
        out_context = params.get("out_context", None)
        super().__init__(actor=actor, **params)
        if views and out_context:
            self._add_outputs(views, out_context)

    @property
    def input_schema(self):
        return self.actor.input_schema

    @property
    def output_schema(self):
        return self.actor.output_schema

    async def _update_spec(self, event: param.parameterized.Event):
        self.out_context = dict(self.out_context, **(await event.obj.render_context()))

    def _add_outputs(self, views: list, context: TContext):
        # Handle Tool specific behaviors
        if isinstance(self.actor, Tool):
            # Handle View/Viewable results regardless of agent type
            rendered = []
            for o in views:
                if not isinstance(o, (View, Viewable)):
                    continue
                if isinstance(o, Viewable):
                    pipeline = None if context is None else context.get('pipeline')
                    o = Panel(object=o, pipeline=pipeline)
                o = LumenEditor(
                    component=o, title=self.title
                )
                message_kwargs = dict(value=o, user=self.actor.name)
                if self.interface:
                    self.interface.stream(**message_kwargs)
                rendered.append(o)
            views = rendered

        # Attach retry controls
        for view in views:
            if not isinstance(view, LumenEditor):
                continue
            view.param.watch(self._update_spec, "spec")

        self.views = self._header + views
        rendered = []
        for view in views:
            out = self._render_output(view)
            if out is not None:
                rendered.append(out)
        self._view.extend(rendered)

    async def _execute(self, context: TContext, **kwargs) -> tuple[list[Any], TContext]:
        views = []
        if self.title and not self._header:
            title = Typography(f"{'#'*self.level} {self.title}", margin=(10, 10, 0, 10))
            views.append(title)
            self.views = self._header = [title]

        messages = self._render_message_history(context)
        subcontext = merge_contexts(self.input_schema, [self.context, context])
        with self.actor.param.update(
            interface=self.interface,
            llm=self.actor.llm or self.llm,
            steps_layout=self.steps_layout
        ):
            try:
                kwargs["step_title"] = self.title
                outputs, out_context = await self.actor.respond(messages, subcontext, **kwargs)
            except MissingContextError:
                # Re-raise MissingContextError to allow retry logic at Plan level
                raise
            except Exception as e:
                self.status = "error"
                raise e
            self._add_outputs(outputs, out_context)
        self.status = "success"
        unprovided = [
            p for p in self.output_schema.__required_keys__
            if p not in out_context
        ]
        if unprovided:
            raise RuntimeError(f"{self.actor.__class__.__name__} failed to provide declared context: {', '.join(unprovided)}.")
        contexts = [self.context, out_context] if self.context else [out_context]
        return outputs, merge_contexts(LWW, contexts)

    def _actor_prompt(self, actor: Actor):
        prompt = Select(
            options=list(actor.prompts), label="Select prompt to modify",
            margin=(10, 0), sizing_mode="stretch_width"
        )
        template = pn.rx(actor.prompts)[prompt]['template']
        block = Select(
            options=template.rx.pipe(get_block_names), label='Select block to modify',
            margin=(10, 0, 10, 10), sizing_mode="stretch_width"
        )

        def update_prompt(prompt, block, event):
            if prompt not in actor.template_overrides:
                actor.template_overrides[prompt] = {}
            actor.template_overrides[prompt][block] = event.new

        def edit_prompt(event):
            source = extract_block_source(template.rx.value, block.value)
            editor = TextAreaInput(
                label=f"Edit {prompt.value} prompt's {block.value} block",
                value=source, sizing_mode='stretch_width', margin=(10, 0), height=300
            )
            cancel = IconButton(
                icon="cancel",
                on_click=lambda _: (layout.remove(edit_layout), actor.template_overrides.pop(prompt.value, {}).pop(block.value, None)),
                styles={'position': 'absolute', 'right': '0px', 'zIndex': "9999"}
            )
            edit_layout = Column(cancel, editor)
            editor.param.watch(partial(update_prompt, prompt.value, block.value), 'value')
            layout.append(edit_layout)

        edit = Button(
            icon="edit_note", icon_size='2em', on_click=edit_prompt, height=54,
            margin=(10, 0, 10, 10)
        )
        layout = Column(Row(prompt, block, edit))
        return layout

    async def prepare(self, context: TContext):
        await self.actor.prepare(context)
        self._prepared = True

    def editor(self, show_title=True):
        """
        Returns the editor for the tasks.

        Arguments
        ----------
        show_title: bool
            Whether to show the title of the tasks.

        Returns
        -------
        The editor for the tasks.
        """
        task_type = type(self.actor).__name__
        if isinstance(self.actor, FunctionTool):
            title, prompt = f"{task_type}: {self.actor.function.__name__}", ''
        else:
            title, prompt = task_type, self._actor_prompt(self.actor)
        return Column(
            *self._render_controls(),
            Divider(sizing_mode="stretch_width", margin=(10, 0)),
            Card(
                prompt, margin=(10, 0, 10, 0), title=title, title_variant="h4",
                sizing_mode="stretch_width"
            )
        )


class Action(ExecutableTask):
    """
    An `Action` implements an execute method that performs some unit of work
    and optionally generates outputs to be rendered.
    """

    render_outputs = param.Boolean(default=True, doc="""
         Whether the outputs should be rendered.""")

    def show(self, **kwargs):
        """Show the Action wrapped in a minimal Report UI."""
        section_title = self.title or self.__class__.__name__
        report = Report(
            Section(self, title=section_title),
            title=section_title,
            auto_execute=True
        )
        return report.show(**kwargs)

    def servable(self, **kwargs):
        """Make the Action servable wrapped in a minimal Report UI."""
        section_title = self.title or self.__class__.__name__
        report = Report(
            Section(self, title=section_title),
            title=section_title,
            auto_execute=True
        )
        return report.servable(**kwargs)

    @final
    async def execute(self, context: TContext | None = None, **kwargs):
        views, out_context = await super().execute(context, **kwargs)
        if self.render_outputs:
            self._view[:] = [self._render_output(out) for out in views]
        if self.title:
            views = [self._title] + views
        self.views = views
        self.status = "success"
        return views, out_context
