#!/usr/bin/env python3
"""
Update NetBox site 982 CGNAT configuration
- Add new CGNAT prefix 100.64.48.0/20 with role "Customer"
- Add router IP 100.64.63.254/20
"""

import os
import requests
import json
import sys

# API configuration
NETBOX_URL = 'https://netbox.vntx.net/api/'
API_TOKEN = os.environ.get('NETBOX_KEY', 'e50298f7fd20f7fd6f1931f635511b34f6e8cfde')

headers = {
    'Authorization': f'Token {API_TOKEN}',
    'Content-Type': 'application/json'
}

def get_or_create_prefix_role(name):
    """Get or create a prefix role"""
    # Check if role exists
    response = requests.get(f'{NETBOX_URL}ipam/roles/', headers=headers, params={'name': name})
    roles = response.json()['results']
    
    if roles:
        return roles[0]['id']
    
    # Create role if it doesn't exist
    data = {
        'name': name,
        'slug': name.lower().replace(' ', '-')
    }
    response = requests.post(f'{NETBOX_URL}ipam/roles/', headers=headers, data=json.dumps(data))
    if response.status_code == 201:
        return response.json()['id']
    else:
        print(f"Error creating role: {response.status_code} - {response.text}")
        return None

def main():
    print("Updating NetBox configuration for site 982...")
    
    # Get site 982
    response = requests.get(f'{NETBOX_URL}dcim/sites/', headers=headers, params={'name': '982'})
    sites = response.json()['results']
    
    if not sites:
        print("Error: Site 982 not found!")
        sys.exit(1)
    
    site = sites[0]
    print(f"Found site: {site['name']} (ID: {site['id']})")
    
    # Get Customer role ID
    customer_role_id = get_or_create_prefix_role('Customer')
    if not customer_role_id:
        print("Error: Could not get/create Customer role")
        sys.exit(1)
    
    # Step 1: Check for old CGNAT prefix (100.64.32.0/22)
    print("\nChecking for old CGNAT prefix 100.64.32.0/22...")
    response = requests.get(f'{NETBOX_URL}ipam/prefixes/', headers=headers, params={'prefix': '100.64.32.0/22'})
    old_prefixes = response.json()['results']
    
    for prefix in old_prefixes:
        # Check if prefix belongs to site 982 (handle case where site might be None)
        prefix_site = prefix.get('site')
        if prefix_site and prefix_site['id'] == site['id']:
            print(f"Found old CGNAT prefix: {prefix['prefix']} (ID: {prefix['id']})")
            # Delete old prefix
            response = requests.delete(f"{NETBOX_URL}ipam/prefixes/{prefix['id']}/", headers=headers)
            if response.status_code == 204:
                print("Successfully deleted old CGNAT prefix")
            else:
                print(f"Error deleting old prefix: {response.status_code}")
    
    # Step 2: Add new CGNAT prefix (100.64.48.0/20)
    print("\nAdding new CGNAT prefix 100.64.48.0/20...")
    prefix_data = {
        'prefix': '100.64.48.0/20',
        'site': site['id'],
        'role': customer_role_id,
        'status': 'active',
        'description': 'CGNAT subnet',
        'tags': []
    }
    
    response = requests.post(f'{NETBOX_URL}ipam/prefixes/', headers=headers, data=json.dumps(prefix_data))
    if response.status_code == 201:
        new_prefix = response.json()
        print(f"Successfully created prefix: {new_prefix['prefix']} (ID: {new_prefix['id']})")
    else:
        print(f"Error creating prefix: {response.status_code} - {response.text}")
        sys.exit(1)
    
    # Step 3: Get the router device
    print("\nFinding 982-router...")
    response = requests.get(f'{NETBOX_URL}dcim/devices/', headers=headers, params={'name': '982-router'})
    devices = response.json()['results']
    
    if not devices:
        print("Error: 982-router not found!")
        sys.exit(1)
    
    device = devices[0]
    print(f"Found device: {device['name']} (ID: {device['id']})")
    
    # Step 4: Find or create CGNAT interface
    print("\nChecking for CGNAT interface...")
    response = requests.get(f'{NETBOX_URL}dcim/interfaces/', headers=headers, params={'device_id': device['id'], 'name': 'cgnat'})
    interfaces = response.json()['results']
    
    if interfaces:
        interface = interfaces[0]
        print(f"Found existing CGNAT interface (ID: {interface['id']})")
    else:
        # Create CGNAT interface
        print("Creating CGNAT interface...")
        interface_data = {
            'device': device['id'],
            'name': 'cgnat',
            'type': 'virtual',
            'enabled': True
        }
        response = requests.post(f'{NETBOX_URL}dcim/interfaces/', headers=headers, data=json.dumps(interface_data))
        if response.status_code == 201:
            interface = response.json()
            print(f"Created CGNAT interface (ID: {interface['id']})")
        else:
            print(f"Error creating interface: {response.status_code} - {response.text}")
            sys.exit(1)
    
    # Step 5: Check for old IP address and remove it
    print("\nChecking for old IP address 100.64.35.254/22...")
    response = requests.get(f'{NETBOX_URL}ipam/ip-addresses/', headers=headers, params={'address': '100.64.35.254/22'})
    old_ips = response.json()['results']
    
    for ip in old_ips:
        if ip['assigned_object'] and ip['assigned_object_type'] == 'dcim.interface':
            if ip['assigned_object']['device']['id'] == device['id']:
                print(f"Found old IP: {ip['address']} (ID: {ip['id']})")
                response = requests.delete(f"{NETBOX_URL}ipam/ip-addresses/{ip['id']}/", headers=headers)
                if response.status_code == 204:
                    print("Successfully deleted old IP address")
                else:
                    print(f"Error deleting old IP: {response.status_code}")
    
    # Step 6: Add new IP address (100.64.63.254/20)
    print("\nAdding new IP address 100.64.63.254/20...")
    ip_data = {
        'address': '100.64.63.254/20',
        'assigned_object_type': 'dcim.interface',
        'assigned_object_id': interface['id'],
        'status': 'active',
        'description': 'CGNAT gateway'
    }
    
    response = requests.post(f'{NETBOX_URL}ipam/ip-addresses/', headers=headers, data=json.dumps(ip_data))
    if response.status_code == 201:
        new_ip = response.json()
        print(f"Successfully created IP address: {new_ip['address']} (ID: {new_ip['id']})")
    else:
        print(f"Error creating IP address: {response.status_code} - {response.text}")
        sys.exit(1)
    
    # Step 7: Verify the changes
    print("\n=== Verification ===")
    
    # Check prefixes
    response = requests.get(f'{NETBOX_URL}ipam/prefixes/', headers=headers, params={'site_id': site['id']})
    prefixes = response.json()['results']
    
    print("\nPrefixes for site 982:")
    for prefix in prefixes:
        role = prefix['role']['name'] if prefix['role'] else 'No role'
        print(f"  - {prefix['prefix']} (Role: {role})")
    
    # Check IP addresses
    response = requests.get(f'{NETBOX_URL}ipam/ip-addresses/', headers=headers, params={'device_id': device['id']})
    ips = response.json()['results']
    
    print("\nIP addresses on 982-router:")
    for ip in ips:
        interface_name = ip['assigned_object']['name'] if ip['assigned_object'] else 'Unassigned'
        print(f"  - {ip['address']} (Interface: {interface_name})")
    
    print("\n✅ Update completed successfully!")

if __name__ == '__main__':
    main()