"""
R8: Smart Workflow Truncation

Provides intelligent truncation of ComfyUI workflow JSON that:
1. Preserves error-related nodes and their connections
2. Keeps nodes within N hops of the error node
3. Falls back to property pruning before node removal
4. Returns metadata about what was truncated
"""
import json
from typing import Optional


def _normalize_error_node_id(error_node_id: Optional[str]) -> Optional[str]:
    """Normalize execution IDs like '12:34:56' to their local node id."""
    if error_node_id is None:
        return None
    text = str(error_node_id)
    if not text:
        return None
    return text.split(':')[-1]


def truncate_workflow_smart(
    workflow_str: str,
    error_node_id: Optional[str] = None,
    max_chars: int = 4000,
    max_hops: int = 2
) -> tuple[str, dict]:
    """
    Smartly truncate a workflow JSON string while preserving relevant context.
    
    Args:
        workflow_str: The workflow JSON as a string (or simplified text)
        error_node_id: ID of the node that caused the error (prioritized)
        max_chars: Maximum character length for output
        max_hops: How many connection hops from error node to keep
        
    Returns:
        Tuple of (truncated_workflow_str, metadata_dict)
        
        metadata_dict contains:
        - original_length: Original string length
        - truncated_length: Final string length
        - original_nodes: Number of nodes before truncation
        - kept_nodes: Number of nodes after truncation
        - truncation_method: "none", "property_pruning", "node_removal", or "char_slice"
        - error_node_preserved: Whether the error node was kept
    """
    if not workflow_str:
        return "", {"truncation_method": "none", "original_length": 0}
    
    original_length = len(workflow_str)
    
    # If already under limit, return as-is
    if original_length <= max_chars:
        return workflow_str, {
            "truncation_method": "none",
            "original_length": original_length,
            "truncated_length": original_length
        }
    
    # Try to parse as JSON
    try:
        workflow = json.loads(workflow_str)
    except json.JSONDecodeError:
        # Not valid JSON, fall back to simple character truncation
        truncated = workflow_str[:max_chars - 50] + "\n\n[... workflow truncated ...]"
        return truncated, {
            "truncation_method": "char_slice",
            "original_length": original_length,
            "truncated_length": len(truncated),
            "reason": "invalid_json"
        }
    
    # Workflow can be dict (API format) or has "nodes" key
    nodes = None
    if isinstance(workflow, dict):
        if "nodes" in workflow:
            nodes = workflow["nodes"]
        elif all(isinstance(v, dict) for v in workflow.values()):
            # Dict of node_id -> node_data format
            nodes = list(workflow.values())
    elif isinstance(workflow, list):
        nodes = workflow
    
    if not nodes:
        # Unknown format, use character truncation
        truncated = workflow_str[:max_chars - 50] + "\n\n[... workflow truncated ...]"
        return truncated, {
            "truncation_method": "char_slice",
            "original_length": original_length,
            "truncated_length": len(truncated),
            "reason": "unknown_format"
        }
    
    original_node_count = len(nodes)
    error_node_preserved = False
    
    # Strategy 1: Property pruning (keep all nodes but remove verbose properties)
    pruned_workflow = _prune_node_properties(workflow, nodes)
    pruned_str = json.dumps(pruned_workflow, ensure_ascii=False, separators=(',', ':'))
    
    if len(pruned_str) <= max_chars:
        return pruned_str, {
            "truncation_method": "property_pruning",
            "original_length": original_length,
            "truncated_length": len(pruned_str),
            "original_nodes": original_node_count,
            "kept_nodes": original_node_count,
            "error_node_preserved": True if error_node_id else None
        }
    
    # Strategy 2: Node removal (keep error node and neighbors)
    normalized_error_node_id = _normalize_error_node_id(error_node_id)
    if normalized_error_node_id:
        priority_nodes = _get_priority_nodes(workflow, normalized_error_node_id, max_hops)
    else:
        # No error node specified, keep first N nodes
        priority_nodes = set(str(n.get("id", i)) for i, n in enumerate(nodes[:20]))
    
    filtered_source_workflow = pruned_workflow if isinstance(workflow, dict) and "nodes" in workflow else workflow
    filtered_source_nodes = (
        filtered_source_workflow.get("nodes", [])
        if isinstance(filtered_source_workflow, dict) and "nodes" in filtered_source_workflow
        else nodes
    )
    filtered_workflow = _filter_nodes(filtered_source_workflow, filtered_source_nodes, priority_nodes)
    filtered_str = json.dumps(filtered_workflow, ensure_ascii=False, separators=(',', ':'))
    
    # Check if error node is preserved
    if normalized_error_node_id:
        for n in (filtered_workflow.get("nodes", []) if isinstance(filtered_workflow, dict) else filtered_workflow):
            if str(n.get("id", "")) == str(normalized_error_node_id):
                error_node_preserved = True
                break
    
    kept_node_count = len(filtered_workflow.get("nodes", filtered_workflow) if isinstance(filtered_workflow, dict) else filtered_workflow)
    
    if len(filtered_str) <= max_chars:
        return filtered_str, {
            "truncation_method": "node_removal",
            "original_length": original_length,
            "truncated_length": len(filtered_str),
            "original_nodes": original_node_count,
            "kept_nodes": kept_node_count,
            "error_node_preserved": error_node_preserved
        }
    
    # Strategy 3: Final fallback - character truncation with note
    truncated = filtered_str[:max_chars - 80] + ',"_truncated":true}'
    # Ensure valid JSON by trying to close it properly
    try:
        json.loads(truncated)
    except json.JSONDecodeError:
        # Can't make valid JSON, just slice
        truncated = filtered_str[:max_chars - 50] + "\n\n[... workflow truncated ...]"
    
    return truncated, {
        "truncation_method": "char_slice",
        "original_length": original_length,
        "truncated_length": len(truncated),
        "original_nodes": original_node_count,
        "kept_nodes": kept_node_count,
        "error_node_preserved": error_node_preserved
    }


