#!/usr/bin/env python3

import requests
import json
from urllib.parse import urljoin

class NetBoxUpdater:
    def __init__(self, url, token):
        self.base_url = url.rstrip('/')
        self.api_url = urljoin(self.base_url + '/', 'api/')
        self.headers = {
            'Authorization': f'Token {token}',
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
        self.session = requests.Session()
        self.session.headers.update(self.headers)
    
    def post(self, endpoint, data):
        """Make POST request to NetBox API"""
        url = urljoin(self.api_url, endpoint.lstrip('/'))
        response = self.session.post(url, json=data)
        if response.status_code not in [200, 201]:
            print(f"Error creating {endpoint}: {response.status_code}")
            print(f"Response: {response.text}")
        response.raise_for_status()
        return response.json()
    
    def patch(self, endpoint, data):
        """Make PATCH request to NetBox API"""
        url = urljoin(self.api_url, endpoint.lstrip('/'))
        response = self.session.patch(url, json=data)
        if response.status_code not in [200, 201]:
            print(f"Error updating {endpoint}: {response.status_code}")
            print(f"Response: {response.text}")
        response.raise_for_status()
        return response.json()
    
    def get(self, endpoint, params=None):
        """Make GET request to NetBox API"""
        url = urljoin(self.api_url, endpoint.lstrip('/'))
        response = self.session.get(url, params=params)
        response.raise_for_status()
        return response.json()


def main():
    # NetBox connection
    import os
    api_token = os.environ.get('NETBOX_KEY', 'e50298f7fd20f7fd6f1931f635511b34f6e8cfde')
    nb = NetBoxUpdater('https://netbox.vntx.net/', api_token)
    
    # Router subnet data with interface mappings
    router_data = [
        {
            'prefix': '10.250.1.24/29',
            'ip': '10.250.1.25/29',
            'interface': 'ether3-climax-11ghz',
            'description': 'Climax 11GHz backhaul link',
            'role': 'infrastructure'
        },
        {
            'prefix': '10.254.254.101/32',
            'ip': '10.254.254.101/32',
            'interface': 'loopback',
            'description': 'Verona router loopback',
            'role': 'loopback'
        },
        {
            'prefix': '100.64.0.0/22',
            'ip': '100.64.3.254/22',
            'interface': 'verona',
            'description': 'Verona main CGNAT subnet',
            'role': 'customer'
        },
        {
            'prefix': '204.110.188.224/27',
            'ip': '204.110.188.254/27',
            'interface': 'verona',
            'description': 'Verona public IP subnet',
            'role': 'customer'
        },
        {
            'prefix': '100.64.12.0/22',
            'ip': '100.64.15.254/22',
            'interface': 'vlan_19_ether6',
            'description': 'Verona VLAN 19 CGNAT',
            'role': 'customer'
        },
        {
            'prefix': '10.10.80.0/20',
            'ip': '10.10.95.254/20',
            'interface': 'vlan_10_ether6',
            'description': 'Verona CPE management subnet 1',
            'role': 'management'
        },
        {
            'prefix': '10.10.0.0/20',
            'ip': '10.10.15.254/20',
            'interface': 'vlan_10_ether6',
            'description': 'Verona CPE management subnet 2',
            'role': 'management'
        },
        {
            'prefix': '10.0.101.0/24',
            'ip': '10.0.101.254/24',
            'interface': 'sfp-sfpplus1-verona-tower-switch',
            'description': 'Verona tower switch connection',
            'role': 'infrastructure'
        },
        {
            'prefix': '10.250.1.144/29',
            'ip': '10.250.1.145/29',
            'interface': 'ether6-switch',
            'description': 'Verona switch interconnect',
            'role': 'infrastructure'
        },
        {
            'prefix': '204.110.191.0/27',
            'ip': '204.110.191.30/27',
            'interface': 'vlan9_sfpplus1',
            'description': 'Verona VLAN 9 public subnet',
            'role': 'customer'
        },
        {
            'prefix': '10.25.1.0/24',
            'ip': '10.25.1.254/24',
            'interface': 'ether1',
            'description': 'Verona ether1 management',
            'role': 'management'
        },
        {
            'prefix': '192.168.99.0/24',
            'ip': '192.168.99.254/24',
            'interface': 'ether10-powerswitch',
            'description': 'Verona power switch management',
            'role': 'management'
        }
    ]
    
    # Get site ID
    site_response = nb.get('dcim/sites/', params={'name': 'Verona'})
    if site_response['count'] == 0:
        print("Error: Verona site not found")
        return
    
    site_id = site_response['results'][0]['id']
    print(f"Found Verona site with ID: {site_id}")
    
    # Get or create prefix roles
    role_mapping = {
        'infrastructure': 'Infrastructure',
        'loopback': 'Loopback',
        'customer': 'Customer',
        'management': 'Management'
    }
    
    role_ids = {}
    for key, name in role_mapping.items():
        role_response = nb.get('ipam/roles/', params={'name': name})
        if role_response['count'] > 0:
            role_ids[key] = role_response['results'][0]['id']
        else:
            # Create role if it doesn't exist
            role_data = {
                'name': name,
                'slug': key
            }
            try:
                created_role = nb.post('ipam/roles/', role_data)
                role_ids[key] = created_role['id']
                print(f"Created role: {name}")
            except:
                print(f"Could not create role: {name}")
                role_ids[key] = None
    
    # Process each subnet
    print("\n=== Creating/Updating Prefixes ===")
    created_prefixes = 0
    failed_prefixes = 0
    
    for item in router_data:
        # Check if prefix already exists
        prefix_response = nb.get('ipam/prefixes/', params={'prefix': item['prefix']})
        
        prefix_data = {
            'prefix': item['prefix'],
            'site': site_id,
            'status': 'active',
            'description': item['description'],
            'is_pool': False
        }
        
        if role_ids.get(item['role']):
            prefix_data['role'] = role_ids[item['role']]
        
        try:
            if prefix_response['count'] == 0:
                # Create new prefix
                created_prefix = nb.post('ipam/prefixes/', prefix_data)
                print(f"Created prefix: {item['prefix']} - {item['description']}")
                created_prefixes += 1
            else:
                # Update existing prefix
                existing_id = prefix_response['results'][0]['id']
                updated_prefix = nb.patch(f'ipam/prefixes/{existing_id}/', prefix_data)
                print(f"Updated prefix: {item['prefix']} - {item['description']}")
                created_prefixes += 1
        except Exception as e:
            print(f"Failed to create/update prefix {item['prefix']}: {e}")
            failed_prefixes += 1
    
    print(f"\nPrefix Summary: {created_prefixes} successful, {failed_prefixes} failed")
    
    # Update IP addresses with interface information
    print("\n=== Updating IP Address Interface Assignments ===")
    updated_ips = 0
    failed_ips = 0
    
    # Get the device for Verona
    device_response = nb.get('dcim/devices/', params={'site_id': site_id})
    if device_response['count'] == 0:
        print("Warning: No device found for Verona site")
        device_id = None
    else:
        device_id = device_response['results'][0]['id']
        device_name = device_response['results'][0]['name']
        print(f"Found device: {device_name} (ID: {device_id})")
    
    # Create a mapping of IP to interface
    ip_to_interface = {item['ip']: item['interface'] for item in router_data}
    ip_to_description = {item['ip']: item['description'] for item in router_data}
    
    # Get all IPs for the site
    ips_response = nb.get('ipam/ip-addresses/', params={'site_id': site_id})
    
    for ip_obj in ips_response['results']:
        ip_address = ip_obj['address']
        
        if ip_address in ip_to_interface:
            interface_name = ip_to_interface[ip_address]
            description = ip_to_description[ip_address]
            
            # Update IP with description
            update_data = {
                'description': f"{interface_name} - {description}"
            }
            
            # If we have a device, try to find or create the interface
            if device_id:
                # Check if interface exists
                interface_response = nb.get('dcim/interfaces/', params={
                    'device_id': device_id,
                    'name': interface_name
                })
                
                if interface_response['count'] == 0:
                    # Create interface
                    interface_data = {
                        'device': device_id,
                        'name': interface_name,
                        'type': 'virtual',  # Using virtual for now
                        'enabled': True
                    }
                    try:
                        created_interface = nb.post('dcim/interfaces/', interface_data)
                        interface_id = created_interface['id']
                        print(f"Created interface: {interface_name}")
                        
                        # Assign IP to interface
                        update_data['assigned_object_type'] = 'dcim.interface'
                        update_data['assigned_object_id'] = interface_id
                    except Exception as e:
                        print(f"Failed to create interface {interface_name}: {e}")
                else:
                    # Use existing interface
                    interface_id = interface_response['results'][0]['id']
                    update_data['assigned_object_type'] = 'dcim.interface'
                    update_data['assigned_object_id'] = interface_id
            
            # Update the IP address
            try:
                updated_ip = nb.patch(f"ipam/ip-addresses/{ip_obj['id']}/", update_data)
                print(f"Updated IP {ip_address} with interface {interface_name}")
                updated_ips += 1
            except Exception as e:
                print(f"Failed to update IP {ip_address}: {e}")
                failed_ips += 1
    
    print(f"\nIP Update Summary: {updated_ips} successful, {failed_ips} failed")
    
    print("\n=== Update Complete ===")
    print(f"Total prefixes processed: {created_prefixes + failed_prefixes}")
    print(f"Total IPs processed: {updated_ips + failed_ips}")


if __name__ == "__main__":
    main()