#!/usr/bin/env python3
"""
CVE Report Splitter
Splits multiple CVE IDs in a single row into separate rows for easier filtering and analysis.
"""

import pandas as pd
import sys
from pathlib import Path

def split_cve_report(input_file, output_file=None, sheet_name='Vulnerabilities', cve_column='CVE ID'):
    """
    Split CVE IDs into separate rows.

    Args:
        input_file: Path to input Excel file
        output_file: Path to output file (default: adds '_Split' to input filename)
        sheet_name: Name of sheet with vulnerability data (default: 'Vulnerabilities')
        cve_column: Name of column containing CVE IDs (default: 'CVE ID')
    """
    input_path = Path(input_file)

    if not input_path.exists():
        print(f"Error: File not found: {input_file}")
        sys.exit(1)

    if output_file is None:
        output_file = input_path.parent / f"{input_path.stem}_Split{input_path.suffix}"

    print(f"Reading: {input_file}")

    try:
        df = pd.read_excel(input_file, sheet_name=sheet_name)
    except ValueError as e:
        print(f"Error: Sheet '{sheet_name}' not found in workbook")
        print(f"Available sheets: {pd.ExcelFile(input_file).sheet_names}")
        sys.exit(1)

    if cve_column not in df.columns:
        print(f"Error: Column '{cve_column}' not found")
        print(f"Available columns: {list(df.columns)}")
        sys.exit(1)

    original_rows = len(df)
    print(f"Original rows: {original_rows}")

    # Split CVE IDs by comma
    df[cve_column] = df[cve_column].astype(str).str.split(',')

    # Explode to create separate rows
    df_exploded = df.explode(cve_column)

    # Clean up CVE IDs
    df_exploded[cve_column] = df_exploded[cve_column].str.strip()
    df_exploded = df_exploded[df_exploded[cve_column].notna()]
    df_exploded = df_exploded[df_exploded[cve_column] != 'nan']
    df_exploded = df_exploded[df_exploded[cve_column] != '']

    # Reset index
    df_exploded = df_exploded.reset_index(drop=True)

    new_rows = len(df_exploded)
    print(f"New rows: {new_rows}")
    print(f"Added {new_rows - original_rows} rows from splitting CVEs")

    # Save output
    df_exploded.to_excel(output_file, index=False, sheet_name=sheet_name)
    print(f"\n✓ Success! Saved to: {output_file}")

    return output_file

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 split_cve_report.py <input_file.xlsx> [output_file.xlsx]")
        print("\nExample:")
        print("  python3 split_cve_report.py 'Vulnerability Workbook.xlsx'")
        print("  python3 split_cve_report.py 'input.xlsx' 'output.xlsx'")
        sys.exit(1)

    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None

    split_cve_report(input_file, output_file)