"""
compile_nef.py — Convert a BIE file to a deployable NEF file using ktc.
"""

import os
import sys
import yaml
import subprocess

try:
    import ktc
    from ktc import ModelConfig
except ImportError:
    raise ImportError("ktc (Kneron Toolchain) is not installed in this environment.")

CONFIG_PATH = "config.yaml"

# ==========================================
# CONFIG LOADER
# ==========================================
def load_config(path: str = CONFIG_PATH) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)

# ==========================================
# CORE FUNCTION (importable)
# ==========================================
def compile_nef(
    bie_path:   str,
    onnx_path:  str,
    output_dir: str,
    platform:   str = "730",
    device_id:  int = 32770,
    os_id:      str = "8b28",
) -> str:
    
    # 1. Clean our own variables
    platform = str(platform).strip()
    os_id = str(os_id).strip()
    
    # 2. FIX: Convert EVERYTHING to absolute paths. 
    # C++ compilers often crash (SIGABRT) on relative paths!
    bie_path   = os.path.abspath(bie_path)
    onnx_path  = os.path.abspath(onnx_path)
    output_dir = os.path.abspath(output_dir).rstrip("/")
    
    os.makedirs(output_dir, exist_ok=True)

    print("=" * 70)
    print(f"  NEF Compilation")
    print(f"  BIE         : {bie_path}")
    print(f"  ONNX        : {onnx_path}")
    print(f"  Output dir  : {output_dir}")
    print(f"  Platform    : {platform}")
    print("=" * 70)

    if not os.path.exists(bie_path):
        raise FileNotFoundError(f"BIE file not found: {bie_path}")
    if not os.path.exists(onnx_path):
        raise FileNotFoundError(f"ONNX file not found: {onnx_path}")

    # 3. Fix the Conda python path issue
    python_bin_dir = os.path.dirname(sys.executable)
    if python_bin_dir not in os.environ.get("PATH", ""):
        os.environ["PATH"] = python_bin_dir + os.pathsep + os.environ.get("PATH", "")

    print(f"\n[Step 1] Creating ModelConfig...")
    km          = ModelConfig(device_id, os_id, platform, onnx_path=onnx_path)
    km.bie_path = bie_path
    print(f"  BIE path set: {bie_path}")

    print(f"\n[Step 2] Compiling to NEF...")

    # ── THE ULTIMATE MONKEY PATCH ──────────────────────────────────────────
    original_run = subprocess.run

    def patched_run(*args, **kwargs):
        clean_args = []
        for arg in args:
            if isinstance(arg, list):
                cmd_list = [a.strip() if isinstance(a, str) else a for a in arg]
                # FIX: Kneron leaves a dangling '-o' at the end of the command which causes SIGABRT
                if cmd_list and cmd_list[-1] == '-o':
                    cmd_list.append(output_dir)
                clean_args.append(cmd_list)
            else:
                clean_args.append(arg)
                
        if "args" in kwargs and isinstance(kwargs["args"], list):
            cmd_list = [a.strip() if isinstance(a, str) else a for a in kwargs["args"]]
            if cmd_list and cmd_list[-1] == '-o':
                cmd_list.append(output_dir)
            kwargs["args"] = cmd_list
            
        try:
            return original_run(*clean_args, **kwargs)
        except subprocess.CalledProcessError as e:
            # If the C++ compiler crashes, intercept the logs and print them clearly!
            print("\n" + "!" * 70)
            print("  FATAL C++ COMPILER CRASH DETECTED")
            print("!" * 70)
            log_file = os.path.join(output_dir, "batch_compile.log")
            if os.path.exists(log_file):
                print(f"  Found log file at {log_file}:\n")
                with open(log_file, "r") as f:
                    print(f.read())
            else:
                print("  No compiler log file found in the output directory.")
            print("!" * 70 + "\n")
            raise

    subprocess.run = patched_run
    # ───────────────────────────────────────────────────────────────────────

    try:
        nef_path = ktc.compile([km], output_dir=output_dir)
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise RuntimeError(f"NEF compilation failed: {e}") from e
    finally:
        subprocess.run = original_run

    nef_size_kb = os.path.getsize(nef_path) / 1024
    print(f"\n{'=' * 70}")
    print(f"  NEF generated: {nef_path}  ({nef_size_kb:.1f} KB)")
    print(f"{'=' * 70}")
    return nef_path

# ==========================================
# CLI ENTRY POINT
# ==========================================
def main():
    cfg = load_config()
    n   = cfg["bie_to_nef"]

    compile_nef(
        bie_path   = n["bie"],
        onnx_path  = n["onnx"],
        output_dir = n["output_dir"],
        platform   = n.get("platform", "730"),
    )

if __name__ == "__main__":
    main()