"""Real cross-compilation + wheel assembly for HarmonyOS/arm64."""

import base64
import hashlib
import os
import re
import shutil
import subprocess
import sys
import tarfile
import zipfile

# ── Paths ─────────────────────────────────────────────────
HERE = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.dirname(HERE)
CROSS_FILE = os.path.join(REPO_ROOT, "experiments", "ohos-aarch64-cross.ini")
OHOS_NDK = "/home/program/tools/command-line-tools/sdk/default/openharmony/native/llvm"
SOURCES_DIR = "/tmp/ecoark-sources"
OUTPUT_DIR = os.path.join(REPO_ROOT, "output", "wheels")


def _check_toolchain():
    if not os.path.isdir(OHOS_NDK):
        print(f"[rebuild] ERROR: OHOS NDK not found at {OHOS_NDK}")
        print("[rebuild] Set OHOS_NDK_PATH or install the NDK first.")
        return False
    for exe in ["clang", "clang++", "ld.lld"]:
        path = os.path.join(OHOS_NDK, "bin", exe)
        if not os.path.isfile(path):
            print(f"[rebuild] Missing tool: {path}")
            return False
    return True


def _ensure_dirs(path):
    os.makedirs(path, exist_ok=True)
    return path


def _download_sdist(pkg_name, version, dest_dir):
    """Download a PyPI source tarball (sdist) for the given package/version."""
    os.makedirs(dest_dir, exist_ok=True)
    url = f"https://files.pythonhosted.org/packages/source/{pkg_name[0]}/{pkg_name}/{pkg_name}-{version}.tar.gz"
    dest = os.path.join(dest_dir, f"{pkg_name}-{version}.tar.gz")
    if os.path.isfile(dest):
        print(f"[rebuild] Source tarball already cached: {dest}")
        return dest
    print(f"[rebuild] Downloading {url} ...")
    ret = subprocess.run(
        ["curl", "-sL", url, "--connect-timeout", "10", "--max-time", "300", "-o", dest],
        capture_output=True, timeout=310,
    )
    if ret.returncode != 0:
        raise RuntimeError(f"Download failed: {ret.stderr.decode()}")
    print(f"[rebuild] Downloaded: {dest} ({os.path.getsize(dest)} bytes)")
    return dest


def _extract_tarball(tarball_path, dest_dir):
    """Extract a .tar.gz to dest_dir and return the top-level dir name."""
    print(f"[rebuild] Extracting {tarball_path} ...")
    with tarfile.open(tarball_path, "r:gz") as tf:
        top = tf.getnames()[0].split("/")[0]
        tf.extractall(path=dest_dir)
    extracted = os.path.join(dest_dir, top)
    print(f"[rebuild] Extracted to: {extracted}")
    return extracted


def _ensure_dir_empty(path):
    if os.path.isdir(path):
        shutil.rmtree(path)
    os.makedirs(path, exist_ok=True)


def _run(cmd, cwd=None, desc=None):
    if desc:
        print(f"[rebuild] {desc} ...")
    print(f"[rebuild] Running: {' '.join(cmd[:4])} ...")
    ret = subprocess.run(cmd, cwd=cwd, capture_output=False, timeout=1800)
    if ret.returncode != 0:
        print(f"[rebuild] FAILED (exit={ret.returncode})")
        return False
    return True


def cross_compile_numpy(source_dir, build_dir, cross_file=None, jobs=None):
    """Cross-compile numpy extensions using vendored meson + OHOS cross file.

    Returns (success, build_dir).
    """
    cf = cross_file or CROSS_FILE
    if not os.path.isfile(cf):
        print(f"[rebuild] Cross file not found: {cf}")
        return False, build_dir

    _ensure_dir_empty(build_dir)

    vendored_meson = os.path.join(source_dir, "vendored-meson", "meson", "meson.py")
    if not os.path.isfile(vendored_meson):
        # try system meson as fallback
        meson_cmd = [sys.executable, "-m", "mesonbuild"]
        if not shutil.which("meson"):
            print(f"[rebuild] vendored meson not found at {vendored_meson}, and no system meson")
            return False, build_dir
        meson_cmd = ["meson"]
    else:
        meson_cmd = [sys.executable, vendored_meson]

    njobs = jobs or os.cpu_count() or 4

    # meson setup
    setup_args = meson_cmd + [
        "setup", build_dir,
        f"--cross-file={cf}",
        "-Dallow-noblas=true",
        "-Ddisable-threading=true",
        "-Ddisable-optimization=true",
    ]
    if not _run(setup_args, cwd=source_dir, desc="meson setup"):
        return False, build_dir

    # meson compile
    compile_args = meson_cmd + ["compile", "-C", build_dir, f"-j{njobs}"]
    if not _run(compile_args, desc="meson compile"):
        return False, build_dir

    print(f"[rebuild] Cross-compilation succeeded. Build dir: {build_dir}")
    return True, build_dir


