streaming_code_block.py

"""
Streaming Code Block for Open Interpreter
Properly updates in place for file creation
"""

import re
import os
from rich.console import Console, Group
from rich.panel import Panel
from rich.syntax import Syntax
from rich.table import Table
from rich.text import Text
from rich.box import MINIMAL, ROUNDED
from rich.align import Align
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn

# Track file creation state globally
file_tracker = {
    'active': False,
    'filepath': None,
    'lines_shown': 0,
    'total_lines': 0,
    'completed_files': [],
    'last_block': None
}

def streaming_refresh(self, cursor=True):
    """Enhanced refresh that streams file creation properly"""
    
    # Skip if no code
    if not self.code and not self.output:
        return
    
    # Detect file creation patterns
    is_file_creation = False
    filepath = None
    
    if 'cat <<' in self.code or 'cat >' in self.code:
        # Extract filepath
        filepath_match = re.search(r'>\s*([^\s<>]+)', self.code)
        if filepath_match:
            filepath = filepath_match.group(1).strip()
            is_file_creation = True
    
    # Handle file creation specially
    if is_file_creation and filepath:
        global file_tracker
        
        # Check if this is a new file or continuation
        if not file_tracker['active'] or file_tracker['filepath'] != filepath:
            # New file creation starting
            file_tracker['active'] = True
            file_tracker['filepath'] = filepath
            file_tracker['lines_shown'] = 0
            file_tracker['last_block'] = self
            
            # Detect language from extension
            ext = os.path.splitext(filepath)[1]
            lang_map = {
                '.py': 'python', '.js': 'javascript', '.html': 'html',
                '.css': 'css', '.json': 'json', '.md': 'markdown',
                '.sh': 'bash', '.yml': 'yaml', '.yaml': 'yaml'
            }
            language = lang_map.get(ext, 'text')
        else:
            language = self.language or 'text'
        
        # Count total lines in current code
        code_lines = self.code.split('\n')
        total_lines = len(code_lines)
        
        # Check if file creation is complete
        is_complete = any(marker in self.code for marker in ['EOF', 'EOL']) or self.output
        
        if is_complete:
            # File creation complete - show summary
            size = len(self.code.encode('utf-8'))
            
            # Store the full code for potential expansion
            self.full_code = self.code
            self.collapsed = True
            
            # Create compact summary
            summary_text = Text()
            summary_text.append("✓ ", style="bold green")
            summary_text.append(filepath, style="cyan")
            summary_text.append(f"\n  {total_lines} lines • {size} bytes", style="dim")
            
            # Create expandable panel with full code in a scrollable syntax block
            # The full code is stored but hidden initially
            full_syntax = Syntax(
                self.code,
                language,
                theme="monokai",
                line_numbers=True,
                word_wrap=False
            )
            
            # Create summary panel
            summary_panel = Panel(
                summary_text,
                title="[bold green]Created[/bold green]",
                box=ROUNDED,
                expand=False,
                style="green",
                padding=(0, 1)
            )
            
            # Store reference to full content for later expansion
            if not hasattr(self, '_stored_files'):
                self._stored_files = {}
            self._stored_files[filepath] = {
                'code': self.code,
                'language': language,
                'lines': total_lines,
                'size': size
            }
            
            # Update display with summary
            self.live.update(summary_panel)
            self.live.refresh()
            
            # Reset tracker
            file_tracker['active'] = False
            file_tracker['filepath'] = None
            file_tracker['last_block'] = None
            file_tracker['completed_files'].append(filepath)
            
            # Show output if any
            if self.output:
                output_panel = Panel(
                    self.output,
                    title="[bold]Output[/bold]",
                    style="yellow",
                    box=MINIMAL
                )
                self.live.update(Group(summary_panel, output_panel))
                self.live.refresh()
        else:
            # File creation in progress - show streaming view
            
            # Show only last 10 lines for context
            display_lines = code_lines[-10:] if len(code_lines) > 10 else code_lines
            display_code = '\n'.join(display_lines)
            
            # Create syntax highlighted code
            syntax = Syntax(
                display_code,
                language,
                theme="monokai",
                line_numbers=True,
                start_line=max(1, total_lines - 9)
            )
            
            # Create progress indicator
            progress_text = Text()
            progress_text.append("📝 Writing: ", style="bold cyan")
            progress_text.append(filepath, style="cyan")
            progress_text.append(f" ({total_lines} lines)", style="dim")
            
            # Add spinning indicator
            if total_lines % 3 == 0:
                indicator = "⠋"
            elif total_lines % 3 == 1:
                indicator = "⠙"
            else:
                indicator = "⠹"
            progress_text.append(f" {indicator}", style="cyan")
            
            # Create panel with current content
            panel = Panel(
                syntax,
                title=progress_text,
                border_style="cyan",
                box=ROUNDED,
                expand=True
            )
            
            # Update display in place
            self.live.update(panel)
            self.live.refresh()
            
            # Track lines shown
            file_tracker['lines_shown'] = total_lines
            file_tracker['total_lines'] = total_lines
        
        return  # Skip normal display
    
    # For non-file creation, use normal display but more compact
    # Original code display logic (simplified)
    code = self.code
    
    # Create a table for code display
    code_table = Table(
        show_header=False,
        show_footer=False,
        box=None,
        padding=0,
        expand=True
    )
    code_table.add_column()
    
    # Add syntax highlighting
    if self.language:
        # Show only last 20 lines if code is long
        code_lines = code.split('\n')
        if len(code_lines) > 20:
            display_code = '\n'.join(code_lines[-20:])
            start_line = len(code_lines) - 19
        else:
            display_code = code
            start_line = 1
            
        syntax = Syntax(
            display_code,
            self.language.lower(),
            theme="monokai",
            line_numbers=True,
            start_line=start_line,
            word_wrap=True
        )
        code_table.add_row(syntax)
    else:
        code_table.add_row(code)
    
    # Add output if present
    if self.output:
        # Truncate output if too long
        output_lines = self.output.split('\n')
        if len(output_lines) > 10:
            display_output = '\n'.join(output_lines[:10]) + f"\n... ({len(output_lines) - 10} more lines)"
        else:
            display_output = self.output
            
        output_panel = Panel(
            display_output,
            title="[bold]Output[/bold]",
            box=MINIMAL,
            style="yellow"
        )
        code_table.add_row(output_panel)
    
    # Create main panel
    title = f"[bold]{self.language}[/bold]" if self.language else "[bold]Code[/bold]"
    if cursor and self.active_line is not None:
        title += f" [dim](line {self.active_line})[/dim]"
    
    panel = Panel(
        code_table,
        title=title,
        box=ROUNDED,
        expand=True
    )
    
    # Update display
    self.live.update(panel)
    self.live.refresh()

# Export the enhanced refresh function
__all__ = ['streaming_refresh', 'file_tracker']