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

# The target URLs. You can freely add any new Intel GPU driver documentation links here.
URLS = [
    "https://dgpu-docs.intel.com/overview/supported-hardware/xe-driver-gpus.html"
]

def get_filename_from_url(url):
    """
    Derives a clean .ids filename from the URL slug.
    Examples: 
      .../legacy-gpus.html        -> legacy.ids
      .../i915-driver-gpus.html   -> i915.ids
      .../xe2-driver-gpus.html    -> xe2.ids (for future updates)
    """
    slug = url.split('/')[-1].replace('.html', '')
    for suffix in ['-driver-gpus', '-gpus', '-driver']:
        if slug.endswith(suffix):
            return slug[:-len(suffix)] + '.ids'
    return slug + '.ids'

def extract_pci_ids(url):
    """
    Scrapes the URL, dynamically locates the PCI ID column, and extracts
    unique, clean 4-digit hexadecimal PCI IDs.
    """
    print(f"Scraping: {url}")
    headers = {
        'User-Agent': 'Mozilla/5.0 (NixOS Intel-detect Tool)'
    }
    
    try:
        response = requests.get(url, headers=headers, timeout=15)
        response.raise_for_status()
    except requests.RequestException as e:
        print(f"Error fetching {url}: {e}", file=sys.stderr)
        return []

    soup = BeautifulSoup(response.text, 'html.parser')
    pci_ids = set()
    
    tables = soup.find_all('table')
    if not tables:
        print(f"Warning: No data tables found on {url}", file=sys.stderr)
        return []

    for table in tables:
        # 1. Inspect headers to locate the exact index of the PCI/Device ID column
        thead = table.find('thead')
        if thead:
            headers_text = [th.get_text(strip=True).lower() for th in thead.find_all('th')]
        else:
            first_tr = table.find('tr')
            headers_text = [th.get_text(strip=True).lower() for th in first_tr.find_all(['th', 'td'])] if first_tr else []

        pci_col_idx = -1
        for idx, text in enumerate(headers_text):
            if any(kwd in text for kwd in ['pci', 'device id', 'dev id', 'did']):
                pci_col_idx = idx
                break
        
        # Soft fallback: check if header just contains 'id'
        if pci_col_idx == -1:
            for idx, text in enumerate(headers_text):
                if 'id' in text:
                    pci_col_idx = idx
                    break

        # 2. Iterate through rows and pull data
        tbody = table.find('tbody')
        rows = tbody.find_all('tr') if tbody else table.find_all('tr')
        
        for row in rows:
            cells = row.find_all('td')
            if not cells:
                continue
            
            # Extract strictly from the matching column if found
            if pci_col_idx != -1 and pci_col_idx < len(cells):
                cell_text = cells[pci_col_idx].get_text(strip=True)
                # Matches standalone 4-char hex strings, ignoring optional '0x' prefix
                matches = re.findall(r'\b(?:0x)?([0-9a-fA-F]{4})\b', cell_text)
                for match in matches:
                    pci_ids.add(match.lower())
            else:
                # Absolute fallback: Scan all cells but guard heavily against false-positive string matches
                for cell in cells:
                    cell_text = cell.get_text(strip=True)
                    matches = re.findall(r'\b(?:0x)?([0-9a-fA-F]{4})\b', cell_text)
                    for match in matches:
                        # Only accept if the cell text is short (an ID cell) or explicitly contains 0x
                        if len(cell_text) < 15 or '0x' in cell_text.lower():
                            pci_ids.add(match.lower())

    return sorted(list(pci_ids))

def main():
    for url in URLS:
        filename = get_filename_from_url(url)
        ids = extract_pci_ids(url)
        
        if ids:
            with open(filename, 'w', encoding='utf-8') as f:
                for pci_id in ids:
                    f.write(f"{pci_id}\n")
            print(f"-> Successfully saved {len(ids)} IDs to {filename}\n")
        else:
            print(f"-> No IDs extracted for {url}. Skipping file write.\n")

if __name__ == "__main__":
    main()