#!/usr/bin/env python3
"""Verify that the Firedrake conda environment is correctly configured.

Checks:
  1. Environment variables (PETSC_DIR, CC, etc.)
  2. Conda-provided libraries (PETSc, MPI, HDF5)
  3. Python imports (petsc4py, mpi4py, firedrake, icepack)
  4. JIT compilation (the most likely failure point)
  5. A minimal PDE solve

Run with:
    conda activate firedrake
    python verify.py
"""

import os
import sys
import importlib
import ctypes
from pathlib import Path

PASS = "\033[92m✓\033[0m"
FAIL = "\033[91m✗\033[0m"
WARN = "\033[93m!\033[0m"
passed = 0
failed = 0
warned = 0


def check(description, condition, detail=""):
    global passed, failed
    if condition:
        print(f"  {PASS} {description}")
        passed += 1
        return True
    else:
        msg = f"  {FAIL} {description}"
        if detail:
            msg += f"  ({detail})"
        print(msg)
        failed += 1
        return False


def warn_check(description, detail=""):
    global warned
    msg = f"  {WARN} {description}"
    if detail:
        msg += f"  ({detail})"
    print(msg)
    warned += 1


# ═══════════════════════════════════════════════════════════════════════════
print("\n1. Environment variables")
print("─" * 40)

conda_prefix = os.environ.get("CONDA_PREFIX", "")
check("CONDA_PREFIX is set", bool(conda_prefix), conda_prefix or "not set")

petsc_dir = os.environ.get("PETSC_DIR", "")
check(
    "PETSC_DIR points to conda prefix",
    petsc_dir == conda_prefix,
    f"PETSC_DIR={petsc_dir!r}, expected {conda_prefix!r}"
)

petsc_arch = os.environ.get("PETSC_ARCH", "")
check(
    "PETSC_ARCH is empty (prefix install)",
    petsc_arch == "",
    f"PETSC_ARCH={petsc_arch!r}"
)

cc = os.environ.get("CC", "")
check("CC is set to mpicc", cc == "mpicc", f"CC={cc!r}")

hdf5_mpi = os.environ.get("HDF5_MPI", "")
check("HDF5_MPI=ON", hdf5_mpi == "ON", f"HDF5_MPI={hdf5_mpi!r}")

pyop2_cflags = os.environ.get("PYOP2_CFLAGS", "")
check("PYOP2_CFLAGS has -fPIC", "-fPIC" in pyop2_cflags, f"PYOP2_CFLAGS={pyop2_cflags!r}")

pyop2_ldflags = os.environ.get("PYOP2_LDFLAGS", "")
check("PYOP2_LDFLAGS has -shared", "-shared" in pyop2_ldflags, f"PYOP2_LDFLAGS={pyop2_ldflags!r}")

# ═══════════════════════════════════════════════════════════════════════════
print("\n2. Shared libraries")
print("─" * 40)

if conda_prefix:
    lib_dir = Path(conda_prefix) / "lib"
    ext = "dylib" if sys.platform == "darwin" else "so"

    for lib_name in ["libpetsc", "libmpi", "libhdf5"]:
        lib_path = lib_dir / f"{lib_name}.{ext}"
        if not lib_path.exists():
            # Try finding any matching file
            matches = list(lib_dir.glob(f"{lib_name}*"))
            found = len(matches) > 0
            detail = str(matches[0]) if found else f"no {lib_name}* in {lib_dir}"
        else:
            found = True
            detail = str(lib_path)
        check(f"{lib_name} found", found, detail)
else:
    warn_check("Skipping library checks (CONDA_PREFIX not set)")

# ═══════════════════════════════════════════════════════════════════════════
print("\n3. Python imports")
print("─" * 40)

modules = [
    ("numpy", "numpy"),
    ("mpi4py", "mpi4py"),
    ("petsc4py", "petsc4py"),
    ("h5py", "h5py"),
    ("ufl", "ufl"),
    ("FIAT", "FIAT"),
    ("finat", "finat"),
    ("loopy", "loopy"),
    ("tsfc", "tsfc"),
    ("pyop2", "pyop2"),
    ("firedrake", "firedrake"),
    ("icepack", "icepack"),
]

imported = {}
for display_name, module_name in modules:
    try:
        mod = importlib.import_module(module_name)
        version = getattr(mod, "__version__", "?")
        check(f"import {display_name}", True, f"v{version}")
        imported[module_name] = mod
    except Exception as e:
        check(f"import {display_name}", False, str(e))

# ═══════════════════════════════════════════════════════════════════════════
print("\n4. PETSc configuration")
print("─" * 40)

