"""
Core decompilation logic for IDA Pro batch decompiler.
This module provides batch decompiling binary files using IDA Pro's idat command-line tool.
"""

import os
import subprocess
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple


class IDABatchDecompiler:
    """
    IDA Pro batch decompiler class.

    This class provides methods for batch decompiling binary files using
    IDA Pro's idat command-line tool.

    Args:
        idat_path: Path to idat executable
        output_dir: Output directory for generated C files
        jobs: Number of parallel processes
        timeout: Timeout in seconds per file
        verbose: Enable verbose output
    """

    def __init__(
        self,
        idat_path: str = "/Applications/IDA Professional 9.1.app/Contents/MacOS/idat",
        output_dir: str = "./output",
        jobs: int = 4,
        timeout: int = 300,
        verbose: bool = False
    ):
        self.idat_path = Path(idat_path)
        self.output_dir = Path(output_dir)
        self.jobs = jobs
        self.timeout = timeout
        self.verbose = verbose

        # Validate idat path
        if not self.idat_path.exists():
            raise FileNotFoundError(f"idat not found: {self.idat_path}")

        # Create output directory
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def cleanup_ida_files(self, directory: Path) -> int:
        """
        Clean up IDA Pro intermediate files and database files in a directory.

        Args:
            directory: The directory to scan for IDA files

        Returns:
            Number of files deleted
        """
        # IDA Pro file patterns to clean up
        # .id0, .id1 - IDA database files (index, segments)
        # .nam - named places database
        # .til - type library files
        # .i64 - flat instruction text file
        # .asm - assembly listing
        # .cor - IDA comment database
        # .trs - trace file
        # .mid - module information
        # *~ - temporary backup files
        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()

                # Check if it's an IDA file
                is_ida_file = ext in ida_extensions or filename.endswith('~')

                if is_ida_file:
                    try:
                        file_path.unlink()
                        deleted_count += 1
                        if self.verbose:
                            print(f"  Cleaned up: {file_path}")
                    except Exception as e:
                        if self.verbose:
                            print(f"  Warning: Could not delete {file_path}: {e}")

        return deleted_count

    def create_idapython_script(self) -> str:
        """
        Create IDAPython script for exporting C code using idaapi.decompile() API.

        This script is based on export_decompile.py and uses the decompile() API
        to get high-quality C code output.

        Returns:
            The script content as string
        """
        # The IDAPython script content - using raw string to handle backslashes
        script_content = r"""import idc
import idaapi
import idautils
import os
import datetime

# Get output directory from environment
output_dir = os.environ.get('IDA_OUTPUT_DIR', '{{output_dir}}')
bin_name = os.path.splitext(idc.get_input_file_path())[0]
output_file = os.path.join(output_dir, bin_name + "_decompiled.c")

os.makedirs(output_dir, exist_ok=True)

def decompile_all_functions_to_single_file(output_path):
    start_time = datetime.datetime.now()

    try:
        input_file = os.path.basename(idaapi.idb_path)
    except:
        try:
            input_file = os.path.basename(idc.get_input_file_path())
        except:
            input_file = "unknown"

    functions = list(idautils.Functions())
    total = len(functions)

    with open(output_path, 'w', encoding='utf-8') as f:
        # Write file header
        f.write("/*\n")
        f.write(" * File: " + os.path.basename(output_path) + "\n")
        f.write(" * Generated from: " + input_file + "\n")
        f.write(" * Generated at: " + datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + "\n")
        f.write(" * Total functions: " + str(total) + "\n")
        f.write(" */\n\n")
        f.write("#include <stdint.h>\n")
        f.write("#include <stddef.h>\n\n")

        success_count = 0
        for i, func_addr in enumerate(functions):
            func_name = idc.get_func_name(func_addr)
            if not func_name:
                func_name = "sub_{:x}".format(func_addr)

            try:
                decomp = idaapi.decompile(func_addr)
                if decomp is not None:
                    c_code = str(decomp)
                    # Write function header comment
                    f.write("/*\n * Function: " + func_name + "\n")
                    f.write(" * Address: 0x{:x}\n".format(func_addr))
                    f.write(" * Order: {}/{}\n".format(i+1, total))
                    f.write(" */\n")
                    f.write(c_code)
                    f.write("\n\n" + "=" * 80 + "\n\n")
                    success_count += 1
                    print("[{}/{}] {}".format(i+1, total, func_name))
            except Exception as e:
                print("Failed to decompile {}: {}".format(func_name, e))

    print("\nExport complete: " + output_file)
    print("Success: {}/{} functions".format(success_count, total))
    return success_count > 0

# Export
success = decompile_all_functions_to_single_file(output_file)

if success:
    print("SUCCESS: " + output_file)
else:
    print("FAILED: No functions were decompiled")

idaapi.qexit(0 if success else 1)
"""

        # Replace the placeholder with actual output directory
        script_content = script_content.replace("{{output_dir}}", str(self.output_dir))

        return script_content

    def decompile_binary(self, binary_path: str) -> Tuple[bool, str]:
        """
        Decompile a single binary file.

        Args:
            binary_path: Path to the binary file

        Returns:
            Tuple of (success, message)
        """
        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 = self.output_dir / f"{bin_name}_decompiled.c"

        # Create temporary script file
        script_content = self.create_idapython_script()
        print(script_content)

        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:
            # Build idat command
            cmd = [
                str(self.idat_path),
                "-A",       # Automatic mode
                f"-S{script_path}",
                str(binary_path)
            ]

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

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

            if output_file.exists():
                return True, f"Generated {output_file}"
            else:
                stderr = result.stderr.strip() if result.stderr else result.stdout.strip()
                return False, f"Output not created: {stderr}"

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

    def decompile_directory(self, input_dir: str) -> Dict[str, Tuple[bool, str]]:
        """
        Decompile all binaries in a directory.

        Args:
            input_dir: Path to input directory

        Returns:
            Dictionary mapping binary paths to (success, message) tuples
        """
        from .discovery import BinaryScanner

        input_dir = Path(input_dir)
        scanner = BinaryScanner()

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

        binaries = scanner.scan_directory(input_dir)

        results: Dict[str, Tuple[bool, str]] = {}

        if self.jobs > 1 and len(binaries) > 1:
            # Parallel processing
            import multiprocessing

            args_list = [
                (str(b),) for b in binaries
            ]

            with multiprocessing.Pool(self.jobs) as pool:
                binary_paths = [str(b) for b in binaries]
                decompile_results = pool.map(
                    self._decompile_wrapper,
                    [(b,) for b in binaries]
                )
                results = dict(zip(binary_paths, decompile_results))
        else:
            # Sequential processing
            for binary in binaries:
                success, message = self.decompile_binary(str(binary))
                results[str(binary)] = (success, message)

        return results

    def _decompile_wrapper(self, args: Tuple[str]) -> Tuple[bool, str]:
        """Wrapper for multiprocessing."""
        return self.decompile_binary(args[0])

    def decompile_batch(
        self,
        binary_paths: List[str],
        parallel: bool = True
    ) -> Dict[str, Tuple[bool, str]]:
        """
        Decompile a list of binary files.

        Args:
            binary_paths: List of binary file paths
            parallel: Enable parallel processing

        Returns:
            Dictionary mapping binary paths to (success, message) tuples
        """
        results: Dict[str, Tuple[bool, str]] = {}

        if parallel and self.jobs > 1 and len(binary_paths) > 1:
            import multiprocessing

            with multiprocessing.Pool(self.jobs) as pool:
                decompile_results = pool.map(
                    self._decompile_wrapper,
                    [(b,) for b in binary_paths]
                )
                results = dict(zip(binary_paths, decompile_results))
        else:
            for binary_path in binary_paths:
                success, message = self.decompile_binary(binary_path)
                results[binary_path] = (success, message)

        return results
