#!/usr/bin/env python3
"""Command-line interface for IDA Pro batch decompiler."""

import argparse
import os
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path


def cleanup_ida_files(directory: Path) -> int:
    """Clean up IDA Pro intermediate files in a directory."""
    ida_extensions = {'.id0', '.id1', '.nam', '.til', '.i64', '.asm', '.cor', '.trs', '.mid'}

    deleted_count = 0
    for root, dirs, files in os.walk(directory):
        for filename in files:
            file_path = Path(root) / filename
            ext = file_path.suffix.lower()

            if ext in ida_extensions or filename.endswith('~'):
                try:
                    file_path.unlink()
                    deleted_count += 1
                except Exception:
                    pass

    return deleted_count


def is_binary_file(path: Path) -> bool:
    """Check if a file is a binary executable or shared library."""
    try:
        result = subprocess.run(
            ['file', str(path)],
            capture_output=True,
            text=True,
            timeout=5
        )
        output = result.stdout.lower()
        return any(kind in output for kind in ['executable', 'shared object', ' Mach-O'])
    except Exception:
        return False


def find_binaries(directory: Path) -> list:
    """Find all binary files in a directory."""
    binaries = []
    for root, dirs, files in os.walk(directory):
        for filename in files:
            file_path = Path(root) / filename
            if file_path.is_file() and is_binary_file(file_path):
                binaries.append(file_path)
    return binaries


def create_idapython_script(output_dir: str) -> str:
    """Create IDAPython script for exporting C code."""
    output_dir_escaped = output_dir.replace("\\", "\\\\")
    lines = []
    lines.append("import idc")
    lines.append("import idaapi")
    lines.append("import idautils")
    lines.append("import os")
    lines.append("import datetime")
    lines.append("")
    lines.append("output_dir = os.environ.get('IDA_OUTPUT_DIR', r'%s')" % output_dir_escaped)
    lines.append("bin_name = os.path.basename(os.path.splitext(idc.get_input_file_path())[0])")
    lines.append("output_file = os.path.join(output_dir, bin_name + \"_decompiled.c\")")
    lines.append("")
    lines.append("os.makedirs(output_dir, exist_ok=True)")
    lines.append("")
    lines.append("def decompile_all_functions_to_single_file(output_path):")
    lines.append("    start_time = datetime.datetime.now()")
    lines.append("")
    lines.append("    try:")
    lines.append("        input_file = os.path.basename(idaapi.idb_path)")
    lines.append("    except:")
    lines.append("        try:")
    lines.append("            input_file = os.path.basename(idc.get_input_file_path())")
    lines.append("        except:")
    lines.append("            input_file = \"unknown\"")
    lines.append("")
    lines.append("    functions = list(idautils.Functions())")
    lines.append("    total = len(functions)")
    lines.append("")
    lines.append("    with open(output_path, 'w', encoding='utf-8') as f:")
    lines.append("        f.write(\"/*\\n\")")
    lines.append("        f.write(\" * File: \" + os.path.basename(output_path) + \"\\n\")")
    lines.append("        f.write(\" * Generated from: \" + input_file + \"\\n\")")
    lines.append("        f.write(\" * Generated at: \" + datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + \"\\n\")")
    lines.append("        f.write(\" * Total functions: \" + str(total) + \"\\n\")")
    lines.append("        f.write(\" */\\n\\n\")")
    lines.append("        f.write(\"#include <stdint.h>\\n\")")
    lines.append("        f.write(\"#include <stddef.h>\\n\\n\")")
    lines.append("")
    lines.append("        success_count = 0")
    lines.append("        for i, func_addr in enumerate(functions):")
    lines.append("            func_name = idc.get_func_name(func_addr)")
    lines.append("            if not func_name:")
    lines.append("                func_name = \"sub_%x\" % func_addr")
    lines.append("")
    lines.append("            try:")
    lines.append("                decomp = idaapi.decompile(func_addr)")
    lines.append("                if decomp is not None:")
    lines.append("                    c_code = str(decomp)")
    lines.append("                    f.write(\"/*\\n * Function: \" + func_name + \"\\n\")")
    lines.append("                    f.write(\" * Address: 0x%x\" % func_addr + \"\\n\")")
    lines.append("                    f.write(\" * Order: %d/%d\" % (i+1, total) + \"\\n\")")
    lines.append("                    f.write(\" */\\n\")")
    lines.append("                    f.write(c_code)")
    lines.append("                    f.write(\"\\n\\n\" + \"=\" * 80 + \"\\n\\n\")")
    lines.append("                    success_count += 1")
    lines.append("                    print(\"[%d/%d] %s\" % (i+1, total, func_name))")
    lines.append("            except Exception as e:")
    lines.append("                print(\"Failed to decompile {}: {}\".format(func_name, e))")
    lines.append("")
    lines.append("    print(\"\\nExport complete: \" + output_file)")
    lines.append("    print(\"Success: %d/%d functions\" % (success_count, total))")
    lines.append("    return success_count > 0")
    lines.append("")
    lines.append("success = decompile_all_functions_to_single_file(output_file)")
    lines.append("")
    lines.append("if success:")
    lines.append("    print(\"SUCCESS: \" + output_file)")
    lines.append("else:")
    lines.append("    print(\"FAILED: No functions were decompiled\")")
    lines.append("")
    lines.append("idaapi.qexit(0 if success else 1)")

    return "\n".join(lines)


