#!/usr/bin/env python3

import os
import sys
import json
import argparse
import requests
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 get_or_create_role(self, name, slug=None):
        """Get or create an IPAM role"""
        if not slug:
            slug = name.lower().replace(' ', '-')
            
        response = self.get('ipam/roles/', params={'slug': slug})
        if response['count'] > 0:
            return response['results'][0]['id']
        
        # Create role
        role_data = {
            'name': name,
            'slug': slug
        }
        created = self.post('ipam/roles/', role_data)
        return created['id']


def determine_prefix_role(interface, description=''):
    """Determine the role of a prefix based on interface name and description"""
    interface_lower = str(interface).lower()
    
    # Infrastructure links
    if any(term in interface_lower for term in ['airfiber', '11ghz', '24ghz', 'backhaul', 'ether3', 'ether4', 'ether5', 'ether6']):
        return 'infrastructure'
    
    # Loopback
    if 'loopback' in interface_lower or interface == 'lo':
        return 'loopback'
    
    # Management
    if any(term in interface_lower for term in ['mgmt', 'management', 'ether1']):
        return 'management'
    
    # Customer bridges
    if 'bridge' in interface_lower or 'pppoe' in interface_lower:
        return 'customer'
    
    # VLANs are often customer-facing
    if 'vlan' in interface_lower:
        return 'customer'
    
    return 'infrastructure'  # Default


