#!/usr/bin/env bash
# ============================================================================
# diagnose.sh — Diagnose linking and configuration issues
#
# Run this if verify.py fails. It prints detailed information about
# the environment that helps debug library conflicts.
#
# Usage:
#   conda activate firedrake
#   bash diagnose.sh
# ============================================================================

set -uo pipefail

echo "═══════════════════════════════════════════════════════════════"
echo "Firedrake conda environment diagnostics"
echo "═══════════════════════════════════════════════════════════════"

echo ""
echo "── Platform ─────────────────────────────────────────────────"
echo "OS:        $(uname -s)"
echo "Arch:      $(uname -m)"
echo "Kernel:    $(uname -r)"
if [[ "$(uname -s)" == "Linux" ]]; then
    echo "Distro:    $(cat /etc/os-release 2>/dev/null | grep PRETTY_NAME | cut -d= -f2 | tr -d '"')"
    echo "glibc:     $(ldd --version 2>&1 | head -1)"
fi

echo ""
echo "── Conda ────────────────────────────────────────────────────"
echo "conda:         $(conda --version 2>/dev/null || echo 'not found')"
echo "mamba:         $(mamba --version 2>/dev/null | head -1 || echo 'not found')"
echo "CONDA_PREFIX:  ${CONDA_PREFIX:-not set}"
echo "Active env:    ${CONDA_DEFAULT_ENV:-none}"

