module LinearSolvePETScMPIExt

# MPI-parallel PETSc extension for LinearSolve.jl.
#
# Each MPI rank runs the same Julia code simultaneously.  When the system
# matrix is a `PSparseMatrix` (from PartitionedArrays.jl) this extension
# intercepts the key hooks exposed by `LinearSolvePETScExt` to build a
# distributed PETSc KSP solve:
#
#   • to_petsc_mat    → MatCreateMPIAIJWithArrays  (owned rows per rank)
#   • ensure_vecs!    → no-op (Vecs are created fresh in run_ksp!)
#   • run_ksp!        → VecCreateGhostWithArray (zero-copy) + KSPSolve
#   • store/check_sparse_pattern!, update_sparse_values! for PSparseMatrix
#
# Two backends are supported:
#   • MPIArray  — true MPI-parallel solve; `comm` comes from the partition.
#   • DebugArray — single-process fallback (e.g. unit tests with with_debug());
#                 uses MPI.COMM_SELF so PETSc still sees a valid communicator.
#
# After the solve, only the OWNED degrees-of-freedom of the solution PVector
# are updated.  Ghost values remain stale and must be synchronised by the
# caller via `PartitionedArrays.consistent!` if needed.

using PartitionedArrays: PartitionedArrays, PSparseMatrix, PVector,
    MPIArray, DebugArray, partition,
    own_length, ghost_length,
    own_to_local, local_to_global, own_to_global, ghost_to_global,
    consistent, own_values

using SparseMatricesCSR: SparseMatrixCSR, getrowptr, getcolval
using LinearAlgebra

using PETSc
import PETSc: MPI, LibPETSc

using LinearSolve: PETScAlgorithm, LinearSolve
using SciMLBase: ReturnCode, SciMLBase

const PETScExt = Base.get_extension(LinearSolve, :LinearSolvePETScExt)

# ── Local-partition accessor ──────────────────────────────────────────────────
#
# MPIArray stores each rank's local data in a `.item` field (accessed
# independently on every rank during an MPI run).
# DebugArray (single-process simulation) stores all parts in `.items`; with
# a single part we just take the first element.

_local_part(a::MPIArray) = a.item
_local_part(a::DebugArray) = a.items[1]

# ── Comm resolution ───────────────────────────────────────────────────────────
#
# Extract the MPI communicator from the PSparseMatrix's array partition.
# Returns COMM_SELF for DebugArray (single-process tests).

function PETScExt.get_comm(alg::PETScAlgorithm, A::PSparseMatrix)
    p = partition(A)
    return p isa MPIArray ? p.comm : MPI.COMM_SELF
end

# ── Sparse dispatch trait ─────────────────────────────────────────────────────

PETScExt.is_sparse_petsc(::PSparseMatrix) = true

# ── Local CSR construction ────────────────────────────────────────────────────
#
# Build the three arrays required by MatCreateMPIAIJWithArrays:
#   rowptr  — 0-based row offsets for the owned rows on this rank
#   colind  — 0-based GLOBAL column indices for every NZ in those rows
#   nzval   — scalar values in CSR order
#
# Only the first `m_own` rows of the local matrix are passed to PETSc.  Ghost
# rows (when present) are dropped — PETSc handles off-rank contributions
# through its internal assembly.
#
# The conversion to SparseMatrixCSR{0,PetscScalar,PetscInt} is inlined here
# so that the CSR{0} fast-path (common with GridapFEM matrices) skips the
# intermediate allocation on the colind loop.

function _build_mpi_csr_arrays(
        local_mat, row_idx, col_idx,
        ::Type{PetscInt}, ::Type{PetscScalar}
    ) where {PetscInt, PetscScalar}

    m_own = own_length(row_idx)

    # Convert the local sparse matrix to CSR{0, PetscScalar, PetscInt}.
    # For a SparseMatrixCSR{0,PetscScalar,PetscInt} this is a no-op.
    csr = convert(SparseMatrixCSR{0, PetscScalar, PetscInt}, local_mat)

    # Row pointer slice for owned rows only.
    rowptr = Vector{PetscInt}(getrowptr(csr)[1:(m_own + 1)])
    n_nz = Int(rowptr[end])

    # Global (0-based) column index translation.
    # local_to_global(col_idx)[j_local_1based] gives the 1-based global column.
    # csr.colval stores 0-based local column indices.
    cols_l2g = local_to_global(col_idx)
    colind = Vector{PetscInt}(undef, n_nz)
    src_cv = getcolval(csr)
    @inbounds for k in 1:n_nz
        colind[k] = PetscInt(cols_l2g[src_cv[k] + 1] - 1)
    end

    nzval = Vector{PetscScalar}(csr.nzval[1:n_nz])

    return rowptr, colind, nzval, m_own, n_nz