def _should_include_source(fn, rel_path):
    """Decide whether a file from the source tree should go into the wheel."""
    ext = os.path.splitext(fn)[1]

    parts = rel_path.split("/")

    # Exclude tests and benchmarks
    if "tests" in parts or "benchmarks" in parts:
        return False

    # Exclude C/Cython source directories
    if "src" in parts and ext in (".py", ".pyx", ".pxd", ".pxi"):
        return False

    # Exclude build system files / source code
    if ext in (".c", ".h", ".cpp", ".cc", ".cxx", ".hpp", ".pyx", ".pxd", ".pxi"):
        return False
    if fn == "meson.build" or fn == "meson_options.txt":
        return False
    if fn.startswith("meson-"):
        return False
    if fn == "CMakeLists.txt" or fn.endswith(".cmake"):
        return False
    if fn.endswith(".in"):
        return False
    if fn == "setup.py" or fn == "setup.cfg":
        return False

    # Exclude __pycache__
    if "__pycache__" in parts:
        return False
    if fn.endswith(".pyc") or fn.endswith(".pyo"):
        return False

    # Exclude hidden files
    if fn.startswith("."):
        return False

    # Exclude top-level config/doc files
    if len(parts) == 1 and ext in (".cfg", ".ini", ".toml", ".txt", ".rst", ".md",
                                    ".yml", ".yaml", ".svg", ".png", ".jpg"):
        return False

    return True


def _soabi_for(target_os, arch):
    """Return the SOABI string for the target platform."""
    if target_os == "harmonyos" or target_os == "ohos":
        return f"cpython-312-{arch}-linux-ohos"
    return f"cpython-312-{arch}-linux-gnu"


def _host_soabi():
    """Detect the build machine's SOABI from sysconfig."""
    import sysconfig
    return sysconfig.get_config_var("EXT_SUFFIX").lstrip(".").rsplit(".", 1)[0]


