#!/usr/bin/env python3
"""
Build script for NimPacket Nim extension.

Compiles the Nim source into a Python extension module (.pyd on Windows, .so on Linux/macOS).

Requirements:
    - Nim compiler (1.6.0+): https://nim-lang.org/install.html
    - nimpy package: nimble install nimpy

Usage (after pip install):
    nimpacket-build             # Build the extension
    nimpacket-build --clean     # Clean build artifacts
    nimpacket-build --test      # Build and run tests

Usage (without pip install):
    python python/build.py      # Build the extension
"""

import os
import sys
import shutil
import subprocess
import platform
import glob as globmod


def get_extension_suffix():
    """Returns the correct extension suffix for the current platform."""
    if platform.system() == "Windows":
        return ".pyd"
    else:
        return ".so"


def find_nim():
    """Finds the Nim compiler."""
    nim_path = shutil.which("nim")
    if nim_path:
        return nim_path

    # Check common installation paths
    common_paths = [
        os.path.expanduser("~/.nimble/bin/nim"),
        os.path.expanduser("~/nim/bin/nim"),
        "/usr/local/bin/nim",
        "/usr/bin/nim",
    ]

    if platform.system() == "Windows":
        common_paths.extend([
            r"C:\Nim\bin\nim.exe",
            r"C:\Users\%USERNAME%\.nimble\bin\nim.exe",
        ])

    for path in common_paths:
        expanded = os.path.expandvars(path)
        if os.path.exists(expanded):
            return expanded

    return None


def check_nimpy():
    """Checks if nimpy is installed."""
    try:
        result = subprocess.run(
            ["nimble", "list", "-i"],
            capture_output=True,
            text=True
        )
        return "nimpy" in result.stdout
    except Exception:
        return False


def install_nimpy():
    """Installs nimpy using nimble."""
    print("Installing nimpy...")
    try:
        subprocess.run(
            ["nimble", "install", "nimpy", "-y"],
            check=True
        )
        return True
    except subprocess.CalledProcessError:
        return False


def _find_src_dir():
    """Locate the src/ directory containing Nim sources.

    Searches in order:
    1. Relative to this file (pip-installed with source included, or dev layout)
    2. Relative to current working directory
    """
    # When installed via pip or in dev layout: this file is at python/nimpacket/_build.py
    # src/ is at ../../src/ relative to this file's directory
    pkg_dir = os.path.dirname(os.path.abspath(__file__))
    candidate = os.path.normpath(os.path.join(pkg_dir, "..", "..", "src"))
    if os.path.isdir(candidate) and os.path.exists(os.path.join(candidate, "nimpacket_py.nim")):
        return candidate

    # Relative to CWD (user ran from project root)
    candidate = os.path.join(os.getcwd(), "src")
    if os.path.isdir(candidate) and os.path.exists(os.path.join(candidate, "nimpacket_py.nim")):
        return candidate

    return None


def build_extension():
    """Builds the Nim Python extension."""
    src_dir = _find_src_dir()
    if src_dir is None:
        print("ERROR: Cannot find src/nimpacket_py.nim")
        print("Make sure you run this from the NimPacket project root,")
        print("or that the source is available in a pip editable install.")
        return False

    nim_source = os.path.join(src_dir, "nimpacket_py.nim")

    # Output goes next to __init__.py in the nimpacket package
    pkg_dir = os.path.dirname(os.path.abspath(__file__))
    extension_suffix = get_extension_suffix()
    output_name = f"nimpacket_py{extension_suffix}"
    output_path = os.path.join(pkg_dir, output_name)

    # Check Nim is available
    nim_path = find_nim()
    if not nim_path:
        print("ERROR: Nim compiler not found!")
        print("Please install Nim from https://nim-lang.org/install.html")
        return False

    print(f"Found Nim at: {nim_path}")

    # Check nimpy is installed
    if not check_nimpy():
        print("nimpy not found, installing...")
        if not install_nimpy():
            print("ERROR: Failed to install nimpy")
            print("Try running: nimble install nimpy")
            return False

    print(f"Source:  {nim_source}")
    print(f"Output:  {output_path}")

    cmd = [
        nim_path, "c",
        "--app:lib",
        f"--out:{output_path}",
        "-d:release",
        "--opt:speed",
        "--threads:on",
        f"--path:{src_dir}",
        nim_source
    ]

    print(f"Building {output_name}...")

    try:
        result = subprocess.run(
            cmd,
            cwd=src_dir,
            capture_output=True,
            text=True
        )

        if result.returncode != 0:
            print("Build failed!")
            if result.stdout.strip():
                print("STDOUT:", result.stdout)
            if result.stderr.strip():
                print("STDERR:", result.stderr)
            return False

        print(f"Build successful! -> {output_path}")
        return True

    except Exception as e:
        print(f"Build error: {e}")
        return False


def clean():
    """Cleans build artifacts."""
    pkg_dir = os.path.dirname(os.path.abspath(__file__))

    patterns = ["*.pyd", "*.so", "*.dylib"]
    cleaned = 0

    for pattern in patterns:
        for f in globmod.glob(os.path.join(pkg_dir, pattern)):
            os.remove(f)
            print(f"Removed: {f}")
            cleaned += 1

    # Clean nimcache
    src_dir = _find_src_dir()
    if src_dir:
        nimcache = os.path.join(src_dir, "nimcache")
        if os.path.exists(nimcache):
            shutil.rmtree(nimcache)
            print(f"Removed: {nimcache}")
            cleaned += 1

    if cleaned == 0:
        print("Nothing to clean.")
    else:
        print(f"Cleaned {cleaned} items.")


def main():
    """Entry point for `nimpacket-build` console script."""
    if len(sys.argv) > 1:
        if sys.argv[1] in ("--clean", "clean"):
            clean()
        elif sys.argv[1] in ("--test", "test"):
            if build_extension():
                print("Build OK. Run your tests manually.")
        elif sys.argv[1] in ("--help", "-h", "help"):
            print(__doc__)
        else:
            print(f"Unknown option: {sys.argv[1]}")
            print("Usage: nimpacket-build [--clean | --test | --help]")
    else:
        build_extension()


if __name__ == "__main__":
    main()