end

# ── Matrix construction ───────────────────────────────────────────────────────
#
# Called during KSP build (Case 1 / Case 2).
# If `pcache` is not nothing (system-matrix path), we stash all MPI-specific
# data in pcache.mpi_data so that:
#   • the GC anchor for rowptr/colind/nzval is maintained
#     (PETSc may borrow these arrays internally)
#   • subsequent Case 3 value-only updates can find the nzval buffer
#   • pattern-change detection has a fingerprint to compare against
#
# When `pcache` is nothing (preconditioner-matrix path) we simply build the
# matrix without storing metadata — Case 3 for a PSparseMatrix prec is not
# yet supported and will trigger a full KSP rebuild instead.

function PETScExt.to_petsc_mat(petsclib, A::PSparseMatrix, pcache)
    part = partition(A)
    comm = part isa MPIArray ? part.comm : MPI.COMM_SELF

    PetscInt = PETSc.inttype(petsclib)
    PetscScalar = PETSc.scalartype(petsclib)

    M, N = size(A)

    local_mat = _local_part(partition(A))
    row_idx = _local_part(partition(axes(A, 1)))
    col_idx = _local_part(partition(axes(A, 2)))

    rowptr, colind, nzval, m_own, n_nz =
        _build_mpi_csr_arrays(local_mat, row_idx, col_idx, PetscInt, PetscScalar)

    n_own_cols = own_length(col_idx)

    mat = LibPETSc.MatCreateMPIAIJWithArrays(
        petsclib, comm,
        PetscInt(m_own), PetscInt(n_own_cols),
        PetscInt(M), PetscInt(N),
        rowptr, colind, nzval
    )
    PETSc.assemble!(mat)

    # Store GC anchors + metadata in the cache when available.
    # We keep the rawlocal CSR structure (before global-col translation) for
    # cheap pattern-change detection in check_pattern_changed.
    if pcache !== nothing
        csr_det = convert(SparseMatrixCSR{0, PetscScalar, PetscInt}, _local_part(partition(A)))
        local_rowptr = Vector{Int}(getrowptr(csr_det)[1:(m_own + 1)])
        local_colval = Vector{Int}(getcolval(csr_det)[1:n_nz])

        pcache.mpi_data = (
            rowptr = rowptr,       # GC anchor (may be borrowed by PETSc)
            colind = colind,       # GC anchor
            nzval = nzval,        # buffer reused for Case 3 updates
            m_own = m_own,
            n_nz = n_nz,
            global_size = (M, N),
            local_rowptr = local_rowptr, # pre-global-translation, for detection
            local_colval = local_colval,
            comm = comm,
        )
    end

    return mat
end

# ── Sparsity pattern storage ──────────────────────────────────────────────────
#
# Called immediately after to_petsc_mat in build_ksp!.
# mpi_data is already populated by to_petsc_mat above, so this function
# just needs to copy the rowptr/colval fingerprint fields into prev_colptr /
# prev_rowval for the shared check_pattern_changed fast path.
# (prev_colptr / prev_rowval are checked to be `=== nothing` as the "never
# built" sentinel in the default serial logic; we reuse them here.)

function PETScExt.store_sparse_pattern!(pcache, A::PSparseMatrix)
    data = pcache.mpi_data
    data === nothing && return
    pcache.prev_colptr = data.local_rowptr
    pcache.prev_rowval = data.local_colval
    return nothing
end

# ── Sparsity pattern change detection ─────────────────────────────────────────