def _prune_node_properties(workflow: dict | list, nodes: list) -> dict | list:
    """Remove verbose/non-essential properties from nodes."""
    PRUNE_KEYS = {"widgets_values", "pos", "size", "order", "mode", "flags", "color", "bgcolor"}

    def trim_value(value):
        if isinstance(value, str) and len(value) > 100:
            return value[:100] + "..."
        if isinstance(value, dict):
            return {k: trim_value(v) for k, v in value.items()}
        if isinstance(value, list):
            return [trim_value(v) for v in value]
        return value

    def prune_node(node: dict) -> dict:
        pruned = {}
        for key, value in node.items():
            if key in PRUNE_KEYS:
                continue
            if key in {"inputs", "properties"} and isinstance(value, dict):
                pruned[key] = {k: trim_value(v) for k, v in value.items()}
            else:
                pruned[key] = trim_value(value)
        return pruned

    pruned_nodes = [prune_node(n) for n in nodes]

    if isinstance(workflow, dict):
        if "nodes" in workflow:
            result = {k: v for k, v in workflow.items() if k != "nodes"}
            result["nodes"] = pruned_nodes
            return result
        return {k: prune_node(v) if isinstance(v, dict) else v for k, v in workflow.items()}
    return pruned_nodes


def _get_priority_nodes(workflow: dict | list, error_node_id: str, max_hops: int) -> set:
    """Get node IDs within max_hops of the error node."""
    nodes = workflow.get("nodes", workflow) if isinstance(workflow, dict) else workflow
    
    # Build adjacency map
    node_map = {}
    adjacency = {}
    
    for node in nodes:
        node_id = str(node.get("id", ""))
        node_map[node_id] = node
        adjacency[node_id] = set()
        
        # Check inputs for connections
        inputs = node.get("inputs", {})
        if isinstance(inputs, dict):
            for inp_key, inp_val in inputs.items():
                # Connection format: [node_id, output_index] or just node_id
                if isinstance(inp_val, list) and len(inp_val) >= 1:
                    adjacency[node_id].add(str(inp_val[0]))
                elif isinstance(inp_val, (int, str)) and str(inp_val) in node_map:
                    adjacency[node_id].add(str(inp_val))
    
    # BFS from error node
    priority = {str(error_node_id)}
    frontier = {str(error_node_id)}
    
    for _ in range(max_hops):
        next_frontier = set()
        for node_id in frontier:
            # Add neighbors
            next_frontier.update(adjacency.get(node_id, set()))
            # Add nodes that connect TO this node
            for nid, neighbors in adjacency.items():
                if node_id in neighbors:
                    next_frontier.add(nid)
        priority.update(next_frontier)
        frontier = next_frontier - priority
    
    return priority


def _filter_links_for_nodes(links: list, keep_ids: set) -> list:
    """Keep only links whose endpoints are preserved nodes."""
    filtered = []
    for link in links:
        if not isinstance(link, list) or len(link) < 4:
            filtered.append(link)
            continue
        origin_id = str(link[1]) if len(link) > 1 else None
        target_id = str(link[3]) if len(link) > 3 else None
        if origin_id in keep_ids and target_id in keep_ids:
            filtered.append(link)
    return filtered


def _preserve_workflow_shell(workflow: dict, filtered_nodes: list, keep_ids: set) -> dict:
    """Preserve UI workflow metadata when pruning nodes."""
    result = {k: v for k, v in workflow.items() if k not in {"nodes", "links"}}
    result["nodes"] = filtered_nodes
    if isinstance(workflow.get("links"), list):
        result["links"] = _filter_links_for_nodes(workflow["links"], keep_ids)
    return result


def _filter_nodes(workflow: dict | list, nodes: list, keep_ids: set) -> dict | list:
    """Filter workflow to only keep specified node IDs."""
    filtered = [n for n in nodes if str(n.get("id", "")) in keep_ids]

    if isinstance(workflow, dict):
        if "nodes" in workflow:
            return _preserve_workflow_shell(workflow, filtered, keep_ids)
        return {k: v for k, v in workflow.items() if str(k) in keep_ids}
    return filtered