echo ""
echo "── Key environment variables ────────────────────────────────"
for var in PETSC_DIR PETSC_ARCH CC CXX FC HDF5_MPI HDF5_DIR \
           PYOP2_CFLAGS PYOP2_LDFLAGS OMP_NUM_THREADS \
           LD_LIBRARY_PATH DYLD_LIBRARY_PATH PATH; do
    val="${!var:-<unset>}"
    # Truncate long PATH values
    if [[ ${#val} -gt 120 ]]; then
        val="${val:0:120}..."
    fi
    printf "  %-22s %s\n" "$var" "$val"
done

echo ""
echo "── Compilers ────────────────────────────────────────────────"
for cmd in cc gcc mpicc mpicxx mpifort; do
    loc=$(which "$cmd" 2>/dev/null || echo "not found")
    printf "  %-10s %s\n" "$cmd" "$loc"
done

echo ""
echo "── MPI ────────────────────────────────────────────────────────"
if command -v mpicc &> /dev/null; then
    echo "mpicc -show:"
    mpicc -show 2>/dev/null || mpicc --showme 2>/dev/null || echo "  (could not determine)"
fi
if command -v mpiexec &> /dev/null; then
    mpiexec --version 2>&1 | head -3
fi

echo ""
echo "── Shared libraries in CONDA_PREFIX/lib ─────────────────────"
if [[ -n "${CONDA_PREFIX:-}" ]]; then
    for lib in libpetsc libslepc libmpi libhdf5 libopenblas libsuperlu_dist \
               libdmumps libscotch libmetis libparmetis; do
        found=$(ls "${CONDA_PREFIX}/lib/${lib}"* 2>/dev/null | head -1)
        if [[ -n "$found" ]]; then
            printf "  %-25s %s\n" "$lib" "$(basename "$found")"
        else
            printf "  %-25s %s\n" "$lib" "NOT FOUND"
        fi
    done
fi

echo ""
echo "── Python packages ──────────────────────────────────────────"
python3 -c "
import sys
print(f'  Python:    {sys.version}')
print(f'  Prefix:    {sys.prefix}')
print()

pkgs = [
    'numpy', 'mpi4py', 'petsc4py', 'slepc4py', 'h5py',
    'ufl', 'FIAT', 'finat', 'tsfc', 'loopy', 'pyop2',
    'firedrake', 'icepack',
]
for name in pkgs:
    try:
        mod = __import__(name)
        ver = getattr(mod, '__version__', '?')
        loc = getattr(mod, '__file__', '?')
        # Shorten path
        loc = loc.replace(sys.prefix, '\$PREFIX')
        print(f'  {name:15s} {ver:20s} {loc}')
    except ImportError as e:
        print(f'  {name:15s} IMPORT FAILED: {e}')
"

echo ""
echo "── petsc4py linking ─────────────────────────────────────────"
if [[ "$(uname -s)" == "Linux" ]]; then
    PETSC4PY_SO=$(python3 -c "import petsc4py; print(petsc4py.__file__)" 2>/dev/null)
    if [[ -n "$PETSC4PY_SO" ]]; then
        PETSC4PY_DIR=$(dirname "$PETSC4PY_SO")
        SO_FILE=$(find "$PETSC4PY_DIR" -name "PETSc*.so" 2>/dev/null | head -1)
        if [[ -n "$SO_FILE" ]]; then
            echo "  petsc4py .so: $SO_FILE"
            echo "  Links against:"
            LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${LD_LIBRARY_PATH:-}" ldd "$SO_FILE" 2>/dev/null | grep -E "petsc|mpi|hdf5|blas|lapack" | sed 's/^/    /'
        fi
    fi
elif [[ "$(uname -s)" == "Darwin" ]]; then
    PETSC4PY_SO=$(python3 -c "import petsc4py; print(petsc4py.__file__)" 2>/dev/null)
    if [[ -n "$PETSC4PY_SO" ]]; then
        PETSC4PY_DIR=$(dirname "$PETSC4PY_SO")
        SO_FILE=$(find "$PETSC4PY_DIR" -name "PETSc*.so" -o -name "PETSc*.dylib" 2>/dev/null | head -1)
        if [[ -n "$SO_FILE" ]]; then
            echo "  petsc4py .so: $SO_FILE"
            echo "  Links against:"
            otool -L "$SO_FILE" 2>/dev/null | grep -E "petsc|mpi|hdf5|blas|lapack" | sed 's/^/    /'
        fi
    fi
fi

echo ""
echo "── JIT compilation test ─────────────────────────────────────"
python3 -c "
import os, tempfile, subprocess, sys

cc = os.environ.get('CC', 'mpicc')
prefix = os.environ.get('CONDA_PREFIX', '')

# Write a minimal C file that links against PETSc
src = '''
#include <petsc.h>
int main(int argc, char **argv) {
    PetscInitialize(&argc, &argv, NULL, NULL);
    PetscPrintf(PETSC_COMM_WORLD, \"PETSc JIT test OK\\\\n\");
    PetscFinalize();
    return 0;
}
'''

with tempfile.NamedTemporaryFile(suffix='.c', mode='w', delete=False) as f:
    f.write(src)
    src_path = f.name

out_path = src_path.replace('.c', '')

# Build command similar to what PyOP2 would use
cmd = [
    cc,
    src_path,
    '-o', out_path,
    f'-I{prefix}/include',
    f'-L{prefix}/lib',
    '-lpetsc',
    f'-Wl,-rpath,{prefix}/lib',
]

print(f'  Compile command: {\" \".join(cmd)}')
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode == 0:
    print('  ✓ Compilation succeeded')
    run = subprocess.run([out_path], capture_output=True, text=True)
    if run.returncode == 0:
        print(f'  ✓ Execution succeeded: {run.stdout.strip()}')
    else:
        print(f'  ✗ Execution failed: {run.stderr.strip()}')
else:
    print(f'  ✗ Compilation failed:')
    print(f'    stdout: {result.stdout.strip()}')
    print(f'    stderr: {result.stderr.strip()}')

os.unlink(src_path)
try:
    os.unlink(out_path)
except:
    pass
" 2>&1

echo ""
echo "═══════════════════════════════════════════════════════════════"
echo "Done. If issues persist, share this output when seeking help."
echo "═══════════════════════════════════════════════════════════════"