#!/usr/bin/env python3
"""
Task 2: Wire batch actions to Container Actions sub-workflow
"""

import json
import uuid

CONTAINER_ACTIONS_WF_ID = "fYSZS5PkH0VSEaT5"

def load_workflow():
    with open('n8n-workflow.json', 'r') as f:
        return json.load(f)

def save_workflow(workflow):
    with open('n8n-workflow.json', 'w') as f:
        json.dump(workflow, f, indent=2)
    print(f"Saved workflow with {len(workflow['nodes'])} nodes")

def find_node(workflow, name):
    for node in workflow['nodes']:
        if node['name'] == name:
            return node
    return None

def remove_node(workflow, node_name):
    """Remove a node and all its connections"""
    # Remove from nodes list
    workflow['nodes'] = [n for n in workflow['nodes'] if n['name'] != node_name]

    # Remove from connections (as source)
    if node_name in workflow['connections']:
        del workflow['connections'][node_name]

    # Remove from connections (as target)
    for source, outputs in list(workflow['connections'].items()):
        for output_key, connections in list(outputs.items()):
            workflow['connections'][source][output_key] = [
                [conn for conn in conn_list if conn.get('node') != node_name]
                for conn_list in connections
            ]

def create_execute_workflow_node(name, workflow_id, position):
    """Create an Execute Workflow node with n8n 1.2 format"""
    return {
        "parameters": {
            "workflowId": {
                "__rl": True,
                "mode": "list",
                "value": workflow_id
            },
            "options": {}
        },
        "id": str(uuid.uuid4()),
        "name": name,
        "type": "n8n-nodes-base.executeWorkflow",
        "typeVersion": 1.2,
        "position": position
    }

def create_code_node(name, code, position):
    """Create a Code node"""
    return {
        "parameters": {
            "jsCode": code
        },
        "id": str(uuid.uuid4()),
        "name": name,
        "type": "n8n-nodes-base.code",
        "typeVersion": 2,
        "position": position
    }

def main():
    print("Loading workflow...")
    workflow = load_workflow()
    initial_count = len(workflow['nodes'])
    print(f"Initial node count: {initial_count}")

    # Current batch action flow:
    # Route Batch Loop Action (outputs 1,2,3) -> Build Batch Action Command ->
    # Execute Batch Container Action -> Check Batch Action Result -> ...

    # We need to replace this with:
    # Route Batch Loop Action -> Prepare Batch Action Input -> Execute Batch Action ->
    # Handle Batch Action Result

    # 1. Create "Prepare Batch Action Input" code node
    prepare_code = '''// Prepare input for Container Actions sub-workflow
const data = $json;
const container = data.container;
const action = data.action;

// Extract container info
const containerId = container.id || container.Id || '';
const containerName = container.name || container.Name || '';

return {
  json: {
    containerId: containerId,
    containerName: containerName,
    action: action,
    chatId: data.chatId,
    messageId: data.progressMessageId || 0,
    responseMode: "inline"
  }
};'''

    prepare_node = create_code_node(
        "Prepare Batch Action Input",
        prepare_code,
        [4000, -200]
    )

    # 2. Create "Execute Batch Action" execute workflow node
    execute_node = create_execute_workflow_node(
        "Execute Batch Action Sub-workflow",
        CONTAINER_ACTIONS_WF_ID,
        [4220, -200]
    )

    # 3. Create "Handle Batch Action Result" code node
    handle_code = '''// Handle action result from sub-workflow
const data = $('Build Progress Message').item.json;
const result = $json;

// Update counters based on result
let successCount = data.successCount || 0;
let failureCount = data.failureCount || 0;
let warningCount = data.warningCount || 0;

if (result.success) {
  successCount++;
} else {
  failureCount++;
}

// Add to results array
const results = data.results || [];
results.push({
  container: data.containerName,
  action: data.action,
  success: result.success,
  message: result.message || ''
});

return {
  json: {
    ...data,
    successCount: successCount,
    failureCount: failureCount,
    warningCount: warningCount,
    results: results
  }
};'''

    handle_node = create_code_node(
        "Handle Batch Action Result Sub",
        handle_code,
        [4440, -200]
    )

    # Add nodes to workflow
    print("\nAdding new nodes:")
    print(f"  - {prepare_node['name']}")
    print(f"  - {execute_node['name']}")
    print(f"  - {handle_node['name']}")

    workflow['nodes'].extend([prepare_node, execute_node, handle_node])

    # Update connections from Route Batch Loop Action
    # Outputs 1, 2, 3 (start, stop, restart) should go to Prepare Batch Action Input
    print("\nUpdating connections:")

    for i in [1, 2, 3]:
        workflow['connections']['Route Batch Loop Action']['main'][i] = [{
            "node": "Prepare Batch Action Input",
            "type": "main",
            "index": 0
        }]
    print("  - Route Batch Loop Action [start/stop/restart] -> Prepare Batch Action Input")

    # Add new connections
    workflow['connections']['Prepare Batch Action Input'] = {
        'main': [[{
            "node": "Execute Batch Action Sub-workflow",
            "type": "main",
            "index": 0
        }]]
    }
    print("  - Prepare Batch Action Input -> Execute Batch Action Sub-workflow")

    workflow['connections']['Execute Batch Action Sub-workflow'] = {
        'main': [[{
            "node": "Handle Batch Action Result Sub",
            "type": "main",
            "index": 0
        }]]
    }
    print("  - Execute Batch Action Sub-workflow -> Handle Batch Action Result Sub")

    workflow['connections']['Handle Batch Action Result Sub'] = {
        'main': [[{
            "node": "Prepare Next Iteration",
            "type": "main",
            "index": 0
        }]]
    }
    print("  - Handle Batch Action Result Sub -> Prepare Next Iteration")

    # Now remove the old nodes that are no longer needed
    print("\nRemoving obsolete nodes:")
    nodes_to_remove = [
        "Build Batch Action Command",
        "Execute Batch Container Action",
        "Check Batch Action Result",
        "Needs Action Call",
        "Execute Batch Action 2",
        "Parse Batch Action 2",
        "Handle Action Result"
    ]

    removed_count = 0
    for node_name in nodes_to_remove:
        if find_node(workflow, node_name):
            print(f"  - Removing: {node_name}")
            remove_node(workflow, node_name)
            removed_count += 1

    # Save
    final_count = len(workflow['nodes'])
    print(f"\nNode count: {initial_count} -> {final_count} ({final_count - initial_count:+d})")
    print(f"Removed: {removed_count} nodes")
    print(f"Added: 3 nodes")
    print(f"Net change: {final_count - initial_count:+d} nodes")

    save_workflow(workflow)
    print("\n✓ Task 2 complete: Batch actions now use Container Actions sub-workflow")

if __name__ == '__main__':
    main()