#!/usr/bin/env python3
"""
Task 1: Wire batch update to Container Update sub-workflow
"""

import json
import uuid

# Workflow IDs from STATE.md
CONTAINER_UPDATE_WF_ID = "7AvTzLtKXM2hZTio92_mC"
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 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}")

    # Find Route Batch Loop Action
    route_node = find_node(workflow, "Route Batch Loop Action")
    if not route_node:
        print("ERROR: Could not find Route Batch Loop Action")
        return

    print(f"\nFound Route Batch Loop Action at {route_node['position']}")

    # The Route Batch Loop Action has 4 outputs:
    # 0: update (currently empty)
    # 1: start
    # 2: stop
    # 3: restart

    # We need to add nodes for the update path
    # Position: Route is at [3760, -500], so place new nodes at x=4000+

    # 1. Create "Prepare Batch Update Input" code node
    #    This prepares the input for the sub-workflow
    prepare_code = '''// Prepare input for Container Update sub-workflow
const data = $json;
const container = data.container;

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

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

    prepare_node = create_code_node(
        "Prepare Batch Update Input",
        prepare_code,
        [4000, -800]
    )

    # 2. Create "Execute Batch Update" execute workflow node
    execute_node = create_execute_workflow_node(
        "Execute Batch Update",
        CONTAINER_UPDATE_WF_ID,
        [4220, -800]
    )

    # 3. Create "Handle Batch Update Result" code node
    #    This processes the result from sub-workflow and prepares for next iteration
    handle_code = '''// Handle update 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: 'update',
  success: result.success,
  message: result.message || ''
});

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

    handle_node = create_code_node(
        "Handle Batch Update Result",
        handle_code,
        [4440, -800]
    )

    # 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])

    # Add connections
    print("\nAdding connections:")

    # Route Batch Loop Action (output 0: update) -> Prepare Batch Update Input
    if 'Route Batch Loop Action' not in workflow['connections']:
        workflow['connections']['Route Batch Loop Action'] = {'main': [[], [], [], []]}

    workflow['connections']['Route Batch Loop Action']['main'][0] = [{
        "node": "Prepare Batch Update Input",
        "type": "main",
        "index": 0
    }]
    print("  - Route Batch Loop Action [update] -> Prepare Batch Update Input")

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

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

    # Handle Batch Update Result -> Prepare Next Iteration (same as other actions)
    workflow['connections']['Handle Batch Update Result'] = {
        'main': [[{
            "node": "Prepare Next Iteration",
            "type": "main",
            "index": 0
        }]]
    }
    print("  - Handle Batch Update Result -> Prepare Next Iteration")

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

    save_workflow(workflow)
    print("\n✓ Task 1 complete: Batch update now uses Container Update sub-workflow")

if __name__ == '__main__':
    main()