def decompile_binary(idat_path: str, binary_path: str, output_dir: str,
                     timeout: int, verbose: bool) -> tuple:
    """Decompile a single binary file."""
    binary_path = Path(binary_path)
    if not binary_path.exists():
        return False, f"File not found: {binary_path}"

    bin_name = binary_path.stem
    output_file = Path(output_dir) / f"{bin_name}_decompiled.c"

    script_content = create_idapython_script(output_dir)

    with tempfile.NamedTemporaryFile(
        mode='w',
        suffix='.py',
        prefix=f'ida_export_{bin_name}_',
        delete=False
    ) as f:
        script_path = f.name
        f.write(script_content)

    try:
        cmd = [
            str(idat_path),
            "-A",
            f"-S{script_path}",
            str(binary_path)
        ]

        if verbose:
            print(f"Processing: {bin_name}")
            print(f"  Command: {' '.join(cmd)}")

        result = subprocess.run(
            cmd,
            timeout=timeout,
            capture_output=True,
            text=True,
            env={**os.environ, 'IDA_OUTPUT_DIR': output_dir}
        )

        if output_file.exists():
            return True, f"Generated {output_file}"
        else:
            stderr = result.stderr.strip() if result.stderr else result.stdout.strip()
            # Add more debug info
            debug_msg = f"Output not created. Expected: {output_file}"
            if stderr:
                debug_msg += f"\nIDAT output: {stderr}"
            return False, debug_msg

    except subprocess.TimeoutExpired:
        return False, f"Timeout after {timeout}s"
    except Exception as e:
        return False, f"Exception: {e}"
    finally:
        try:
            os.unlink(script_path)
        except:
            pass


def main():
    parser = argparse.ArgumentParser(
        description="IDA Pro Batch Decompiler - Batch decompile binaries to C code",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  ida-decompile -i ./binaries -o ./output -p /Applications/IDA\\ Professional\\ 9.1.app/Contents/MacOS/idat
  ida-decompile -i /path/to/binary -o ./decompiled -p /path/to/idat
  ida-decompile -i ./binaries -o ./output -p /path/to/idat -j 8 -t 600
        """
    )

    parser.add_argument(
        '--version',
        action='store_true',
        help='Show version'
    )
    parser.add_argument(
        '-i', '--input-dir',
        required=True,
        help='Input directory or binary file path'
    )
    parser.add_argument(
        '-o', '--output-dir',
        required=True,
        help='Output directory for C files'
    )
    parser.add_argument(
        '-p', '--idat-path',
        required=True,
        help='Path to idat executable'
    )
    parser.add_argument(
        '-j', '--jobs',
        type=int,
        default=4,
        help='Number of parallel processes (default: 4)'
    )
    parser.add_argument(
        '-t', '--timeout',
        type=int,
        default=300,
        help='Timeout per file in seconds (default: 300)'
    )
    parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        help='Show verbose output'
    )

    args = parser.parse_args()

    if args.version:
        try:
            from . import __version__
            print(f"ida-decompile version {__version__}")
        except:
            print("ida-decompile version 0.2.1")
        sys.exit(0)

    input_path = Path(args.input_dir)
    if not input_path.exists():
        print(f"ERROR: Input path not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'=' * 60}")
    print("IDA Pro Batch Decompiler")
    print(f"{'=' * 60}")
    print(f"\nInput: {input_path}")
    print(f"Output: {output_dir}")
    print(f"IDAT: {args.idat_path}")
    print(f"Jobs: {args.jobs}")
    print(f"Timeout: {args.timeout}s")
    print()

    if not input_path.is_file():
        print(f"Cleaning up IDA files in {input_path}...")
        cleaned_count = cleanup_ida_files(input_path)
        if cleaned_count > 0:
            print(f"Cleaned up {cleaned_count} IDA files")
        print()

    if input_path.is_file():
        print(f"Processing single file: {input_path.name}")
        success, message = decompile_binary(
            args.idat_path, str(input_path), str(output_dir),
            args.timeout, args.verbose
        )
        status = "OK" if success else "FAIL"
        print(f"[{status}] {input_path.name}: {message}")
        success_count = 1 if success else 0
        fail_count = 0 if success else 1
    else:
        binaries = find_binaries(input_path)

        if not binaries:
            print("\nNo binary files found.")
            return

        print(f"Found {len(binaries)} binary file(s)\n")

        print(f"{'=' * 60}")
        print("Starting decomposition...")
        print(f"{'=' * 60}\n")

        success_count = 0
        fail_count = 0

        if args.jobs > 1 and len(binaries) > 1:
            import multiprocessing

            def decompile_wrapper(args):
                return decompile_binary(*args)

            decompile_args = [
                (args.idat_path, str(b), str(output_dir), args.timeout, args.verbose)
                for b in binaries
            ]

            with multiprocessing.Pool(args.jobs) as pool:
                results = pool.map(decompile_wrapper, decompile_args)
                for binary_path, (success, msg) in zip(binaries, results):
                    status = "OK" if success else "FAIL"
                    print(f"[{status}] {binary_path.name}: {msg}")
                    if success:
                        success_count += 1
                    else:
                        fail_count += 1
        else:
            for binary in binaries:
                success, msg = decompile_binary(
                    args.idat_path, str(binary), str(output_dir),
                    args.timeout, args.verbose
                )
                status = "OK" if success else "FAIL"
                print(f"[{status}] {binary.name}: {msg}")
                if success:
                    success_count += 1
                else:
                    fail_count += 1

    print(f"\n{'=' * 60}")
    print("Processing complete!")
    print(f"{'=' * 60}")
    print(f"Success: {success_count}")
    print(f"Failed: {fail_count}")
    print(f"Total: {success_count + fail_count}")
    print(f"\nOutput directory: {output_dir}")


if __name__ == "__main__":
    main()
