#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
"""
Diff two GSP firmware versions.
Usage: python3 diff_gsp_firmware.py <old_firmware> <new_firmware> <output_report>
"""

import sys
import os
import subprocess
import json
import hashlib

def extract_fwimage(firmware_path, output_path):
    """Extract .fwimage section from GSP firmware ELF."""
    cmd = ["readelf", "-x", ".fwimage", firmware_path]
    result = subprocess.run(cmd, capture_output=True, text=True)
    # This is a hex dump, we need to parse it
    # Better to use objcopy
    cmd = ["objcopy", "-O", "binary", "-j", ".fwimage", firmware_path, output_path]
    result = subprocess.run(cmd, capture_output=True)
    return result.returncode == 0

def extract_version(firmware_path):
    """Extract version from .fwversion section."""
    cmd = ["readelf", "-x", ".fwversion", firmware_path]
    result = subprocess.run(cmd, capture_output=True, text=True)
    # Parse hex dump
    for line in result.stdout.split('\n'):
        if '0x' in line and len(line.strip()) > 10:
            parts = line.split()
            if len(parts) >= 5:
                hex_str = ''.join(parts[1:5])
                try:
                    version = bytes.fromhex(hex_str).decode('ascii').strip('\x00')
                    return version
                except:
                    pass
    return "unknown"

def run_ghidra_analysis(fwimage_path, project_dir, project_name):
    """Run Ghidra analysis on fwimage."""
    analyze_headless = "/snap/ghidra/47/ghidra/support/analyzeHeadless"
    if not os.path.exists(analyze_headless):
        analyze_headless = "/opt/ghidra/support/analyzeHeadless"
    
    script_path = "/home/rmcore/Git/integdev_gpu_drv/open-iov/tools/ghidra/gsp-tu104-570.144/scripts"
    
    cmd = [
        analyze_headless,
        project_dir,
        project_name,
        "-import", fwimage_path,
        "-processor", "RISCV:LE:64:default",
        "-cspec", "gcc",
        "-overwrite",
        "-scriptPath", script_path,
        "-postScript", "analyze_gsp_fwimage.py"
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=7200)
    return result.returncode == 0

def load_analysis_results(output_dir):
    """Load analysis results from JSON files."""
    results = {}
    for fname in ["functions.json", "rpc_functions.json", "vgpu_strings.json", 
                  "wpr_functions.json", "mem_functions.json", "irq_functions.json",
                  "analysis_summary.json"]:
        path = os.path.join(output_dir, fname)
        if os.path.exists(path):
            with open(path, 'r') as f:
                results[fname.replace('.json', '')] = json.load(f)
    return results

def diff_functions(old_funcs, new_funcs):
    """Compare function lists."""
    old_names = {f['name']: f for f in old_funcs}
    new_names = {f['name']: f for f in new_funcs}
    
    added = [f for name, f in new_names.items() if name not in old_names]
    removed = [f for name, f in old_names.items() if name not in new_names]
    modified = []
    
    for name in set(old_names.keys()) & set(new_names.keys()):
        if old_names[name]['size'] != new_names[name]['size'] or \
           old_names[name]['entry'] != new_names[name]['entry']:
            modified.append({
                'name': name,
                'old': old_names[name],
                'new': new_names[name]
            })
    
    return added, removed, modified

