#!/usr/bin/env python3

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

class NetBoxManager:
    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 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 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 create_site(self, name, slug=None, status='active', comments='', physical_address=''):
        """Create a new site"""
        if not slug:
            slug = name.lower().replace(' ', '-').replace('_', '-')
        
        site_data = {
            'name': name,
            'slug': slug,
            'status': status,
            'comments': comments
        }
        
        if physical_address:
            site_data['physical_address'] = physical_address
        
        return self.post('dcim/sites/', site_data)
    
    def create_device(self, name, device_type_id, role_id, site_id, status='active', primary_ip=None):
        """Create a new device"""
        device_data = {
            'name': name,
            'device_type': device_type_id,
            'role': role_id,
            'site': site_id,
            'status': status
        }
        
        if primary_ip:
            device_data['primary_ip4'] = primary_ip
        
        return self.post('dcim/devices/', device_data)
    
    def create_ip_address(self, address, status='active', description='', role=None):
        """Create a new IP address"""
        ip_data = {
            'address': address,
            'status': status,
            'description': description
        }
        
        if role:
            ip_data['role'] = role
        
        return self.post('ipam/ip-addresses/', ip_data)
    
    def get_or_create_manufacturer(self, name):
        """Get or create a manufacturer"""
        response = self.get('dcim/manufacturers/', params={'name': name})
        if response['count'] > 0:
            return response['results'][0]['id']
        
        # Create manufacturer
        mfg_data = {
            'name': name,
            'slug': name.lower()
        }
        created = self.post('dcim/manufacturers/', mfg_data)
        return created['id']
    
    def get_or_create_device_type(self, model, manufacturer_id):
        """Get or create a device type"""
        response = self.get('dcim/device-types/', params={'model': model})
        if response['count'] > 0:
            return response['results'][0]['id']
        
        # Create device type
        dt_data = {
            'manufacturer': manufacturer_id,
            'model': model,
            'slug': model.lower().replace(' ', '-').replace('/', '-')
        }
        created = self.post('dcim/device-types/', dt_data)
        return created['id']
    
    def get_or_create_device_role(self, name, slug=None):
        """Get or create a device role"""
        if not slug:
            slug = name.lower().replace(' ', '-')
            
        response = self.get('dcim/device-roles/', params={'name': name})
        if response['count'] > 0:
            return response['results'][0]['id']
        
        # Create device role
        role_data = {
            'name': name,
            'slug': slug,
            'color': '2196f3'  # Blue color
        }
        created = self.post('dcim/device-roles/', role_data)
        return created['id']


