#!/usr/bin/env python3
"""
Complete update script for site 982 CGNAT configuration
This script handles the complete update process including verification
"""

import os
import requests
import json
import sys
import time

# 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',
    'Accept': 'application/json'
}

def delete_prefix_by_cidr(cidr):
    """Delete a prefix by CIDR notation"""
    response = requests.get(f'{NETBOX_URL}ipam/prefixes/', headers=headers, params={'prefix': cidr})
    prefixes = response.json()['results']
    
    deleted = False
    for prefix in prefixes:
        print(f"Deleting prefix {prefix['prefix']} (ID: {prefix['id']})")
        response = requests.delete(f"{NETBOX_URL}ipam/prefixes/{prefix['id']}/", headers=headers)
        if response.status_code == 204:
            print("  ✓ Deleted successfully")
            deleted = True
        else:
            print(f"  ✗ Error: {response.status_code}")
    
    return deleted

def delete_ip_by_address(address):
    """Delete an IP address by address string"""
    response = requests.get(f'{NETBOX_URL}ipam/ip-addresses/', headers=headers, params={'address': address})
    ips = response.json()['results']
    
    deleted = False
    for ip in ips:
        print(f"Deleting 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("  ✓ Deleted successfully")
            deleted = True
        else:
            print(f"  ✗ Error: {response.status_code}")
    
    return deleted

def main():
    print("=== Site 982 CGNAT Update Script ===\n")
    
    # Step 1: Get site 982
    print("1. Finding site 982...")
    response = requests.get(f'{NETBOX_URL}dcim/sites/', headers=headers, params={'name': '982'})
    sites = response.json()['results']
    
    if not sites:
        print("✗ Site 982 not found!")
        sys.exit(1)
    
    site = sites[0]
    site_id = site['id']
    print(f"✓ Found site: {site['name']} (ID: {site_id})")
    
    # Step 2: Get the router
    print("\n2. Finding 982-router...")
    response = requests.get(f'{NETBOX_URL}dcim/devices/', headers=headers, params={'name': '982-router', 'site_id': site_id})
    devices = response.json()['results']
    
    if not devices:
        print("✗ 982-router not found!")
        sys.exit(1)
    
    device = devices[0]
    device_id = device['id']
    print(f"✓ Found device: {device['name']} (ID: {device_id})")
    
    # Step 3: Clean up old configurations
    print("\n3. Cleaning up old configurations...")
    
    # Delete old CGNAT prefix
    print("   Removing old CGNAT prefix 100.64.32.0/22...")
    delete_prefix_by_cidr('100.64.32.0/22')
    
    # Delete any existing 100.64.48.0/20 prefix (from previous attempts)
    print("   Removing any existing 100.64.48.0/20 prefix...")
    delete_prefix_by_cidr('100.64.48.0/20')
    
    # Delete old IP address
    print("   Removing old IP 100.64.35.254/22...")
    delete_ip_by_address('100.64.35.254/22')
    
    # Delete any existing new IP (from previous attempts)
    print("   Removing any existing 100.64.63.254/20...")
    delete_ip_by_address('100.64.63.254/20')
    
    # Step 4: Create or find CGNAT interface
    print("\n4. Setting up 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]
        interface_id = interface['id']
        print(f"✓ Found existing CGNAT interface (ID: {interface_id})")
    else:
        # Create interface
        interface_data = {
            'device': device_id,
            'name': 'cgnat',
            'type': 'virtual',
            'enabled': True,
            'description': 'CGNAT interface'
        }
        response = requests.post(f'{NETBOX_URL}dcim/interfaces/', headers=headers, json=interface_data)
        if response.status_code == 201:
            interface = response.json()
            interface_id = interface['id']
            print(f"✓ Created CGNAT interface (ID: {interface_id})")
        else:
            print(f"✗ Error creating interface: {response.status_code}")
            print(response.text)
            sys.exit(1)
    
    # Step 5: Get or create Customer role
    print("\n5. Setting up Customer role...")
    response = requests.get(f'{NETBOX_URL}ipam/roles/', headers=headers, params={'name': 'Customer'})
    roles = response.json()['results']
    
    if roles:
        role = roles[0]
        role_id = role['id']
        print(f"✓ Found Customer role (ID: {role_id})")
    else:
        # Create role
        role_data = {
            'name': 'Customer',
            'slug': 'customer'
        }
        response = requests.post(f'{NETBOX_URL}ipam/roles/', headers=headers, json=role_data)
        if response.status_code == 201:
            role = response.json()
            role_id = role['id']
            print(f"✓ Created Customer role (ID: {role_id})")
        else:
            print(f"✗ Error creating role: {response.status_code}")
            sys.exit(1)
    
    # Step 6: Create the new CGNAT prefix
    print("\n6. Creating new CGNAT prefix 100.64.48.0/20...")
    
    # Try creating without site first, then update
    prefix_data = {
        'prefix': '100.64.48.0/20',
        'role': role_id,
        'status': 'active',
        'description': 'CGNAT subnet for site 982',
        'tags': []
    }
    
    response = requests.post(f'{NETBOX_URL}ipam/prefixes/', headers=headers, json=prefix_data)
    if response.status_code == 201:
        prefix = response.json()
        prefix_id = prefix['id']
        print(f"✓ Created prefix (ID: {prefix_id})")
        
        # Now try to assign site
        print("   Assigning prefix to site 982...")
        # Try different approaches
        
        # Approach 1: PATCH with just site
        patch_data = {'site': site_id}
        response = requests.patch(f'{NETBOX_URL}ipam/prefixes/{prefix_id}/', headers=headers, json=patch_data)
        
        if response.status_code != 200:
            # Approach 2: PUT with all fields
            put_data = {
                'prefix': '100.64.48.0/20',
                'site': site_id,
                'role': role_id,
                'status': 'active',
                'description': 'CGNAT subnet for site 982'
            }
            response = requests.put(f'{NETBOX_URL}ipam/prefixes/{prefix_id}/', headers=headers, json=put_data)
        
        if response.status_code in [200, 201]:
            updated_prefix = response.json()
            if updated_prefix.get('site'):
                print(f"   ✓ Prefix assigned to site {updated_prefix['site']['name']}")
            else:
                print("   ⚠ Warning: Site assignment may have failed")
        else:
            print(f"   ⚠ Warning: Could not assign site: {response.status_code}")
    else:
        print(f"✗ Error creating prefix: {response.status_code}")
        print(response.text)
        sys.exit(1)
    
    # Step 7: Create the new IP address
    print("\n7. Creating 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 for site 982',
        'tags': []
    }
    
    response = requests.post(f'{NETBOX_URL}ipam/ip-addresses/', headers=headers, json=ip_data)
    if response.status_code == 201:
        ip = response.json()
        ip_id = ip['id']
        print(f"✓ Created IP address (ID: {ip_id})")
    else:
        print(f"✗ Error creating IP: {response.status_code}")
        print(response.text)
        sys.exit(1)
    
    # Step 8: Final verification
    print("\n=== VERIFICATION ===")
    
    # Check prefixes
    print("\nPrefixes with 100.64.48.0/20:")
    response = requests.get(f'{NETBOX_URL}ipam/prefixes/', headers=headers, params={'prefix': '100.64.48.0/20'})
    prefixes = response.json()['results']
    for prefix in prefixes:
        site_info = f"Site: {prefix['site']['name']}" if prefix.get('site') else "Site: Not assigned"
        role_info = f"Role: {prefix['role']['name']}" if prefix.get('role') else "Role: None"
        print(f"  - {prefix['prefix']} ({site_info}, {role_info})")
    
    # Check IPs
    print("\nIP addresses on 982-router:")
    response = requests.get(f'{NETBOX_URL}ipam/ip-addresses/', headers=headers, params={'device_id': device_id})
    ips = response.json()['results']
    for ip in ips:
        interface_info = f"Interface: {ip['assigned_object']['name']}" if ip.get('assigned_object') else "Unassigned"
        print(f"  - {ip['address']} ({interface_info})")
    
    print("\n✅ Update completed!")
    print("\nNOTE: If the prefix is not showing as assigned to site 982, this may be a")
    print("permission or API limitation. The prefix and IP have been created successfully.")

if __name__ == '__main__':
    main()