def assemble_wheel(source_dir, build_dir, output_dir, pkg_name, version,
                   target_os="harmonyos", arch="arm64",
                   soabi_override=None):
    """Assemble a .whl from pure-Python source + cross-compiled .so files.

    Returns path to the created .whl file.
    """
    os.makedirs(output_dir, exist_ok=True)

    target_soabi = soabi_override or _soabi_for(target_os, arch)
    host_soabi = _host_soabi()

    plat_tag = f"{target_os}_{arch}"
    wheel_name = f"{pkg_name}-{version}-cp312-cp312-{plat_tag}.whl"
    wheel_path = os.path.join(output_dir, wheel_name)

    dist_info = f"{pkg_name}-{version}.dist-info"
    records = []

    def add_entry(zf, arcname, data_or_path):
        if isinstance(data_or_path, str) and os.path.isfile(data_or_path):
            with open(data_or_path, "rb") as f:
                data = f.read()
        elif isinstance(data_or_path, str):
            data = data_or_path.encode()
        else:
            data = data_or_path
        zf.writestr(arcname, data)
        digest = hashlib.sha256(data).digest()
        b64 = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
        records.append(f"{arcname},sha256={b64},{len(data)}")

    def add_symlink(zf, arcname, target):
        info = zipfile.ZipInfo(arcname)
        info.external_attr = 0o120777 << 16
        zf.writestr(info, target)
        records.append(f"{arcname},,0")

    with zipfile.ZipFile(wheel_path, "w", zipfile.ZIP_DEFLATED) as zf:
        # ── dist-info metadata ──────────────────────────
        add_entry(zf, f"{dist_info}/METADATA", (
            f"Metadata-Version: 2.1\n"
            f"Name: {pkg_name}\n"
            f"Version: {version}\n"
            f"Summary: {pkg_name} \u2014 EcoArk HarmonyOS experimental cross-build\n"
            f"Description: EcoArk HarmonyOS experimental cross-build for {target_os}/{arch}\n"
            f"  Built from source: {source_dir}\n"
            f"  Cross-compiled with: OHOS NDK clang 15.0.4\n"
            f"  This is an EXPERIMENTAL build \u2014 not yet formally released.\n"
            f"License: See package license\n"
            f"Requires-Python: >=3.9\n"
        ))
        add_entry(zf, f"{dist_info}/WHEEL", (
            f"Wheel-Version: 1.0\n"
            f"Generator: ecoark-pipeline 0.1\n"
            f"Root-Is-Purelib: false\n"
            f"Tag: cp312-cp312-{plat_tag}\n"
        ))
        add_entry(zf, f"{dist_info}/top_level.txt", f"{pkg_name}\n")
        add_entry(zf, f"{dist_info}/INSTALLER", "ecoark-pipeline\n")
        add_entry(zf, f"{dist_info}/entry_points.txt",
                  "[console_scripts]\nf2py = numpy.f2py.f2py2e:main\n")

        # ── Copy Python source files ─────────────────────
        pkg_src = os.path.join(source_dir, pkg_name)
        if os.path.isdir(pkg_src):
            for dirpath, dirnames, filenames in os.walk(pkg_src):
                rel = os.path.relpath(dirpath, source_dir).replace("\\", "/")
                # Prune __pycache__
                if "__pycache__" in dirnames:
                    dirnames.remove("__pycache__")
                for fn in filenames:
                    src = os.path.join(dirpath, fn)
                    arcname = f"{rel}/{fn}"
                    if _should_include_source(fn, arcname):
                        add_entry(zf, arcname, src)

        # Copy top-level .py files too (e.g., _config.py)
        for fn in os.listdir(source_dir):
            if fn.endswith(".py"):
                src = os.path.join(source_dir, fn)
                if os.path.isfile(src):
                    add_entry(zf, fn, src)

        # ── Copy compiled .so files (rename SOABI) ──────
        build_pkg = os.path.join(build_dir, pkg_name)
        so_count = 0
        if os.path.isdir(build_pkg):
            for dirpath, dirnames, filenames in os.walk(build_pkg):
                rel = os.path.relpath(dirpath, build_dir).replace("\\", "/")
                for fn in filenames:
                    if not fn.endswith(".so"):
                        continue
                    src = os.path.join(dirpath, fn)
                    new_fn = fn.replace(host_soabi, target_soabi)
                    arcname = f"{rel}/{new_fn}"
                    add_entry(zf, arcname, src)
                    so_count += 1

        # ── Also check for .so files in sub-packages directly
        # in the build dir (some builds put them differently)
        for dirpath, dirnames, filenames in os.walk(build_dir):
            rel = os.path.relpath(dirpath, build_dir).replace("\\", "/")
            for fn in filenames:
                if not fn.endswith(".so"):
                    continue
                if not fn.startswith(pkg_name) and rel not in (".", ""):
                    continue
                src = os.path.join(dirpath, fn)
                new_fn = fn.replace(host_soabi, target_soabi)
                arcname = f"{rel}/{new_fn}"
                # Only add if not already added
                if arcname not in [r.split(",")[0] for r in records]:
                    add_entry(zf, arcname, src)
                    so_count += 1

        if so_count == 0:
            print(f"[rebuild] WARNING: no .so files found in build dir {build_dir}")

        # ── RECORD (final) ─────────────────────────────
        record_content = "\n".join(records) + "\n"
        add_entry(zf, f"{dist_info}/RECORD", record_content)

    print(f"[rebuild] Wheel: {wheel_path}")
    print(f"[rebuild] Size  : {os.path.getsize(wheel_path)} bytes")
    print(f"[rebuild] .so   : {so_count} files (SOABI: {host_soabi} -> {target_soabi})")
    return wheel_path


