#!/usr/bin/env python3
import os
import re
import sys
import requests
from bs4 import BeautifulSoup

def generate_nvidia_ids(url, output_dir="."):
    print(f"Fetching NVIDIA chip manifests from: {url}")
    headers = {'User-Agent': 'Mozilla/5.0 (NixOS NVIDIA-detect Tool)'}

    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
    except Exception as e:
        print(f"Execution failed - unable to fetch manifest: {e}", file=sys.stderr)
        sys.exit(1)

    soup = BeautifulSoup(response.text, 'html.parser')
    driver_maps = {"current": []}
    current_driver = "current"
    
    body = soup.body if soup.body else soup
    
    # Linear DOM traversal utilizing structural and literal textual boundaries
    for element in body.descendants:
        # 1. Identify text blocks capable of changing the tracking context
        if element.name in ['h1', 'h2', 'h3', 'h4', 'h5', 'p', 'div']:
            text = element.get_text(" ", strip=True)
            
            # Hard reset to stable driver branch when hitting the primary section header
            if "Current NVIDIA GPUs" in text:
                current_driver = "current"
                continue
            
            # Switch to legacy branch tracking when encountering driver spec strings
            match = re.search(r'The\s+(\d+)(?:\.\d+)?(?:\.xx)?\s+driver\s+supports', text, re.IGNORECASE)
            if match:
                ver = match.group(1)
                current_driver = f"legacy_{ver}"
                if current_driver == "legacy_470":
                    break
                if current_driver not in driver_maps:
                    driver_maps[current_driver] = []
                continue
                    
        # 2. Process data rows inside standard HTML tables
        elif element.name == 'tr':
            cells = [td.get_text(strip=True) for td in element.find_all(['td', 'th'])]
            if len(cells) < 2:
                continue
            
            # Evaluate columns looking for the primary 4-digit hex PCI ID
            for cell in cells[1:]:
                cleaned = cell.strip()
                if cleaned.lower().startswith('0x'):
                    cleaned = cleaned[2:]
                
                words = cleaned.split()
                if not words:
                    continue
                
                first_word = words[0]
                if re.match(r'^[0-9a-fA-F]{4}$', first_word):
                    pci_id = first_word.lower()
                    if pci_id not in driver_maps[current_driver]:
                        driver_maps[current_driver].append(pci_id)
                    break 
                    
        # 3. Fallback processing for older/minimalist pre-formatted text layouts
        elif element.name == 'pre':
            for line in element.get_text().split('\n'):
                line = line.strip()
                if not line:
                    continue
                
                if "Current NVIDIA GPUs" in line:
                    current_driver = "current"
                    continue
                
                match = re.search(r'The\s+(\d+)(?:\.\d+)?(?:\.xx)?\s+driver\s+supports', line, re.IGNORECASE)
                if match:
                    ver = match.group(1)
                    current_driver = f"legacy_{ver}"
                    if current_driver not in driver_maps:
                        driver_maps[current_driver] = []
                    continue
                
                parts = re.split(r'[,.]', line)
                if len(parts) >= 2:
                    val = parts[1].strip()
                    if val.lower().startswith('0x'):
                        val = val[2:]
                    words = val.split()
                    if words:
                        first_word = words[0]
                        if re.match(r'^[0-9a-fA-F]{4}$', first_word):
                            pci_id = first_word.lower()
                            if pci_id not in driver_maps[current_driver]:
                                driver_maps[current_driver].append(pci_id)

    # Clean file generation
    os.makedirs(output_dir, exist_ok=True)
    for driver, ids in driver_maps.items():
        if not ids:
            continue
        filename = os.path.join(output_dir, f"{driver}.ids")
        with open(filename, "w") as f:
            for pci_id in sorted(ids):
                f.write(f"{pci_id}\n")
        print(f"Generated {filename} containing {len(ids)} mapped PCI IDs.")

if __name__ == "__main__":
    version = "595.84"
    if len(sys.argv) > 1:
        version = sys.argv[1]
        target_url = f"https://download.nvidia.com/XFree86/Linux-x86_64/{version}/README/supportedchips.html"
    else:
        print("Please provide the driver version under following format : xxx.xx, for example : 595.84")
    generate_nvidia_ids(target_url)