"""
Flow Analysis Widget using Textual DataTable
Hierarchical flow display with automatic formatting and responsive layout
"""

from textual.widgets import DataTable, Static
from textual.containers import Vertical
from textual.reactive import reactive
from typing import TYPE_CHECKING, List, Optional, Tuple
from rich.text import Text

if TYPE_CHECKING:
    from ....analysis.core import EthernetAnalyzer
    from ....models import FlowStats


class FlowAnalysisWidget(Vertical):
    """
    Enhanced Flow Analysis using Textual DataTable
    
    Features:
    - Automatic column sizing and alignment
    - Hierarchical sub-rows for protocol breakdown
    - Rich text formatting with colors
    - Mouse and keyboard navigation
    - Real-time data updates
    """
    
    selected_flow_index = reactive(0)
    
    def __init__(self, analyzer: 'EthernetAnalyzer', **kwargs):
        super().__init__(**kwargs)
        self.analyzer = analyzer
        self.flow_table = None
        self.flows_list = []
        
    def compose(self):
        """Create the widget layout"""
        yield Static("FLOW ANALYSIS", id="flow-title")
        
        # Main flow data table
        flow_table = DataTable(id="flows-table")
        flow_table.cursor_type = "row"
        flow_table.zebra_stripes = True
        
        # Add columns with proper alignment
        flow_table.add_columns(
            "#",           # Flow number (right-aligned)
            "Source",      # IP:port (left-aligned)  
            "Proto",       # Transport protocol (left-aligned)
            "Destination", # IP:port (left-aligned)
            "Extended",    # Extended protocol (left-aligned)
            "Frame Type",  # Frame type (left-aligned)
            "Pkts",        # Packet count (right-aligned)
            "Volume",      # Data volume (right-aligned)
            "Timing",      # Inter-arrival timing (right-aligned)
            "Quality"      # Quality metric (right-aligned)
        )
        
        self.flow_table = flow_table
        yield flow_table
    
    def on_mount(self) -> None:
        """Initialize the widget when mounted"""
        self.refresh_data()
    
    def refresh_data(self) -> None:
        """Refresh the flow data in the table"""
        if not self.flow_table:
            return
        
        # Preserve cursor and scroll positions
        cursor_row = self.flow_table.cursor_row
        cursor_column = self.flow_table.cursor_column
        scroll_x = self.flow_table.scroll_x
        scroll_y = self.flow_table.scroll_y
        selected_row_key = None
        if self.flow_table.rows and cursor_row < len(self.flow_table.rows):
            selected_row_key = list(self.flow_table.rows.keys())[cursor_row]
            
        # Clear existing data
        self.flow_table.clear()
        
        # Get updated flows list
        self.flows_list = self._get_flows_list()
        
        # Populate table with hierarchical data
        for i, flow in enumerate(self.flows_list):
            # Add main flow row
            main_row = self._create_flow_row(i + 1, flow)
            row_key = self.flow_table.add_row(*main_row, key=f"flow_{i}")
            
            # Mark enhanced flows with special styling
            if flow.enhanced_analysis.decoder_type != "Standard":
                # Note: DataTable doesn't have set_row_style, using CSS classes instead
                pass
            
            # Add sub-rows for protocol/frame type breakdown
            protocol_combinations = self._get_protocol_frame_combinations(flow)
            for j, (extended_proto, frame_type, count, percentage) in enumerate(protocol_combinations):
                sub_row = self._create_sub_row(flow, extended_proto, frame_type, count, percentage)
                sub_key = self.flow_table.add_row(*sub_row, key=f"flow_{i}_sub_{j}")
                # Note: DataTable doesn't have set_row_style, using CSS classes instead
        
        # Restore cursor position
        if selected_row_key and selected_row_key in self.flow_table.rows:
            row_index = list(self.flow_table.rows.keys()).index(selected_row_key)
            self.flow_table.move_cursor(row=row_index, column=cursor_column, animate=False)
        elif self.flow_table.row_count > 0:
            # If original selection not found, try to maintain row position
            new_row = min(cursor_row, self.flow_table.row_count - 1)
            self.flow_table.move_cursor(row=new_row, column=cursor_column, animate=False)
        
        # Restore scroll position
        self.flow_table.scroll_to(x=scroll_x, y=scroll_y, animate=False)
    
    def _create_flow_row(self, flow_num: int, flow: 'FlowStats') -> List[Text]:
        """Create main flow row with rich text formatting"""
        
        # Format source with potential truncation
        source = f"{flow.src_ip}:{flow.src_port}"
        source_text = Text(source[:22] + "..." if len(source) > 25 else source)
        
        # Transport protocol
        protocol_text = Text(flow.transport_protocol, style="bold cyan")
        
        # Format destination  
        destination = f"{flow.dst_ip}:{flow.dst_port}"
        dest_text = Text(destination[:22] + "..." if len(destination) > 25 else destination)
        
        # Extended protocol summary
        combinations = self._get_protocol_frame_combinations(flow)
        extended_text = Text(f"{len(combinations)} types", style="yellow")
        
        # Frame type summary
        frame_text = Text("Mixed" if len(flow.frame_types) > 1 else "Single", style="blue")
        
        # Packet count (right-aligned)
        packets_text = Text(str(flow.frame_count), justify="right", style="white")
        
        # Volume with units (right-aligned)
        volume = self._format_bytes(flow.total_bytes)
        volume_text = Text(volume, justify="right", style="magenta")
        
        # Timing (right-aligned)
        if flow.avg_inter_arrival > 0:
            timing_ms = flow.avg_inter_arrival * 1000
            if timing_ms >= 1000:
                timing = f"{timing_ms/1000:.1f}s"
            else:
                timing = f"{timing_ms:.1f}ms"
        else:
            timing = "N/A"
        timing_text = Text(timing, justify="right", style="cyan")
        
        # Quality indicator (right-aligned)
        if flow.enhanced_analysis.decoder_type != "Standard":
            if flow.enhanced_analysis.avg_frame_quality > 0:
                quality = f"{flow.enhanced_analysis.avg_frame_quality:.0f}%"
                quality_style = "bold green"
            else:
                quality = "Enhanced"
                quality_style = "green"
        else:
            outlier_pct = len(flow.outlier_frames) / flow.frame_count * 100 if flow.frame_count > 0 else 0
            if outlier_pct > 5:
                quality = f"{outlier_pct:.0f}% Out"
                quality_style = "red"
            else:
                quality = "Normal"
                quality_style = "green"
        
        quality_text = Text(quality, justify="right", style=quality_style)
        
        return [
            Text(str(flow_num), justify="right"),
            source_text,
            protocol_text, 
            dest_text,
            extended_text,
            frame_text,
            packets_text,
            volume_text,
            timing_text,
            quality_text
        ]
    
    def _create_sub_row(self, flow: 'FlowStats', extended_proto: str, frame_type: str, count: int, percentage: float) -> List[Text]:
        """Create sub-row for protocol/frame type combination"""
        
        # Empty columns for inheritance from parent flow
        empty = Text("")
        
        # Extended protocol
        extended_text = Text(extended_proto if extended_proto != '-' else "", style="dim yellow")
        
        # Frame type
        frame_text = Text(frame_type, style="dim blue")
        
        # Packet count for this combination
        count_text = Text(str(count), justify="right", style="dim white")
        
        # Volume estimation
        volume_bytes = int(flow.total_bytes * (percentage / 100))
        volume = self._format_bytes(volume_bytes)
        volume_text = Text(volume, justify="right", style="dim magenta")
        
        # Timing for this frame type
        if frame_type in flow.frame_types and flow.frame_types[frame_type].avg_inter_arrival > 0:
            timing_ms = flow.frame_types[frame_type].avg_inter_arrival * 1000
            if timing_ms >= 1000:
                timing = f"{timing_ms/1000:.1f}s"
            else:
                timing = f"{timing_ms:.1f}ms"
        else:
            timing = "-"
        timing_text = Text(timing, justify="right", style="dim cyan")
        
        # Percentage as quality
        quality_text = Text(f"{percentage:.1f}%", justify="right", style="dim")
        
        return [
            empty,          # Flow number
            empty,          # Source  
            empty,          # Protocol
            empty,          # Destination
            extended_text,  # Extended protocol
            frame_text,     # Frame type
            count_text,     # Packet count
            volume_text,    # Volume
            timing_text,    # Timing
            quality_text    # Percentage
        ]
    
    def _get_protocol_frame_combinations(self, flow: 'FlowStats') -> List[Tuple[str, str, int, float]]:
        """Get distinct extended protocol/frame type combinations for a flow"""
        combinations = []
        total_packets = flow.frame_count
        
        # Group frame types by extended protocol
        protocol_frames = {}
        
        if flow.frame_types:
            for frame_type, ft_stats in flow.frame_types.items():
                # Determine extended protocol for this frame type
                extended_proto = self._get_extended_protocol_for_frame(flow, frame_type)
                
                if extended_proto not in protocol_frames:
                    protocol_frames[extended_proto] = []
                
                protocol_frames[extended_proto].append((frame_type, ft_stats.count))
        else:
            # No frame types, just show the flow-level extended protocol
            extended_proto = self._get_extended_protocol(flow)
            protocol_frames[extended_proto] = [("General", total_packets)]
        
        # Convert to list of tuples with percentages
        for extended_proto, frame_list in protocol_frames.items():
            for frame_type, count in frame_list:
                percentage = (count / total_packets * 100) if total_packets > 0 else 0
                combinations.append((extended_proto, frame_type, count, percentage))
        
        # Sort by count (descending)
        combinations.sort(key=lambda x: x[2], reverse=True)
        return combinations
    
    def _get_extended_protocol_for_frame(self, flow: 'FlowStats', frame_type: str) -> str:
        """Get extended protocol for a specific frame type"""
        if frame_type.startswith('CH10') or frame_type == 'TMATS':
            return 'CH10'
        elif frame_type.startswith('PTP'):
            return 'PTP'
        elif frame_type == 'IENA':
            return 'IENA'
        elif frame_type == 'NTP':
            return 'NTP'
        else:
            return self._get_extended_protocol(flow)
    
    def _get_extended_protocol(self, flow: 'FlowStats') -> str:
        """Get extended protocol (Chapter 10, PTP, IENA, etc.)"""
        if flow.detected_protocol_types:
            # Look for specialized protocols
            enhanced_protocols = {'CHAPTER10', 'CH10', 'PTP', 'IENA'}
            found_enhanced = flow.detected_protocol_types & enhanced_protocols
            if found_enhanced:
                protocol = list(found_enhanced)[0]
                # Simplify display names
                if protocol in ['CHAPTER10', 'CH10']:
                    return 'CH10'
                return protocol
        
        # Check for other common protocols
        if flow.detected_protocol_types and 'NTP' in flow.detected_protocol_types:
            return 'NTP'
        
        return '-'
    
    def _format_bytes(self, bytes_count: int) -> str:
        """Format byte count with appropriate units"""
        if bytes_count >= 1_000_000_000:
            return f"{bytes_count / 1_000_000_000:.1f}GB"
        elif bytes_count >= 1_000_000:
            return f"{bytes_count / 1_000_000:.1f}MB"
        elif bytes_count >= 1_000:
            return f"{bytes_count / 1_000:.1f}KB"
        else:
            return f"{bytes_count}B"
    
    def _get_flows_list(self) -> List['FlowStats']:
        """Get flows sorted by importance for flow analysis"""
        flows_list = list(self.analyzer.flows.values())
        
        # Sort by: Enhanced protocols first, then outliers, then packet count
        flows_list.sort(key=lambda x: (
            x.enhanced_analysis.decoder_type != "Standard",
            len(x.outlier_frames),
            x.frame_count
        ), reverse=True)
        
        return flows_list
    
    def get_selected_flow(self) -> Optional['FlowStats']:
        """Get currently selected flow"""
        if not self.flow_table or not self.flows_list:
            return None
        
        cursor_row = self.flow_table.cursor_row
        if 0 <= cursor_row < len(self.flows_list):
            return self.flows_list[cursor_row]
        
        return None
    
    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        """Handle row selection in the data table"""
        # Extract flow index from row key
        if event.row_key and event.row_key.startswith("flow_"):
            try:
                # Parse "flow_N" or "flow_N_sub_M" to get flow index
                parts = event.row_key.split("_")
                flow_index = int(parts[1])
                self.selected_flow_index = flow_index
            except (IndexError, ValueError):
                pass