def process_router_data(json_file, site_name, dry_run=False):
    """Process router data JSON and update NetBox"""
    
    # Load router data
    with open(json_file, 'r') as f:
        router_data = json.load(f)
    
    # Get API token
    api_token = os.environ.get('NETBOX_KEY', 'e50298f7fd20f7fd6f1931f635511b34f6e8cfde')
    nb = NetBoxUpdater('https://netbox.vntx.net/', api_token)
    
    print(f"=== Processing Router Data for {site_name} ===\n")
    
    # Get site ID
    site_response = nb.get('dcim/sites/', params={'name': site_name})
    if site_response['count'] == 0:
        print(f"Error: Site '{site_name}' not found in NetBox")
        print(f"Please create the site first using: python3 create_site_and_router.py {site_name} {router_data['host']}")
        return False
    
    site_id = site_response['results'][0]['id']
    print(f"Found site: {site_name} (ID: {site_id})")
    
    # Get device
    device_response = nb.get('dcim/devices/', params={'site_id': site_id})
    if device_response['count'] == 0:
        print(f"Error: No device found for site '{site_name}'")
        return False
    
    device = device_response['results'][0]
    device_id = device['id']
    device_name = device['name']
    print(f"Found device: {device_name} (ID: {device_id})\n")
    
    # Get or create roles
    role_mapping = {
        'infrastructure': 'Infrastructure',
        'loopback': 'Loopback',
        'customer': 'Customer',
        'management': 'Management'
    }
    
    if not dry_run:
        print("Setting up IPAM roles...")
        role_ids = {}
        for key, name in role_mapping.items():
            role_ids[key] = nb.get_or_create_role(name, key)
            print(f"  ✓ {name}")
        print()
    
    # Process subnets (non-dynamic only)
    static_subnets = [s for s in router_data['subnets'] if not s.get('dynamic', False)]
    
    print(f"=== Processing {len(static_subnets)} Static Subnets ===")
    
    prefixes_to_create = []
    ips_to_create = []
    
    for subnet in static_subnets:
        address = subnet['address']
        network = subnet['network']
        interface = subnet['interface']
        comment = subnet.get('comment', '')
        
        # Skip PPPoE dynamic IPs
        if str(interface).startswith('<pppoe-'):
            continue
        
        # Determine if it's a host IP or a subnet
        if address.endswith('/32'):
            # It's a host IP
            role = determine_prefix_role(interface, comment)
            ips_to_create.append({
                'address': address,
                'interface': interface,
                'description': comment or f"{device_name} - {interface}",
                'role': role
            })
        else:
            # It's a subnet
            prefix_bits = address.split('/')[-1]
            prefix = f"{network}/{prefix_bits}"
            role = determine_prefix_role(interface, comment)
            
            prefixes_to_create.append({
                'prefix': prefix,
                'interface': interface,
                'description': comment or f"{site_name} - {interface}",
                'role': role,
                'gateway_ip': address
            })
    
    if dry_run:
        print("\n=== DRY RUN - Would create the following: ===\n")
        
        print(f"Prefixes ({len(prefixes_to_create)}):")
        for p in prefixes_to_create:
            print(f"  - {p['prefix']} ({p['role']}) - {p['description']}")
        
        print(f"\nIP Addresses ({len(ips_to_create)}):")
        for ip in ips_to_create:
            print(f"  - {ip['address']} ({ip['role']}) - {ip['description']}")
        
        return True
    
    # Create prefixes
    print(f"\n=== Creating/Updating {len(prefixes_to_create)} Prefixes ===")
    created_prefixes = 0
    failed_prefixes = 0
    
    for item in prefixes_to_create:
        # 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
                
            # Create gateway IP if needed
            if item.get('gateway_ip'):
                gw_response = nb.get('ipam/ip-addresses/', params={'address': item['gateway_ip']})
                if gw_response['count'] == 0:
                    gw_data = {
                        'address': item['gateway_ip'],
                        'status': 'active',
                        'description': f"{item['interface']} gateway - {item['description']}",
                        'role': 'anycast' if item['role'] == 'infrastructure' else None
                    }
                    try:
                        nb.post('ipam/ip-addresses/', gw_data)
                        print(f"  ✓ Created gateway IP: {item['gateway_ip']}")
                    except:
                        pass
                        
        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")
    
    # Create IP addresses
    print(f"\n=== Creating/Updating {len(ips_to_create)} IP Addresses ===")
    created_ips = 0
    failed_ips = 0
    
    for item in ips_to_create:
        # Check if IP already exists
        ip_response = nb.get('ipam/ip-addresses/', params={'address': item['address']})
        
        ip_data = {
            'address': item['address'],
            'status': 'active',
            'description': item['description']
        }
        
        if item['role'] == 'loopback':
            ip_data['role'] = 'loopback'
        
        try:
            if ip_response['count'] == 0:
                # Create new IP
                created_ip = nb.post('ipam/ip-addresses/', ip_data)
                print(f"✓ Created IP: {item['address']} - {item['description']}")
                created_ips += 1
            else:
                # Update existing IP
                existing_id = ip_response['results'][0]['id']
                updated_ip = nb.patch(f'ipam/ip-addresses/{existing_id}/', ip_data)
                print(f"✓ Updated IP: {item['address']} - {item['description']}")
                created_ips += 1
                
        except Exception as e:
            print(f"✗ Failed to create/update IP {item['address']}: {e}")
            failed_ips += 1
    
    print(f"\nIP Summary: {created_ips} successful, {failed_ips} failed")
    
    # Create interfaces on the device
    print(f"\n=== Creating Device Interfaces ===")
    created_interfaces = 0
    
    # Get unique interfaces from both subnets and active interfaces
    interface_set = set()
    for subnet in static_subnets:
        if not str(subnet['interface']).startswith('<'):
            interface_set.add(str(subnet['interface']))
    
    for iface in router_data.get('interfaces', []):
        if not str(iface['name']).startswith('<'):
            interface_set.add(str(iface['name']))
    
    for interface_name in sorted(interface_set):
        # Check if interface exists
        interface_response = nb.get('dcim/interfaces/', params={
            'device_id': device_id,
            'name': interface_name
        })
        
        if interface_response['count'] == 0:
            # Determine interface type
            if 'vlan' in interface_name.lower():
                iface_type = 'virtual'
            elif 'bridge' in interface_name.lower():
                iface_type = 'bridge'
            elif 'loopback' in interface_name.lower() or interface_name == 'lo':
                iface_type = 'virtual'
            else:
                iface_type = '1000base-t'  # Default to gigabit ethernet
            
            interface_data = {
                'device': device_id,
                'name': interface_name,
                'type': iface_type,
                'enabled': True
            }
            
            try:
                created_interface = nb.post('dcim/interfaces/', interface_data)
                print(f"✓ Created interface: {interface_name} ({iface_type})")
                created_interfaces += 1
            except Exception as e:
                print(f"✗ Failed to create interface {interface_name}: {e}")
    
    print(f"\nCreated {created_interfaces} new interfaces")
    
    print(f"\n=== Update Complete ===")
    print(f"Site: {site_name}")
    print(f"Device: {device_name}")
    print(f"Prefixes: {created_prefixes}/{len(prefixes_to_create)}")
    print(f"IPs: {created_ips}/{len(ips_to_create)}")
    print(f"Interfaces: {created_interfaces}")
    
    return True


def main():
    parser = argparse.ArgumentParser(description='Update NetBox with router data from JSON file')
    parser.add_argument('json_file', help='Path to router data JSON file')
    parser.add_argument('site_name', help='Name of the site in NetBox')
    parser.add_argument('--dry-run', action='store_true', help='Show what would be created without making changes')
    
    args = parser.parse_args()
    
    # Check if file exists
    if not os.path.exists(args.json_file):
        print(f"Error: File '{args.json_file}' not found")
        return 1
    
    # Process the data
    success = process_router_data(args.json_file, args.site_name, args.dry_run)
    
    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())