from datetime import datetime
from collections import defaultdict
from rich.text import Text
from rich.console import Console
from v0.relationalai.util.format import format_duration

# Create a console for rendering Rich text to ANSI codes
# Force terminal mode and no pager to ensure ANSI codes are always generated
_console = Console(force_terminal=True, legacy_windows=False)

def render_line(line_data):
    """Convert structured line data to rendered string with optional ANSI color codes.
    
    Args:
        line_data: dict with 'text' (str) and 'color' (str|None) fields.
                   color can be 'green', or None
    
    Returns:
        str: Plain text or ANSI-colored text
    """
    color = line_data.get("color")
    if color:
        text = Text(line_data["text"], style=color)
        with _console.capture() as capture:
            _console.print(text, end="")
        return capture.get()
    return line_data["text"]

def parse_ts(ts):
    ts = ts.replace("Z", "")

    # Handle fractional seconds - normalize to 6 digits (microseconds)
    if "." in ts:
        parts = ts.split(".")
        if len(parts) == 2:
            # Pad or truncate fractional seconds to 6 digits
            fractional = parts[1].ljust(6, '0')[:6]
            ts = f"{parts[0]}.{fractional}"

    return datetime.fromisoformat(ts)

def duration_ms(task):
    if "start_time" not in task or "end_time" not in task:
        return 0
    return int((parse_ts(task["end_time"]) - parse_ts(task["start_time"])).total_seconds() * 1000)

def build_graph(tasks):
    children = defaultdict(list)
    parents = defaultdict(list)

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

    return parents, children

def find_roots(tasks, parents):
    # roots = nodes with no parents (or parents not in tasks)
    return [tid for tid in tasks if tid not in parents or not any(p in tasks for p in parents[tid])]

def node_label(tid, task):
    """Create structured label data for a task node.
    
    Returns:
        list of dicts with 'text' (str) and 'color' (str|None) fields.
        First item is the task line, subsequent items are warning details.
    """
    name = task.get("task_name", tid)
    is_finished = "start_time" in task and "end_time" in task
    has_warnings = "warnings" in task and task["warnings"]
    warning_count = 0
    warning_text = ""
    text_color = 'green' if is_finished else None
    
    lines = []

    if has_warnings:
        warning_count = len(task["warnings"])
        warning_text = f" ⚠️  {warning_count} warning" + ("s" if warning_count > 1 else "")
    
    if is_finished:
        dur_seconds = duration_ms(task) / 1000.0
        duration_str = format_duration(dur_seconds)
        
        lines.append({"text": f"{name}  ({duration_str}){warning_text}", "color": text_color})
    else:
        lines.append({"text": f"{name}  (Evaluating...){warning_text}", "color": text_color})
    
    # Add warning details as separate lines (for both finished and evaluating tasks)
    if has_warnings:
        for warning_code, warning_data in task["warnings"].items():
            message = warning_data.get("message", "")
            lines.append({"text": f"  ⚠️  {warning_code}: {message}", "color": text_color})
    
    return lines

def build_tree(tid, tasks, children, prefix="", is_last=True, seen=None):
    """Build tree structure and return as list of structured line data dicts.
    
    Returns:
        list of dicts with 'text' (str) and 'color' (str|None) fields
    """
    if seen is None:
        seen = set()

    if tid not in tasks:
        return []

    if tid in seen:
        return []
    seen.add(tid)

    task = tasks[tid]
    connector = "└─ " if is_last else "├─ "
    label_lines = node_label(tid, task)
    
    # Build lines with proper prefix
    lines = []
    for i, label_data in enumerate(label_lines):
        if i == 0:
            # First line gets the tree connector
            full_text = prefix + connector + label_data["text"]
        else:
            # Warning detail lines get continuation prefix
            continuation_prefix = prefix + ("   " if is_last else "│  ")
            full_text = continuation_prefix + label_data["text"]
        lines.append({"text": full_text, "color": label_data["color"]})

    new_prefix = prefix + ("   " if is_last else "│  ")

    if tid not in children:
        return lines

    child_list = [c for c in children[tid] if c in tasks and should_include_task(c, tasks[c])]
    for i, child in enumerate(child_list):
        last = i == len(child_list) - 1
        lines.extend(build_tree(child, tasks, children, new_prefix, last, seen))
    
    return lines

def should_include_task(task_id, task):
    """Filter out error tasks that aren't useful for execution flow analysis"""
    task_name = task.get("task_name", task_id)
    return not task_name.startswith(":error_pyrel_error_attrs")

def format_execution_tree(progress):
    """Format execution tree and return as a string."""
    if not progress or "tasks" not in progress:
        return ""

    tasks = progress["tasks"]
    parents, children = build_graph(tasks)
    roots = find_roots(tasks, parents)

    # Filter out error tasks from roots
    filtered_roots = [root for root in roots if should_include_task(root, tasks[root])]

    # Sort roots: in-progress tasks first, then finished tasks, alphabetically within each group
    def sort_key(task_id):
        task = tasks[task_id]
        is_finished = "start_time" in task and "end_time" in task
        task_name = task.get("task_name", task_id)
        return (is_finished, task_name)
    
    sorted_roots = sorted(filtered_roots, key=sort_key)

    # Build structured data first
    all_line_data = []
    for root in sorted_roots:
        # Add root node (may include warning lines)
        label_lines = node_label(root, tasks[root])
        all_line_data.extend(label_lines)
            
        # Add children
        if root in children:
            child_list = [c for c in children[root] if c in tasks and should_include_task(c, tasks[c])]
            for i, child in enumerate(child_list):
                last = i == len(child_list) - 1
                all_line_data.extend(build_tree(child, tasks, children, "", last, set()))
        
        # Empty line after each root tree
        all_line_data.append({"text": "", "color": None})

    # Render structured data to strings
    rendered_lines = [render_line(line_data) for line_data in all_line_data]
    return "\n\n" + "\n".join(rendered_lines)