function PETScExt.check_pattern_changed(pcache, A::PSparseMatrix)
    # First solve: mpi_data not yet set.
    pcache.mpi_data === nothing && return true

    data = pcache.mpi_data

    # Fast global-size check.
    size(A) != data.global_size && return true

    local_mat = _local_part(partition(A))
    row_idx = _local_part(partition(axes(A, 1)))
    m_own = own_length(row_idx)
    m_own != data.m_own && return true

    # Detailed structure check for CSR{0} local matrices.
    if local_mat isa SparseMatrixCSR
        rp = getrowptr(local_mat)
        cv = getcolval(local_mat)
        old_rp = data.local_rowptr
        old_cv = data.local_colval

        # Identity shortcut: same array objects ⇒ same pattern.
        rp === old_rp && cv === old_cv && return false

        length(rp) < m_own + 1 && return true
        Int(rp[m_own + 1]) != data.n_nz && return true

        # Element-wise comparison of owned-row prefix.
        @views rp[1:(m_own + 1)] != old_rp && return true
        n_nz = data.n_nz
        length(cv) < n_nz && return true
        @views cv[1:n_nz] != old_cv && return true
        return false
    else
        # Conservative fallback for non-CSR local matrices: always rebuild.
        # Rebuilding the KSP is safe; only efficiency suffers.
        # On the common path (CSR{0} local mats) the fast check above is used.
        return true
    end
end

# ── In-place value update (Case 3) ───────────────────────────────────────────
#
# Only the system matrix (pcache.petsc_A) is supported — a separate
# PSparseMatrix preconditioner would need its own mpi_data slot, which is not
# yet implemented.
#
# Fast path for CSR{0} local matrices: copy nzval[1:n_nz] directly into the
# pre-allocated buffer, then call MatUpdateMPIAIJWithArray.
# General path: convert to CSR{0} first, then copy.

function PETScExt.update_sparse_values!(
        petsclib, PA, pcache, A::PSparseMatrix; assemble::Bool = true
    )
    if PA !== pcache.petsc_A
        error(
            "Case 3 value updates for a PSparseMatrix preconditioner matrix " *
                "are not yet supported. Use `alg.prec_matrix = nothing` and let " *
                "PETSc handle preconditioning internally."
        )
    end

    data = pcache.mpi_data
    data === nothing && error("mpi_data not initialised — this is an internal bug.")

    PetscScalar = PETSc.scalartype(petsclib)
    PetscInt = PETSc.inttype(petsclib)

    local_mat = _local_part(partition(A))
    n_nz = data.n_nz
    nzval_buf = data.nzval

    if local_mat isa SparseMatrixCSR{0}
        # Fast path: local nzval are already in PETSc row-major order for
        # owned rows — one copyto! from the first n_nz elements.
        copyto!(nzval_buf, view(local_mat.nzval, 1:n_nz))
    else
        # General path: convert to CSR{0,PetscScalar,PetscInt} first.
        csr = convert(SparseMatrixCSR{0, PetscScalar, PetscInt}, local_mat)
        copyto!(nzval_buf, view(csr.nzval, 1:n_nz))
    end

    LibPETSc.MatUpdateMPIAIJWithArray(petsclib, PA, nzval_buf)
    assemble && PETSc.assemble!(PA)
    return nothing
end

# ── Distributed vector creation ───────────────────────────────────────────────
#
# With VecCreateGhostWithArray, Vecs are created fresh in run_ksp! and
# alias the PVector's local storage directly (zero-copy).  This function
# only tracks the owned-size for partition-change detection and cleans up
# any stale Vecs left from a previous call path.

function PETScExt.ensure_vecs!(pcache, petsclib, comm, b::PVector)
    row_idx = _local_part(partition(axes(b, 1)))
    n_own = own_length(row_idx)
    if pcache.vec_n != n_own
        pcache.petsc_x !== nothing && PETSc.destroy(pcache.petsc_x)
        pcache.petsc_b !== nothing && PETSc.destroy(pcache.petsc_b)
        pcache.petsc_x = nothing
        pcache.petsc_b = nothing
        pcache.vec_n = n_own
    end
    return nothing
end

