"""
FINAL CORRECTED _renderWithBudgetFromStructure - Version 3

This file contains the CORRECT implementation that should REPLACE
the existing methods in jsonContinuation.py

KEY BEHAVIOR:
1. Budget is allocated from CUT → ROOT (not top-down)
2. Cut-near values get priority
3. When budget < 50: summary_mode enabled, non-path containers → <object>/<array>
4. Path containers always render their structure

COPY THESE METHODS INTO YOUR JsonAnalyzer CLASS:
- _renderWithBudgetFromStructure (REPLACE existing)
- _buildPathFromCutToRootV3 (ADD)
- _collectAllValuesWithDistance (ADD)
- _renderNodeV3 (ADD)
- _renderObjectV3 (ADD)
- _renderArrayV3 (ADD)
- _renderValueV3 (ADD)
"""

from typing import List, Set
from dataclasses import dataclass, field


@dataclass
class BudgetAllocation:
    """Tracks which nodes have been allocated budget"""
    allocated_node_ids: Set[int] = field(default_factory=set)
    path_node_ids: Set[int] = field(default_factory=set)
    summary_mode: bool = False


# =============================================================================
# METHODS TO COPY INTO JsonAnalyzer CLASS
# =============================================================================

def _renderWithBudgetFromStructure(self, structure: dict, cutPos: int) -> str:
    """
    Render structure with budget logic - allocate from CUT to ROOT.
    
    ALGORITHM:
    
    Phase 1: Build path from cut to root
    - Find the cut element (truncated value or deepest incomplete node)
    - Build ordered path: [cut_element, parent, grandparent, ..., root]
    
    Phase 2: Allocate budget
    - Collect ALL value nodes with their distance to cut
    - Sort by distance (smaller = closer to cut = higher priority)
    - Allocate budget to values in this order
    - When budget < 50: enable summary_mode (affects containers only)
    
    Phase 3: Render
    - PATH containers: always render structure
    - NON-PATH containers in summary_mode: render as <object>/<array>
    - Values: render if allocated, else type hint
    
    Returns:
        Rendered JSON string with budget constraints applied
    """
    # Phase 1: Build path from cut to root
    pathFromCutToRoot = []
    self._buildPathFromCutToRootV3(structure, cutPos, [], pathFromCutToRoot)
    
    pathNodeIds = set(id(node) for node in pathFromCutToRoot)
    
    # Phase 2: Collect ALL values and allocate budget
    allValues = []
    self._collectAllValuesWithDistance(structure, cutPos, allValues)
    
    # Sort by distance (smaller = closer to cut = higher priority)
    allValues.sort(key=lambda x: x['distance'])
    
    # Initialize allocation tracker
    allocation = BudgetAllocation(
        path_node_ids=pathNodeIds,
        allocated_node_ids=set(),
        summary_mode=False
    )
    
    remainingBudget = self.budgetLimit
    
    # Phase 2a: Allocate PATH values first (truncated values are always rendered)
    pathValues = [item for item in allValues if id(item['node']) in pathNodeIds]
    for item in pathValues:
        node = item['node']
        nodeType = node.get('type')
        
        if nodeType == 'truncated_value':
            allocation.allocated_node_ids.add(id(node))
            continue
        
        if nodeType != 'value':
            continue
        
        rawValue = node.get('raw', '')
        valueSize = len(rawValue)
        
        if valueSize <= remainingBudget:
            allocation.allocated_node_ids.add(id(node))
            remainingBudget -= valueSize
            
            if remainingBudget < 50:
                allocation.summary_mode = True
    
    # Phase 2b: Allocate NON-PATH values (skip if path already triggered summary mode)
    if not allocation.summary_mode:
        nonPathValues = [item for item in allValues if id(item['node']) not in pathNodeIds]
        for item in nonPathValues:
            node = item['node']
            nodeType = node.get('type')
            
            if nodeType != 'value':
                continue
            
            rawValue = node.get('raw', '')
            valueSize = len(rawValue)
            
            if valueSize <= remainingBudget:
                allocation.allocated_node_ids.add(id(node))
                remainingBudget -= valueSize
                
                if remainingBudget < 50 and not allocation.summary_mode:
                    allocation.summary_mode = True
    
    # Phase 3: Render with allocation info
    return self._renderNodeV3(structure, 0, allocation)


def _buildPathFromCutToRootV3(self, node: dict, cutPos: int, currentPath: list, resultPath: list) -> bool:
    """
    Recursively find the path from root to cut element, then reverse it.
    Result path is ordered: [cut_element, parent, ..., root]
    """
    nodeType = node.get('type')
    startPos = node.get('start_pos', 0)
    endPos = node.get('end_pos', cutPos + 1)
    
    pathWithCurrent = currentPath + [node]
    
    for child in node.get('children', []):
        if self._buildPathFromCutToRootV3(child, cutPos, pathWithCurrent, resultPath):
            return True
    
    if nodeType == 'truncated_value':
        resultPath.clear()
        resultPath.extend(reversed(pathWithCurrent))
        return True
    
    if nodeType == 'value' and startPos <= cutPos <= endPos:
        resultPath.clear()
        resultPath.extend(reversed(pathWithCurrent))
        return True
    
    if nodeType in ('object', 'array') and not node.get('complete') and startPos <= cutPos:
        resultPath.clear()
        resultPath.extend(reversed(pathWithCurrent))
        return True
    
    if nodeType == 'root' and not resultPath:
        resultPath.clear()
        resultPath.extend(reversed(pathWithCurrent))
        return True
    
    return False