def create_site_and_router(site_name, router_ip, router_name=None, manufacturer='MikroTik', 
                          device_type='RouterBOARD', device_role='Router', 
                          site_comments='', physical_address=''):
    """
    Create a site and router in NetBox
    
    Args:
        site_name: Name of the site (e.g., 'Verona', 'Climax')
        router_ip: Loopback IP of the router (e.g., '10.254.254.101')
        router_name: Name for the router device (defaults to '{site_name}-router')
        manufacturer: Device manufacturer (default: MikroTik)
        device_type: Device model/type (default: RouterBOARD)
        device_role: Role of the device (default: Router)
        site_comments: Comments for the site
        physical_address: Physical address of the site
    """
    # Get API token from environment variable or use the one from CLAUDE.md
    api_token = os.environ.get('NETBOX_KEY', 'e50298f7fd20f7fd6f1931f635511b34f6e8cfde')
    
    # Initialize NetBox client
    nb = NetBoxManager('https://netbox.vntx.net/', api_token)
    
    # Default router name if not provided
    if not router_name:
        router_name = f"{site_name.lower()}-router"
    
    # Default site comments if not provided
    if not site_comments:
        site_comments = f"{site_name} tower site with {manufacturer} router"
    
    print(f"=== Creating {site_name} Site and Router in NetBox ===\n")
    
    # Check if site exists
    site_response = nb.get('dcim/sites/', params={'name': site_name})
    
    if site_response['count'] == 0:
        # Create site
        print(f"Creating {site_name} site...")
        try:
            site = nb.create_site(
                name=site_name,
                slug=site_name.lower(),
                status='active',
                comments=site_comments,
                physical_address=physical_address
            )
            site_id = site['id']
            print(f"✓ Created {site_name} site (ID: {site_id})")
        except Exception as e:
            print(f"✗ Failed to create {site_name} site: {e}")
            return None, None
    else:
        site_id = site_response['results'][0]['id']
        print(f"✓ {site_name} site already exists (ID: {site_id})")
    
    # Check if router already exists
    device_response = nb.get('dcim/devices/', params={'name': router_name, 'site_id': site_id})
    
    if device_response['count'] == 0:
        print(f"\nCreating {router_name} device...")
        
        # Get or create manufacturer
        print(f"  - Setting up manufacturer ({manufacturer})...")
        manufacturer_id = nb.get_or_create_manufacturer(manufacturer)
        
        # Get or create device type
        print(f"  - Setting up device type ({device_type})...")
        device_type_id = nb.get_or_create_device_type(device_type, manufacturer_id)
        
        # Get or create device role
        print(f"  - Setting up device role ({device_role})...")
        role_id = nb.get_or_create_device_role(device_role)
        
        # Create the device
        try:
            device = nb.create_device(
                name=router_name,
                device_type_id=device_type_id,
                role_id=role_id,
                site_id=site_id,
                status='active'
            )
            device_id = device['id']
            print(f"✓ Created {router_name} device (ID: {device_id})")
            
            # Add primary IP
            print(f"\n  - Adding primary IP address ({router_ip}/32)...")
            
            # Check if IP already exists
            ip_response = nb.get('ipam/ip-addresses/', params={'address': f"{router_ip}/32"})
            
            if ip_response['count'] == 0:
                # Create IP address
                try:
                    ip_data = nb.create_ip_address(
                        address=f"{router_ip}/32",
                        status='active',
                        description=f'{router_name} loopback',
                        role='loopback'
                    )
                    ip_id = ip_data['id']
                    print(f"✓ Created IP address {router_ip}/32")
                except Exception as e:
                    print(f"✗ Failed to create IP address: {e}")
                    ip_id = None
            else:
                ip_id = ip_response['results'][0]['id']
                print(f"✓ IP address {router_ip}/32 already exists")
            
            # Set as primary IP for the device
            if ip_id:
                try:
                    device_update = {
                        'primary_ip4': ip_id
                    }
                    nb.patch(f"dcim/devices/{device_id}/", device_update)
                    print("✓ Set as primary IP for device")
                except Exception as e:
                    print(f"✗ Failed to set primary IP: {e}")
            
        except Exception as e:
            print(f"✗ Failed to create {router_name}: {e}")
            return site_id, None
    else:
        device_id = device_response['results'][0]['id']
        print(f"✓ {router_name} already exists (ID: {device_id})")
    
    print(f"\n=== Summary ===")
    print(f"Site: {site_name} (ID: {site_id})")
    print(f"Device: {router_name} (ID: {device_id})")
    
    return site_id, device_id


def main():
    parser = argparse.ArgumentParser(description='Create a site and router in NetBox')
    parser.add_argument('site_name', help='Name of the site (e.g., Verona, Climax)')
    parser.add_argument('router_ip', help='Loopback IP of the router (e.g., 10.254.254.101)')
    parser.add_argument('--router-name', help='Name for the router (default: {site}-router)')
    parser.add_argument('--manufacturer', default='MikroTik', help='Device manufacturer')
    parser.add_argument('--device-type', default='RouterBOARD', help='Device model/type')
    parser.add_argument('--device-role', default='Router', help='Device role')
    parser.add_argument('--site-comments', help='Comments for the site')
    parser.add_argument('--physical-address', help='Physical address of the site')
    
    args = parser.parse_args()
    
    site_id, device_id = create_site_and_router(
        site_name=args.site_name,
        router_ip=args.router_ip,
        router_name=args.router_name,
        manufacturer=args.manufacturer,
        device_type=args.device_type,
        device_role=args.device_role,
        site_comments=args.site_comments,
        physical_address=args.physical_address
    )
    
    if site_id and device_id:
        print(f"\nSuccess! You can now add network data for this site.")
        return 0
    else:
        print(f"\nPartial success or failure. Check the output above.")
        return 1


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