"""
pycutwm.compile_package.py
------------

Main routines for compiling the pycutwm package.

Functions:
    - write_grid_on_file: Overwrite the SetGrid.cuh file as long as the grid 
    were modified.
    - compile_cutwm: Compile the cuTWM CUDA/C++ code. 
 
Author: Alfredo Daniel Sánchez Rossi
Date: 2025-25-12
License: MIT
"""


import subprocess
import platform
import os
import json
import subprocess
import time
import glob
import re
import sys

from .utils import load_config, save_config

def write_grid_on_file(config_path='../config.json', cuTWM_path='mypath'):
    with open(config_path, 'r') as f:
        config = json.load(f)
    # Extract values from the file
    grid = config["grid"]["grid_points"]
    kernel = config["grid"]["kernel_dim"]

    NX = grid["NX"]
    NY = grid["NY"]
    NZ = grid["NZ"]
    NT = grid["NT"]
    BLKT = kernel["BLKT"]
    BLKX = kernel["BLKX"]
    BLKY = kernel["BLKY"]

    # Set global variables in header content
    content = (
        "#ifndef _SETGRIDCUH\n"
        "#define _SETGRIDCUH\n\n"
        f"const uint32_t NX = {NX};\n"
        f"const uint32_t NY = {NY};\n"
        f"const uint32_t NZ = {NZ};\n"
        f"#ifdef DISPERSION\n"
        f"const uint32_t NT = {NT};\n"
        f"#else\n"
        f"const uint32_t NT = 1;\n"
        f"#endif\n"
        f"const uint32_t BLKT = {BLKT};\n"
        f"const uint32_t BLKX = {BLKX};\n"
        f"const uint32_t BLKY = {BLKY};\n"
        f"const uint32_t SIZE = NX*NY*NT;\n"
        "\n\n#endif // _SETGRIDCUH\n"
    )

    # Write the header file
    with open(cuTWM_path+'headers/SetGrid.cuh', 'w') as f:
        f.write(content)
    # print(f"File {cuTWM_path+'headers/SetGrid.cuh'} filled with the corresponding grid values.")


def compile_cutwm(
    dispersion=False, twm_process="sfg",
    cuTWM_path='mypath', config_path='../config.json',
    compiler="nvcc", opt=3,
    hdf5_include_linux="/usr/include/hdf5/serial",
    hdf5_lib_linux="/usr/lib/x86_64-linux-gnu/hdf5/serial",
    hdf5_include_win="C:/HDF5/include",
    hdf5_lib_win="C:/HDF5/lib"
):
    """
    Compile the cuTWM executable file with the given flags,
    autodetecting Linux or Windows and using the appropriate linker flags.
    """
    # print("*-----------------------------------------------------------------*")
    
    write_grid_on_file(config_path=config_path, cuTWM_path=cuTWM_path)

    # Detect system
    system = platform.system().lower()

    # Auto-detect compute capability (works in both systems if nvidia-smi is in PATH)
    try:
        compute_cap = subprocess.check_output(
            "nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d '.'",
            shell=True
        ).decode().strip()
    except Exception as e:
        compute_cap = "61"  # fallback
        print("Warning: could not auto-detect compute capability, using 61 as default.")

    flags = [
        compiler, cuTWM_path+"cuTWM.cu", "-w",
        f"-O{opt}",
        f"-gencode=arch=compute_{compute_cap},code=sm_{compute_cap}",
        f"-gencode=arch=compute_{compute_cap},code=compute_{compute_cap}"
    ]

    # Inserta defines según flags
    if dispersion:
        flags.insert(2, "-DDISPERSION")

    if twm_process == "sfg":
        flags.insert(3, "-DSFG")
    elif twm_process == "shg":
        flags.insert(3, "-DSHG")
    elif twm_process == "opg":
        flags.insert(3, "-DOPG")
    else:
        raise RuntimeError("Compilation failed: invalid twm_process")

    # Flags por sistema operativo
    if system == "linux":
        flags += [
            "-lcufftw", "-lcufft",
            "-lhdf5_cpp", "-lhdf5_cpp", "-lhdf5", "-lcurl" , "-lz",
            "-o", cuTWM_path+"cuTWM" 
        ]
    elif system == "windows":
        flags += [
            f"-I{hdf5_include_win}",
            f"-L{hdf5_lib_win}",
            "hdf5_cpp.lib", "hdf5.lib",
            "-o", cuTWM_path+"cuTWM.exe"
        ]
    elif system == "darwin":  # MacOS
        raise RuntimeError("Compilation for MacOS is not implemented.")
    else:
        raise RuntimeError(f"Unsupported OS: {system}")

    print("*-----------------------------------------------------------------*")
    print("Compiling with:", " ".join(flags))
    result = subprocess.run(flags)
    if result.returncode != 0:
        raise RuntimeError("Compilation failed")
    print("Compilation succeeded.")