"""
Transaction progress updater that manages displaying execution tree as subtasks.

This module handles displaying processed transaction progress data,
keeping display concerns separate from data processing.
"""
from typing import Any, Dict, List, Optional

from relationalai.tools.cli_controls import TaskProgress, NotebookTaskProgress
from relationalai.tools.txn_progress_tree import (
    ProcessedTask,
    ProgressSummaryItem,
    build_progress_tree,
    get_ivm_relative_update_size,
    get_loop_iteration,
)
from relationalai.util.format import format_duration


def should_include_task(task: ProcessedTask) -> bool:
    """Filter out error tasks that aren't useful for display."""
    return not task.task_name.startswith(":error_pyrel_error_attrs")


def format_task_tree(task: ProcessedTask, base_indent: str = "") -> str:
    """
    Format a task tree as a multi-line string for display.
    
    This is a pure function for easy testing. The base_indent parameter
    accounts for the progress system's prefix (5 spaces = "   ➜ ").
    
    Args:
        task: Root task to format
        base_indent: Base indentation (empty for first line, "     " for subsequent)
        
    Returns:
        Multi-line string representation of the task tree
    """
    lines = _format_task_tree_lines(task, indent_level=0, base_indent=base_indent)
    return "\n".join(lines)


def _format_task_tree_lines(
    task: ProcessedTask,
    indent_level: int,
    base_indent: str
) -> List[str]:
    """
    Recursively build lines for a task tree with proper indentation.
    
    Args:
        task: Task node to format
        indent_level: Current indentation level (0 for root)
        base_indent: Base indentation for the first line of this task
        
    Returns:
        List of formatted lines for this task and its children
    """
    lines = []
    
    # Calculate indentation for this line
    # Each indent level adds 2 spaces
    task_indent = base_indent + ("  " * indent_level)
    
    # Format and add this task
    description = _format_task_description(task)
    lines.append(task_indent + description)
    
    # After the first line, all subsequent lines need full base_indent
    # The progress system adds "   ➜ " (5 chars) before the first line only
    child_base_indent = "     " if not base_indent else base_indent
    
    # Add warnings as children
    for warning_code, message in task.warnings:
        warning_indent = child_base_indent + ("  " * (indent_level + 1))
        lines.append(f"{warning_indent}⚠️  {warning_code}: {message}")
    
    ivm_relative_update_size = get_ivm_relative_update_size(task)
    if ivm_relative_update_size:
        indent = child_base_indent + ("  " * (indent_level + 1))
        lines.append(f"{indent}Updating relation incrementally. Updated rows: {ivm_relative_update_size:.1f}%")

    loop_iteration = get_loop_iteration(task)
    if loop_iteration:
        if task.end_time:
            return [f"Loop {task.task_name}: terminated after {loop_iteration} iterations"]
        else:
            return [f"Loop {task.task_name} running on iteration {loop_iteration}"]

    # Add children
    for child in task.children:
        if should_include_task(child):
            child_lines = _format_task_tree_lines(child, indent_level + 1, child_base_indent)
            lines.extend(child_lines)
    
    return lines


def _format_task_description(task: ProcessedTask) -> str:
    """Format a single task description for display."""
    name = task.task_name

    has_warnings = len(task.warnings) > 0
    if has_warnings:
        warning_count = len(task.warnings)
        warning_text = f" ⚠️  {warning_count} warning" + ("s" if warning_count > 1 else "")
    else:
        warning_text = ""
    
    if task.duration is not None:
        duration_str = format_duration(task.duration)
        return f"{name} ({duration_str}){warning_text}"
    else:
        return f"{name} (Evaluating...){warning_text}"


class TxnProgressUpdater:
    """Manages displaying transaction execution tree as subtasks in a TaskProgress."""
    
    def __init__(self, progress: TaskProgress | NotebookTaskProgress | Any, enabled: bool = False):
        """
        Initialize the updater.
        
        Args:
            progress: The TaskProgress or NotebookTaskProgress instance to update
            enabled: Whether to display internal transaction progress
        """
        self.progress = progress
        self.enabled = enabled
    
    def update(self, internal_txn_progress: Optional[Dict]) -> None:
        """
        Update the progress display with transaction execution tree.
        
        Args:
            internal_txn_progress: The internal transaction progress dict with 'tasks' key
        """
        if not self.enabled or internal_txn_progress is None:
            return
        
        if "tasks" not in internal_txn_progress:
            return
        
        # Process data into immutable tree structure
        task_tree = build_progress_tree(internal_txn_progress["tasks"])
        
        # Display the processed tasks
        self._display_tasks(task_tree)
    
    def _display_tasks(self, task_tree: List[ProcessedTask]) -> None:
        """
        Display task tree by updating the progress display.
        
        Only root tasks are added as sub-tasks. Each root's description
        includes its entire subtree rendered as a multi-line string with
        proper indentation that accounts for the progress system's prefix.
        
        Args:
            task_tree: Immutable tree of processed tasks to display
        """
        current_root_ids = set()
        
        # Process each root task
        for root in task_tree:
            if not should_include_task(root):
                continue
                
            current_root_ids.add(root.task_id)
            
            # Build the full multi-line description for this root and its children
            description = format_task_tree(root, base_indent="")
            
            # Get or create the progress task for this root, using the transaction task_id
            if root.task_id in self.progress.list_sub_tasks():
                self.progress.update_sub_task(root.task_id, description)
            else:
                self.progress.add_sub_task(description, task_id=root.task_id)
            
            # Mark as complete if the root task is finished
            if root.duration is not None:
                completed_tasks = self.progress.get_completed_tasks()
                if root.task_id not in completed_tasks:
                    self.progress.complete_sub_task(root.task_id, record_time=False)
        
        # Remove root tasks that are no longer roots. This happens when a child task's parent appears.
        tasks_to_remove = set(self.progress.list_sub_tasks()) - current_root_ids
        for task_id in tasks_to_remove:
            self.progress.remove_sub_task(task_id, animate=False)


def format_progress_summary(
    summary_items: List[ProgressSummaryItem],
    total_duration: float
) -> str:
    """
    Format progress summary items for display.
    
    Args:
        summary_items: List of summary items to format
        total_duration: Total transaction duration in seconds
        
    Returns:
        Formatted multi-line string showing task names, durations, and percentages
    """
    if not summary_items or total_duration <= 0:
        return ""
    
    sorted_items = sorted(summary_items, key=lambda item: item.wall_clock_duration, reverse=True)

    lines = []
    for item in sorted_items:
        if item.loop_iteration:
            lines.append(f"  Loop {item.task_name}: terminated after {item.loop_iteration} iterations")
        else:
            duration_str = format_duration(item.wall_clock_duration)
            percentage = (item.wall_clock_duration / total_duration) * 100
            percentage_str = f"{percentage:.1f}%"
            lines.append(f"  {item.task_name}: {duration_str} ({percentage_str})")
            if item.ivm_relative_update_size:
                lines.append(f"    Relation incrementally updated. Updated rows: {item.ivm_relative_update_size:.1f}%")
    
    return "\n".join(lines)