def build_numpy_wheel(version="1.26.4", output_dir=None, rebuild=False):
    """Full numpy cross-compilation + wheel assembly pipeline.

    Steps:
      1. Download + extract numpy source sdist from PyPI
      2. Cross-compile with vendored meson + OHOS cross file
      3. Assemble wheel from Python source + cross-compiled .so files
    """
    if not _check_toolchain():
        return None

    out_dir = output_dir or OUTPUT_DIR
    sources = _ensure_dirs(SOURCES_DIR)

    # Download + extract
    tarball = _download_sdist("numpy", version, sources)
    src_dir = _extract_tarball(tarball, sources)

    build_dir = os.path.join(sources, f"numpy-{version}-build")
    if os.path.isdir(build_dir) and not rebuild:
        print(f"[rebuild] Build dir already exists: {build_dir}")
        print("[rebuild] Set rebuild=True or remove the dir to rebuild from scratch")
    else:
        ok, build_dir = cross_compile_numpy(src_dir, build_dir)
        if not ok:
            print("[rebuild] Cross-compilation FAILED")
            return None

    # Assemble wheel
    wheel_path = assemble_wheel(src_dir, build_dir, out_dir,
                                pkg_name="numpy", version=version,
                                target_os="harmonyos", arch="arm64")
    return wheel_path


def rebuild_from_source(source_dir, pkg_name, output_dir, version="0.0.0",
                        target_os="harmonyos", arch="arm64"):
    """Generic entry point: rebuild a source package for HarmonyOS.

    For now, dispatches to package-specific builders.
    Falls back to stub wheel generation for unknown packages.
    """
    if pkg_name == "numpy":
        # Use default numpy version if no real version detected (e.g. git clone failed)
        if version == "0.0.0":
            version = "1.26.4"
        return build_numpy_wheel(version=version, output_dir=output_dir)
    return _stub_wheel(source_dir, pkg_name, output_dir, version, target_os, arch)


def rebuild_native_wheel(wheel_path: str, output_dir: str,
                         target_os: str = "harmonyos", arch: str = "arm64") -> str:
    """Entry point for wheel-to-wheel rebuild (not yet implemented)."""
    print(f"[rebuild] Native wheel rebuild from wheel not yet supported: {wheel_path}")
    print("[rebuild] Use --manifest or rebuild_from_source() instead.")
    basename = os.path.basename(wheel_path)
    output_path = os.path.join(output_dir, basename)
    os.makedirs(output_dir, exist_ok=True)
    shutil.copy2(wheel_path, output_path)
    return output_path


def _stub_wheel(source_dir, pkg_name, output_dir, version, target_os, arch):
    """Generate a stub wheel with metadata only."""
    os.makedirs(output_dir, exist_ok=True)
    wheel_name = f"{pkg_name}-{version}-cp312-cp312-{target_os}_{arch}.whl"
    wheel_path = os.path.join(output_dir, wheel_name)
    dist_info = f"{pkg_name}-{version}.dist-info"
    with zipfile.ZipFile(wheel_path, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr(f"{dist_info}/METADATA",
                    f"Metadata-Version: 2.1\nName: {pkg_name}\nVersion: {version}\n"
                    f"Summary: EcoArk HarmonyOS rebuild (stub)\n")
        zf.writestr(f"{dist_info}/WHEEL",
                    f"Wheel-Version: 1.0\nGenerator: ecoark-pipeline\n"
                    f"Root-Is-Purelib: false\nTag: cp312-cp312-{target_os}_{arch}\n")
        zf.writestr(f"{dist_info}/RECORD", "")
        zf.writestr(f"{pkg_name}/__init__.py",
                    f"def hello(): return 'built by EcoArk pipeline (stub)'\n")
        zf.writestr(f"{pkg_name}/_ecoark_stub.txt",
                    f"Package: {pkg_name}\nSource: {source_dir}\n"
                    f"Target: {target_os}/{arch}\nStub wheel (replace with real build).\n")
    print(f"[rebuild] Stub wheel: {wheel_path}")
    return wheel_path


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "numpy":
        build_numpy_wheel(version=sys.argv[2] if len(sys.argv) > 2 else "1.26.4")
    else:
        print("Usage: python -m pipeline.rebuild numpy [version]")
        print("       python -m pipeline.rebuild numpy 1.26.4")