def generate_diff_report(old_firmware, new_firmware, old_results, new_results, output_path):
    """Generate markdown diff report."""
    old_version = extract_version(old_firmware)
    new_version = extract_version(new_firmware)
    
    old_funcs = old_results.get('functions', [])
    new_funcs = new_results.get('functions', [])
    
    added, removed, modified = diff_functions(old_funcs, new_funcs)
    
    old_rpc = {f['name'] for f in old_results.get('rpc_functions', [])}
    new_rpc = {f['name'] for f in new_results.get('rpc_functions', [])}
    
    rpc_added = new_rpc - old_rpc
    rpc_removed = old_rpc - new_rpc
    
    report = f"""# GSP Firmware Diff Report

## Versions
- **Old**: {old_version} ({os.path.basename(old_firmware)})
- **New**: {new_version} ({os.path.basename(new_firmware)})

## Function Changes
- **Total Functions**: {len(old_funcs)} → {len(new_funcs)} ({len(new_funcs) - len(old_funcs):+d})
- **Added**: {len(added)}
- **Removed**: {len(removed)}
- **Modified**: {len(modified)}

## RPC Handler Changes
- **Added RPC Functions**: {len(rpc_added)}
- **Removed RPC Functions**: {len(rpc_removed)}

## Added Functions
"""
    for f in added[:50]:
        report += f"- `{f['name']}` at {f['entry']} (size: {f['size']})\n"
    
    if len(added) > 50:
        report += f"\n... and {len(added) - 50} more\n"
    
    report += "\n## Removed Functions\n"
    for f in removed[:50]:
        report += f"- `{f['name']}` at {f['entry']} (size: {f['size']})\n"
    
    if len(removed) > 50:
        report += f"\n... and {len(removed) - 50} more\n"
    
    report += "\n## Modified Functions\n"
    for f in modified[:50]:
        report += f"- `{f['name']}`: {f['old']['entry']}→{f['new']['entry']}, size {f['old']['size']}→{f['new']['size']}\n"
    
    if len(modified) > 50:
        report += f"\n... and {len(modified) - 50} more\n"
    
    report += "\n## Added RPC Functions\n"
    for name in sorted(rpc_added):
        report += f"- `{name}`\n"
    
    report += "\n## Removed RPC Functions\n"
    for name in sorted(rpc_removed):
        report += f"- `{name}`\n"
    
    # vGPU strings comparison
    old_strings = {s['value'] for s in old_results.get('vgpu_strings', [])}
    new_strings = {s['value'] for s in new_results.get('vgpu_strings', [])}
    
    report += "\n## vGPU String Changes\n"
    report += f"- **Added**: {len(new_strings - old_strings)}\n"
    report += f"- **Removed**: {len(old_strings - new_strings)}\n"
    
    for s in sorted(new_strings - old_strings)[:20]:
        report += f"  + `{s}`\n"
    for s in sorted(old_strings - new_strings)[:20]:
        report += f"  - `{s}`\n"
    
    with open(output_path, 'w') as f:
        f.write(report)
    
    print(f"Diff report saved to {output_path}")

def main():
    if len(sys.argv) < 4:
        print("Usage: diff_gsp_firmware.py <old_firmware> <new_firmware> <output_report>")
        sys.exit(1)
    
    old_firmware = sys.argv[1]
    new_firmware = sys.argv[2]
    output_report = sys.argv[3]
    
    # Extract .fwimage sections
    old_fwimage = "/tmp/old_fwimage.bin"
    new_fwimage = "/tmp/new_fwimage.bin"
    
    print("Extracting .fwimage sections...")
    if not extract_fwimage(old_firmware, old_fwimage):
        print("Failed to extract old .fwimage")
        sys.exit(1)
    if not extract_fwimage(new_firmware, new_fwimage):
        print("Failed to extract new .fwimage")
        sys.exit(1)
    
    # Run Ghidra analysis on both
    print("Analyzing old firmware...")
    old_project_dir = "/tmp/ghidra_old"
    if not run_ghidra_analysis(old_fwimage, old_project_dir, "gsp_old"):
        print("Failed to analyze old firmware")
        sys.exit(1)
    
    print("Analyzing new firmware...")
    new_project_dir = "/tmp/ghidra_new"
    if not run_ghidra_analysis(new_fwimage, new_project_dir, "gsp_new"):
        print("Failed to analyze new firmware")
        sys.exit(1)
    
    # Load results
    old_results = load_analysis_results("/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP")
    # For new firmware, we'd need a separate output dir
    new_results = load_analysis_results("/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/NEW-GSP")
    
    # Generate diff report
    generate_diff_report(old_firmware, new_firmware, old_results, new_results, output_report)
    
    print("Done!")

if __name__ == "__main__":
    main()