#!/usr/bin/env python3
"""Stage built native .so files for EcoArk pythonrunner main-so integration.

Extracts .so files from an existing HarmonyOS-cross-compiled wheel and copies
them into an integration output directory with a README describing where each
module should be placed in the EcoArk pythonrunner environment.
"""
import argparse
import json
import os
import shutil
import sys
import zipfile


def stage_integration(wheel_path, output_dir, recipe=None):
    os.makedirs(output_dir, exist_ok=True)

    if not os.path.isfile(wheel_path):
        print(f"[stage] ERROR: wheel not found: {wheel_path}", file=sys.stderr)
        return 1

    staged = []
    with zipfile.ZipFile(wheel_path, "r") as z:
        for name in z.namelist():
            if not name.endswith(".so"):
                continue
            dest = os.path.join(output_dir, name)
            os.makedirs(os.path.dirname(dest), exist_ok=True)
            with z.open(name) as src, open(dest, "wb") as dst:
                shutil.copyfileobj(src, dst)
            staged.append(name)
            print(f"[stage]  {name}")

    readme_path = os.path.join(output_dir, "README.md")
    with open(readme_path, "w") as f:
        f.write(f"# Integration Staging for EcoArk pythonrunner main-so linking\n\n")
        f.write(f"Source wheel: `{wheel_path}`\n\n")
        f.write(f"## Staged .so files ({len(staged)} total)\n\n")
        for s in staged:
            f.write(f"- `{s}`\n")
        f.write(f"\n## EcoArk target hint\n\n")
        f.write("Copy each .so into EcoArk pythonrunner's native module tree\n")
        f.write("under `lib/python3.12/site-packages/` preserving subdirectory\n")
        f.write("structure so that `import numpy` resolves via main-so linking.\n")
        so_list_path = os.path.join(output_dir, "integration_manifest.json")
        manifest = {
            "pkg": "numpy",
            "staged_count": len(staged),
            "ecoark_target": "lib/python3.12/site-packages/",
            "files": staged,
        }
        if recipe and os.path.isfile(recipe):
            with open(recipe) as rf:
                manifest["recipe"] = json.load(rf)
    with open(so_list_path, "w") as f:
        json.dump(manifest, f, indent=2)

    print(f"\n[stage] Staged {len(staged)} .so files to {output_dir}")
    print(f"[stage] Manifest: {so_list_path}")
    print(f"[stage] README:   {readme_path}")
    return 0


def main():
    parser = argparse.ArgumentParser(
        description="Stage native .so files from a built HarmonyOS wheel for EcoArk integration"
    )
    parser.add_argument("wheel", nargs="?", help="Path to built HarmonyOS .whl file")
    parser.add_argument("--output-dir", "-o", default="output/integration",
                        help="Integration output directory (default: output/integration)")
    parser.add_argument("--recipe", help="Path to integration recipe JSON (optional)")
    parser.add_argument("--list-recipes", action="store_true",
                        help="List available recipes and exit")
    args = parser.parse_args()

    if args.list_recipes:
        recipes_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "recipes")
        print("Available integration recipes:")
        if os.path.isdir(recipes_dir):
            for fn in sorted(os.listdir(recipes_dir)):
                if fn.endswith(".json"):
                    rp = os.path.join(recipes_dir, fn)
                    with open(rp) as f:
                        d = json.load(f)
                    print(f"  {fn}  ({d.get('pkg','?')} {d.get('version','?')})")
        else:
            print("  (no recipes directory)")
        return 0

    if not args.wheel:
        # Default: look for numpy wheel in output/wheels
        default = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            "output", "wheels", "numpy-1.26.4-cp312-cp312-harmonyos_arm64.whl"
        )
        if os.path.isfile(default):
            args.wheel = default
            print(f"[stage] Using default wheel: {default}")
        else:
            parser.error("no wheel specified and default numpy wheel not found")

    return stage_integration(args.wheel, os.path.abspath(args.output_dir), recipe=args.recipe)


if __name__ == "__main__":
    sys.exit(main())