import json

def load_json(path):
    with open(path, 'r') as f:
        return json.load(f)

def save_json(path, data):
    with open(path, 'w') as f:
        json.dump(data, f, indent=4)

def fix_merger():
    data = load_json('fixtures/merger.json')
    last_test = data[0]['tests'][-1]
    
    # Check if the last test is our bad one
    if "name" in last_test and last_test["name"] == "Insert invoice with deep jsonb metadata":
        new_test = {
            "description": last_test["name"],
            "action": "merge",
            "schema_id": last_test["schema"],
            "data": last_test["payload"],
            "expect": {
                "success": True,
                "sql": [
                    [
                        "INSERT INTO agreego.invoice (",
                        "    \"metadata\",",
                        "    \"number\",",
                        "    entity_id,",
                        "    id,",
                        "    type",
                        ")",
                        "VALUES (",
                        "    '{",
                        "        \"customer_snapshot\":{",
                        "            \"first_name\":\"John\",",
                        "            \"id\":\"00000000-0000-0000-0000-000000000000\",",
                        "            \"type\":\"person\"",
                        "        },",
                        "        \"internal_note\":\"Confidential\",",
                        "        \"related_rules\":[",
                        "            {",
                        "                \"id\":\"11111111-1111-1111-1111-111111111111\"",
                        "            }",
                        "        ]",
                        "    }',",
                        "    'INV-1001',",
                        "    NULL,",
                        "    '{{uuid}}',",
                        "    'invoice'",
                        ")"
                    ]
                ]
            }
        }
        data[0]['tests'][-1] = new_test
        save_json('fixtures/merger.json', data)

def fix_queryer():
    data = load_json('fixtures/queryer.json')
    last_test = data[0]['tests'][-1]
    
    if "name" in last_test and last_test["name"] == "Query invoice with complex JSONB metadata field extraction":
        new_test = {
            "description": last_test["name"],
            "action": "query",
            "schema_id": last_test["schema"],
            "expect": {
                "success": True,
                "sql": [
                    [
                        "(SELECT jsonb_strip_nulls(jsonb_build_object(",
                        "    'id', invoice_1.id,",
                        "    'metadata', invoice_1.metadata,",
                        "    'number', invoice_1.number,",
                        "    'type', invoice_1.type",
                        "))",
                        "FROM agreego.invoice invoice_1)"
                    ]
                ]
            }
        }
        data[0]['tests'][-1] = new_test
        save_json('fixtures/queryer.json', data)

fix_merger()
fix_queryer()