"""
Transaction progress data processor.

Processes raw transaction progress data into an immutable tree structure,
separating data transformation from display concerns.
"""
from typing import Dict, List, Optional
from collections import defaultdict
from dataclasses import dataclass

from dateutil.parser import isoparse


@dataclass(frozen=True)
class ProcessedTask:
    """Immutable tree representation of a processed task."""
    task_id: str
    task_name: str
    duration: Optional[float]  # Duration in seconds (None if not finished)
    start_time: Optional[float]  # Start timestamp in seconds since epoch (None if not started)
    end_time: Optional[float]  # End timestamp in seconds since epoch (None if not finished)
    data: Optional[Dict]  # Optional data field for additional info
    warnings: tuple[tuple[str, str], ...]  # Tuple of (warning_code, message)
    children: tuple['ProcessedTask', ...]  # Immutable children


@dataclass(frozen=True)
class ProgressSummaryItem:
    """Summary item for a root task showing its wall clock contribution."""
    task_name: str
    wall_clock_duration: float  # Wall clock time in seconds
    ivm_relative_update_size: Optional[float] = None  # Relative update size in percent
    loop_iteration: Optional[int] = None  # Loop iteration count (for loop tasks)


def _duration_s(task: Dict) -> Optional[float]:
    """Calculate task duration in seconds. Return None if not finished."""
    if "start_time" not in task or "end_time" not in task:
        return None

    start = isoparse(task["start_time"])
    end = isoparse(task["end_time"])
    duration = (end - start).total_seconds()
    if duration < 0:
        # The server reuses loop names across loop iterations, so a
        # polled state can have start_time from the last iteration and
        # end_time from an earlier iteration, producing inverted timestamps.
        return None
    return duration


def _build_graph(tasks: Dict) -> Dict:
    """Build parent-child relationship graph from tasks."""
    children = defaultdict(list)

    for tid, t in tasks.items():
        if "origin_id" in t:
            parent = t["origin_id"]
            children[parent].append(tid)

    return children


def _find_roots(tasks: Dict) -> list[str]:
    """Find root tasks (those with no parents)."""
    return [tid for tid in tasks if "origin_id" not in tasks[tid] or tasks[tid]["origin_id"] not in tasks]


def build_progress_tree(tasks: Dict) -> List[ProcessedTask]:
    """
    Build an immutable tree structure from raw task data.
    
    Args:
        tasks: Raw task data from transaction progress
        
    Returns:
        List of root ProcessedTask tree nodes
    """
    children = _build_graph(tasks)
    roots = _find_roots(tasks)
    
    # Build tree structure
    processed = []
    for root_id in roots:
        node = _build_task_node(root_id, tasks[root_id], children, tasks)
        if node:
            processed.append(node)
    
    return processed


def _build_task_node(
    task_id: str,
    task: Dict,
    children_map: Dict,
    all_tasks: Dict
) -> Optional[ProcessedTask]:
    """
    Recursively build a task tree node with its children.
    
    Args:
        task_id: The transaction task ID
        task: The task data
        children_map: Map of parent IDs to child IDs
        all_tasks: All tasks in the transaction
        
    Returns:
        ProcessedTask tree node or None if task should be filtered
    """
    # Extract raw task data
    task_name = task.get("task_name", task_id)
    duration = _duration_s(task)
    
    # Extract timestamps
    start_time = None
    end_time = None
    if "start_time" in task:
        start_time = isoparse(task["start_time"]).timestamp()
    if "end_time" in task:
        end_time = isoparse(task["end_time"]).timestamp()
    
    # Extract data
    data = task.get("data")

    # Extract warnings
    warning_list = []
    if "warnings" in task and task["warnings"]:
        for warning_code, warning_data in task["warnings"].items():
            message = warning_data.get("message", "")
            warning_list.append((warning_code, message))
    
    # Process children
    child_nodes = []
    if task_id in children_map:
        child_list = children_map[task_id]
        for child_id in child_list:
            child_node = _build_task_node(
                child_id,
                all_tasks[child_id],
                children_map,
                all_tasks
            )
            if child_node:
                child_nodes.append(child_node)
    
    return ProcessedTask(
        task_id=task_id,
        task_name=task_name,
        duration=duration,
        start_time=start_time,
        end_time=end_time,
        data=data,
        warnings=tuple(warning_list),
        children=tuple(child_nodes)
    )


def _collect_time_intervals(task: ProcessedTask) -> List[tuple[float, float]]:
    """
    Collect all time intervals (start, end) for a task and its descendants.
    
    Args:
        task: The processed task node
        
    Returns:
        List of (start_timestamp, end_timestamp) tuples in seconds since epoch
    """
    intervals = []
    
    # Get this task's interval if it has timestamps
    if task.start_time is not None and task.end_time is not None and task.start_time <= task.end_time:
        intervals.append((task.start_time, task.end_time))
    
    # Recursively collect children's intervals
    for child in task.children:
        intervals.extend(_collect_time_intervals(child))
    
    return intervals


def _merge_intervals(intervals: List[tuple[float, float]]) -> List[tuple[float, float]]:
    """
    Merge overlapping time intervals to compute wall clock time.
    
    Args:
        intervals: List of (start, end) tuples
        
    Returns:
        List of merged non-overlapping intervals
    """
    if not intervals:
        return []
    
    # Sort by start time
    sorted_intervals = sorted(intervals)
    merged = [sorted_intervals[0]]
    
    for start, end in sorted_intervals[1:]:
        last_start, last_end = merged[-1]
        
        # If intervals overlap or are adjacent, merge them
        if start <= last_end:
            merged[-1] = (last_start, max(last_end, end))
        else:
            # Non-overlapping, add as new interval
            merged.append((start, end))
    
    return merged


def _compute_wall_clock_duration(task: ProcessedTask) -> float:
    """
    Compute wall clock duration for a task and all its descendants.
    
    For parallel tasks, counts overlapping time only once.
    For sequential tasks, sums up the durations.
    
    Args:
        task: The processed task node
        
    Returns:
        Wall clock duration in seconds (0 if no valid timestamps)
    """
    intervals = _collect_time_intervals(task)
    merged = _merge_intervals(intervals)
    
    # Sum up the merged intervals to get total wall clock time
    return sum(end - start for start, end in merged)

def get_ivm_relative_update_size(task: ProcessedTask) -> Optional[float]:
    """
    Extract the IVM relative update size from task data if available.
    
    Args:
        task: The processed task node
    Returns:
        IVM relative update size in percent, or None if not available
    """
    return task.data.get("ivm_relative_update_size") if task.data else None

def get_loop_iteration(task: ProcessedTask) -> Optional[int]:
    """
    Extract the loop iteration number from task data if available.

    Args:
        task: The processed task node
    Returns:
        Loop iteration number, or None if not available
    """
    return task.data.get("loop_iteration") if task.data else None

def compute_progress_summary(
    task_tree: List[ProcessedTask]
) -> List[ProgressSummaryItem]:
    """
    Compute progress summary showing wall clock duration of each root task.
    
    Args:
        task_tree: List of root ProcessedTask nodes
        
    Returns:
        List of ProgressSummaryItem with task names and durations
    """
    summary = []
    for task in task_tree:
        wall_clock = _compute_wall_clock_duration(task)
        
        # Skip tasks that have no timestamp data
        if wall_clock == 0 and task.end_time is None:
            continue

        summary.append(ProgressSummaryItem(
            task_name=task.task_name,
            wall_clock_duration=wall_clock,
            ivm_relative_update_size=get_ivm_relative_update_size(task),
            loop_iteration=get_loop_iteration(task),
        ))
    
    return summary