if "petsc4py" in imported:
    from petsc4py import PETSc as petsc

    # Check that PETSc has key external packages
    has_mumps = petsc.Sys.hasExternalPackage("mumps")
    check("PETSc has MUMPS", has_mumps)

    has_hypre = petsc.Sys.hasExternalPackage("hypre")
    check("PETSc has hypre", has_hypre)

    has_superlu = petsc.Sys.hasExternalPackage("superlu_dist")
    check("PETSc has SuperLU_dist", has_superlu)

    # Scalar type
    import numpy as np
    scalar = petsc.ScalarType
    is_real = scalar in (float, np.float64, np.float32)
    check("PETSc scalar type is real", is_real, f"scalar={scalar}")
else:
    warn_check("Skipping PETSc config checks (import failed)")

# ═══════════════════════════════════════════════════════════════════════════
print("\n5. JIT compilation test")
print("─" * 40)

if "firedrake" in imported:
    try:
        import firedrake

        # This triggers the full JIT pipeline:
        # UFL form → TSFC → loopy → C code → compile → dlopen
        mesh = firedrake.UnitSquareMesh(4, 4)
        V = firedrake.FunctionSpace(mesh, "CG", 1)
        u = firedrake.TrialFunction(V)
        v = firedrake.TestFunction(V)
        a = firedrake.inner(firedrake.grad(u), firedrake.grad(v)) * firedrake.dx
        A = firedrake.assemble(a)

        check("JIT compilation works", True, "assembled a Laplacian matrix")
    except Exception as e:
        check("JIT compilation works", False, str(e))
        # Try to find and display the actual compiler error
        import glob
        err_files = sorted(
            glob.glob("/tmp/pyop2-tempcache-*/**/*.err", recursive=True),
            key=lambda f: Path(f).stat().st_mtime,
            reverse=True,
        )
        log_files = sorted(
            glob.glob("/tmp/pyop2-tempcache-*/**/*.log", recursive=True),
            key=lambda f: Path(f).stat().st_mtime,
            reverse=True,
        )
        for label, files in [("COMPILE ERRORS", err_files), ("COMPILE LOG", log_files)]:
            if files:
                content = Path(files[0]).read_text().strip()
                if content:
                    print(f"\n    ── {label}: {files[0]} ──")
                    for line in content.splitlines()[-20:]:
                        print(f"    {line}")
else:
    warn_check("Skipping JIT test (firedrake import failed)")

# ═══════════════════════════════════════════════════════════════════════════
print("\n6. Minimal PDE solve")
print("─" * 40)

if "firedrake" in imported:
    try:
        import firedrake
        from firedrake import *

        mesh = UnitSquareMesh(8, 8)
        V = FunctionSpace(mesh, "CG", 1)

        u = TrialFunction(V)
        v = TestFunction(V)

        x, y = SpatialCoordinate(mesh)
        f = sin(pi * x) * sin(pi * y)

        a = inner(grad(u), grad(v)) * dx
        L = f * v * dx
        bc = DirichletBC(V, 0, "on_boundary")

        u_sol = Function(V)
        solve(a == L, u_sol, bcs=bc)

        # Check the solution is reasonable
        import numpy as np
        u_data = u_sol.dat.data_ro
        check(
            "Poisson solve produces valid output",
            np.all(np.isfinite(u_data)) and np.max(np.abs(u_data)) > 0,
            f"max|u| = {np.max(np.abs(u_data)):.6f}",
        )
    except Exception as e:
        check("Poisson solve", False, str(e))
else:
    warn_check("Skipping PDE solve (firedrake import failed)")

# ═══════════════════════════════════════════════════════════════════════════
print("\n7. icepack smoke test")
print("─" * 40)

if "icepack" in imported and "firedrake" in imported:
    try:
        import icepack
        import firedrake

        # Check that core model classes exist and instantiate
        model = icepack.models.IceShelf()
        check("icepack.models.IceShelf()", True)

        model = icepack.models.IceStream()
        check("icepack.models.IceStream()", True)

        # Check rate factor function
        A = icepack.rate_factor(260.0)
        check(
            "icepack.rate_factor(260 K)",
            A > 0,
            f"A = {A:.4e}",
        )
    except Exception as e:
        check("icepack smoke test", False, str(e))
else:
    warn_check("Skipping icepack test (import failed)")

# ═══════════════════════════════════════════════════════════════════════════
print("\n" + "═" * 50)
print(f"Results:  {PASS} {passed} passed   {FAIL} {failed} failed   {WARN} {warned} warnings")
print("═" * 50)

sys.exit(0 if failed == 0 else 1)