def _collectAllValuesWithDistance(self, node: dict, cutPos: int, result: list, depth: int = 0):
    """Collect ALL value nodes with their distance to cut point."""
    nodeType = node.get('type')
    
    if nodeType in ('value', 'truncated_value'):
        endPos = node.get('end_pos', cutPos)
        distance = cutPos - endPos
        result.append({
            'node': node,
            'distance': distance,
            'depth': depth
        })
    
    for child in node.get('children', []):
        self._collectAllValuesWithDistance(child, cutPos, result, depth + 1)


def _renderNodeV3(self, node: dict, depth: int, allocation) -> str:
    """Render a node with budget allocation info."""
    nodeType = node.get('type')
    
    if nodeType == 'root':
        parts = []
        for child in node.get('children', []):
            parts.append(self._renderNodeV3(child, depth, allocation))
        return '\n'.join(parts)
    
    elif nodeType == 'object':
        return self._renderObjectV3(node, depth, allocation)
    
    elif nodeType == 'array':
        return self._renderArrayV3(node, depth, allocation)
    
    elif nodeType == 'value':
        return self._renderValueV3(node, depth, allocation)
    
    elif nodeType == 'truncated_value':
        keyPrefix = f'"{node.get("key")}": ' if node.get('key') else ''
        return f"{keyPrefix}{node.get('raw', '')}"
    
    return ''


def _renderObjectV3(self, node: dict, depth: int, allocation) -> str:
    """Render object - summary mode non-path objects become <object>."""
    indentStr = "  " * depth
    innerIndent = "  " * (depth + 1)
    
    keyPrefix = f'"{node.get("key")}": ' if node.get('key') else ''
    children = node.get('children', [])
    isOnPath = id(node) in allocation.path_node_ids
    
    if allocation.summary_mode and not isOnPath:
        return f"{keyPrefix}<object>"
    
    if not children:
        return f"{keyPrefix}{{}}" if node.get('complete') else f"{keyPrefix}{{"
    
    parts = [f"{keyPrefix}{{"]
    
    for i, child in enumerate(children):
        childRendered = self._renderNodeV3(child, depth + 1, allocation)
        isLast = (i == len(children) - 1)
        isTruncated = child.get('type') == 'truncated_value'
        
        if isLast or isTruncated:
            parts.append(f"{innerIndent}{childRendered}")
        else:
            parts.append(f"{innerIndent}{childRendered},")
    
    if node.get('complete'):
        parts.append(f"{indentStr}}}")
    
    return '\n'.join(parts)


def _renderArrayV3(self, node: dict, depth: int, allocation) -> str:
    """Render array - summary mode non-path arrays become <array>."""
    indentStr = "  " * depth
    innerIndent = "  " * (depth + 1)
    
    keyPrefix = f'"{node.get("key")}": ' if node.get('key') else ''
    children = node.get('children', [])
    isOnPath = id(node) in allocation.path_node_ids
    
    if allocation.summary_mode and not isOnPath:
        return f"{keyPrefix}<array>"
    
    if not children:
        return f"{keyPrefix}[]" if node.get('complete') else f"{keyPrefix}["
    
    parts = [f"{keyPrefix}["]
    
    for i, child in enumerate(children):
        childRendered = self._renderNodeV3(child, depth + 1, allocation)
        isLast = (i == len(children) - 1)
        isTruncated = child.get('type') == 'truncated_value'
        
        if isLast or isTruncated:
            parts.append(f"{innerIndent}{childRendered}")
        else:
            parts.append(f"{innerIndent}{childRendered},")
    
    if node.get('complete'):
        parts.append(f"{indentStr}]")
    
    return '\n'.join(parts)


def _renderValueV3(self, node: dict, depth: int, allocation) -> str:
    """Render value - if allocated render full, else type hint."""
    keyPrefix = f'"{node.get("key")}": ' if node.get('key') else ''
    rawValue = node.get('raw', '""')
    valueType = node.get('value_type', 'string')
    
    typeHints = {
        'string': '<str>',
        'number': '<number>',
        'boolean': '<boolean>',
        'null': '<null>'
    }
    typeHint = typeHints.get(valueType, '<value>')
    
    if id(node) in allocation.allocated_node_ids:
        return f"{keyPrefix}{rawValue}"
    else:
        return f"{keyPrefix}{typeHint}"


# =============================================================================
# ALSO ADD THIS IMPORT AT THE TOP OF YOUR FILE
# =============================================================================
# from dataclasses import dataclass, field
# from typing import Set

# And add the BudgetAllocation class inside your file or as a nested class