#!/usr/bin/env python3

import os
import requests
import json
from urllib.parse import urljoin

# Get API token
api_token = os.environ.get('NETBOX_KEY', 'e50298f7fd20f7fd6f1931f635511b34f6e8cfde')

# NetBox API setup
base_url = 'https://netbox.vntx.net/'
api_url = urljoin(base_url, 'api/')
headers = {
    'Authorization': f'Token {api_token}',
    'Content-Type': 'application/json',
    'Accept': 'application/json'
}

print("=== Updating 380 Edge Router in NetBox ===\n")

# Edge router specific data
edge_device_id = 9
site_id = 8

# Update primary IP assignment for edge router
print("Assigning primary IP to edge router...")
try:
    # First create interface on edge router
    interface_data = {
        'device': edge_device_id,
        'name': 'loopback',
        'type': 'virtual',
        'enabled': True
    }
    response = requests.post(f"{api_url}dcim/interfaces/", headers=headers, json=interface_data)
    if response.status_code == 201:
        interface_id = response.json()['id']
        print(f"✓ Created loopback interface (ID: {interface_id})")
        
        # Assign IP to interface
        ip_update = {
            'assigned_object_type': 'dcim.interface',
            'assigned_object_id': interface_id
        }
        # Find the IP
        ip_response = requests.get(f"{api_url}ipam/ip-addresses/", 
                                 headers=headers, 
                                 params={'address': '10.254.254.254/32'})
        if ip_response.json()['count'] > 0:
            ip_id = ip_response.json()['results'][0]['id']
            patch_response = requests.patch(f"{api_url}ipam/ip-addresses/{ip_id}/", 
                                          headers=headers, 
                                          json=ip_update)
            if patch_response.status_code == 200:
                print("✓ Assigned IP to loopback interface")
                
                # Set as primary IP for device
                device_update = {'primary_ip4': ip_id}
                device_response = requests.patch(f"{api_url}dcim/devices/{edge_device_id}/", 
                                               headers=headers, 
                                               json=device_update)
                if device_response.status_code == 200:
                    print("✓ Set as primary IP for edge router")
except Exception as e:
    print(f"Error: {e}")

# Create interfaces for edge router
print("\n=== Creating Edge Router Interfaces ===")
interfaces = [
    'sfp-sfpplus7-core-direct',
    'sfp-sfpplus8-server-switch', 
    'sfp-sfpplus11-preseem',
    'sfp-sfpplus12-spectrum',
    'cgnat',
    'vlan9_sfpplus8'
]

for interface in interfaces:
    # Check if already exists
    check = requests.get(f"{api_url}dcim/interfaces/", 
                        headers=headers,
                        params={'device_id': edge_device_id, 'name': interface})
    
    if check.json()['count'] == 0:
        if 'vlan' in interface or interface == 'cgnat':
            itype = 'virtual'
        else:
            itype = '10gbase-x-sfpp'  # SFP+ interfaces
            
        interface_data = {
            'device': edge_device_id,
            'name': interface,
            'type': itype,
            'enabled': True
        }
        
        response = requests.post(f"{api_url}dcim/interfaces/", headers=headers, json=interface_data)
        if response.status_code == 201:
            print(f"✓ Created interface: {interface}")
        else:
            print(f"✗ Failed to create interface: {interface}")

print("\n✓ Edge router update complete!")