# ── KSP solve with zero-copy PVector I/O ─────────────────────────────────────
#
# Creates short-lived Vecs via VecCreateGhostWithArray that alias the
# PVector's local storage directly — no per-solve copy for b or x.
#
# Layout (PartitionedArrays OwnAndGhostIndices):
#   local array = [ owned_0 … owned_{n_own-1} | ghost_0 … ghost_{n_ghost-1} ]
#
# PETSc reads the RHS from local_b[1:n_own] and writes the solution back
# into local_u[1:n_own] in place.  Ghost values remain stale; the caller
# must call PartitionedArrays.consistent!(u) if ghost synchronisation is
# needed.
#
# Mixed-precision fallback: if PetscScalar ≠ eltype(local_u), a typed
# copy is made and the solution is written back manually after the solve.

function PETScExt.run_ksp!(pcache, petsclib, alg, b::PVector, u::PVector)
    PetscInt = PETSc.inttype(petsclib)
    PetscScalar = PETSc.scalartype(petsclib)

    data = pcache.mpi_data
    comm = data.comm

    row_idx = _local_part(partition(axes(b, 1)))
    n_own = PetscInt(own_length(row_idx))
    N = PetscInt(length(b))
    n_ghost = PetscInt(ghost_length(row_idx))
    ghosts = convert(Vector{PetscInt}, ghost_to_global(row_idx)) .- PetscInt(1)

    # Obtain PetscScalar-typed arrays aliasing the PVector's local storage.
    # When element types match, the `isa` check is a no-op (same array object).
    local_b = _local_part(partition(b))
    local_u = _local_part(partition(u))
    arr_b = local_b isa Vector{PetscScalar} ? local_b : Vector{PetscScalar}(local_b)
    arr_u = local_u isa Vector{PetscScalar} ? local_u : Vector{PetscScalar}(local_u)

    petsc_b = LibPETSc.VecCreateGhostWithArray(
        petsclib, comm, n_own, N, n_ghost, ghosts, arr_b
    )
    petsc_x = LibPETSc.VecCreateGhostWithArray(
        petsclib, comm, n_own, N, n_ghost, ghosts, arr_u
    )
    try
        if alg.transposed
            LibPETSc.KSPSolveTranspose(petsclib, pcache.ksp, petsc_b, petsc_x)
        else
            LibPETSc.KSPSolve(petsclib, pcache.ksp, petsc_b, petsc_x)
        end

        # When arr_u === local_u, PETSc wrote the solution directly into the
        # PVector's owned slots — nothing to copy back.
        # If a type-conversion buffer was used, copy the owned DOFs back manually.
        if arr_u !== local_u
            own_idxs = own_to_local(row_idx)
            @inbounds for (i, j) in enumerate(own_idxs)
                local_u[j] = arr_u[i]
            end
        end
    finally
        PETSc.destroy(petsc_x)
        PETSc.destroy(petsc_b)
    end
    return nothing
end

function PETScExt.postsolve_solution_check(cache, alg, pcache, A::PSparseMatrix, u::PVector)
    u_consistent = consistent(u, partition(axes(A, 2))) |> fetch
    Au = similar(u_consistent, axes(A, 1))
    LinearAlgebra.mul!(Au, A, u_consistent)

    local_Au = _local_part(own_values(Au))
    local_b = _local_part(own_values(cache.b))
    local_resid_sq = zero(typeof(abs2(zero(eltype(local_Au)))))
    local_b_sq = zero(typeof(abs2(zero(eltype(local_b)))))
    @inbounds for i in eachindex(local_Au, local_b)
        local_resid_sq += abs2(local_Au[i] - local_b[i])
        local_b_sq += abs2(local_b[i])
    end

    data = pcache.mpi_data
    comm = data === nothing ? MPI.COMM_SELF : data.comm
    resid_sq = MPI.Allreduce(local_resid_sq, +, comm)
    b_sq = MPI.Allreduce(local_b_sq, +, comm)
    res_norm = sqrt(resid_sq)
    b_norm = sqrt(b_sq)
    tol = cache.abstol + cache.reltol * b_norm

    if res_norm > tol
        return SciMLBase.build_linear_solution(
            alg, u, nothing, cache; retcode = ReturnCode.APosterioriSafetyFailure
        )
    end
    return nothing
end

end # module
