#!/usr/bin/env python3
from __future__ import annotations

import os
import shutil
import subprocess
import sys
from pathlib import Path


DEFAULT_DEPTH = "2"
DEFAULT_ROOTS = [
    ("entry", "0x005a313b"),
    ("bootstrap", "0x00484440"),
]


def parse_args(argv: list[str]) -> tuple[Path, Path, str, list[tuple[str, str]]]:
    if len(argv) < 3:
        print(
            "usage: export_startup_map.py <exe-path> <output-dir> [--depth N] [--root NAME:ADDR ...]",
            file=sys.stderr,
        )
        raise SystemExit(2)

    exe_path = Path(argv[1]).resolve()
    output_dir = Path(argv[2]).resolve()
    depth = DEFAULT_DEPTH
    roots = list(DEFAULT_ROOTS)

    index = 3
    while index < len(argv):
        token = argv[index]
        if token == "--depth":
            if index + 1 >= len(argv):
                print("--depth requires a value", file=sys.stderr)
                raise SystemExit(2)
            depth = argv[index + 1]
            index += 2
            continue

        if token == "--root":
            if index + 1 >= len(argv):
                print("--root requires NAME:ADDR", file=sys.stderr)
                raise SystemExit(2)
            root_spec = argv[index + 1]
            if ":" not in root_spec:
                print("--root value must be NAME:ADDR", file=sys.stderr)
                raise SystemExit(2)
            name, address = root_spec.split(":", 1)
            if not name or not address:
                print("--root value must be NAME:ADDR", file=sys.stderr)
                raise SystemExit(2)
            if roots == DEFAULT_ROOTS:
                roots = []
            roots.append((name, address))
            index += 2
            continue

        print("unknown argument: %s" % token, file=sys.stderr)
        raise SystemExit(2)

    return exe_path, output_dir, depth, roots


def main() -> int:
    repo_root = Path(__file__).resolve().parents[2]
    exe_path, output_dir, depth, roots = parse_args(sys.argv)
    output_dir.mkdir(parents=True, exist_ok=True)

    ghidra_home = Path.home() / "software" / "ghidra"
    analyze_headless = ghidra_home / "support" / "analyzeHeadless"
    if not analyze_headless.exists():
        print("analyzeHeadless not found at %s" % analyze_headless, file=sys.stderr)
        return 1

    project_dir = repo_root / "ghidra_projects"
    project_dir.mkdir(parents=True, exist_ok=True)
    project_name = "rt3-1.06"
    script_path = repo_root / "tools" / "ghidra" / "scripts"
    ghidra_runtime = repo_root / ".ghidra"
    ghidra_home_dir = ghidra_runtime / "home"
    ghidra_config_dir = ghidra_runtime / "config"
    ghidra_cache_dir = ghidra_runtime / "cache"
    ghidra_java_prefs = ghidra_runtime / "java-prefs"
    ghidra_home_dir.mkdir(parents=True, exist_ok=True)
    ghidra_config_dir.mkdir(parents=True, exist_ok=True)
    ghidra_cache_dir.mkdir(parents=True, exist_ok=True)
    ghidra_java_prefs.mkdir(parents=True, exist_ok=True)

    java_bin = shutil.which("java")
    if java_bin is None:
        print("java not found on PATH", file=sys.stderr)
        return 1
    java_home = Path(java_bin).resolve().parents[1]

    command = [
        str(analyze_headless),
        str(project_dir),
        project_name,
        "-import",
        str(exe_path),
        "-overwrite",
        "-scriptPath",
        str(script_path),
        "-postScript",
        "ExportStartupFunctions.java",
        str(output_dir),
        depth,
    ]
    command.extend(["%s:%s" % (name, address) for name, address in roots])
    command.extend([
        "-analysisTimeoutPerFile",
        "600",
    ])

    env = os.environ.copy()
    env["HOME"] = str(ghidra_home_dir)
    env["XDG_CONFIG_HOME"] = str(ghidra_config_dir)
    env["XDG_CACHE_HOME"] = str(ghidra_cache_dir)
    env["JAVA_HOME"] = str(java_home)
    env["_JAVA_OPTIONS"] = "-Djava.util.prefs.userRoot=%s" % ghidra_java_prefs

    subprocess.run(command, check=True, env=env)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())