from __future__ import annotations

import time
from dataclasses import dataclass, field
import warnings
from typing import Any, Callable, Iterable, List, Sequence, TYPE_CHECKING, TypeAlias

import numpy as np
import jax.numpy as jnp

from ..core.assembly import FormKernel, ResidualForm, assemble_bilinear_form
from ..core.solver import spdirect_solve_cpu, spdirect_solve_gpu
from .cg import cg_solve, cg_solve_jax
from .petsc import petsc_shell_solve
from .sparse import FluxSparseMatrix
from .dirichlet import DirichletBC, expand_dirichlet_solution
from .newton import newton_solve
from .result import SolverResult
from .history import NewtonIterRecord, LoadStepResult
from ..tools.timer import SectionTimer, NullTimer

if TYPE_CHECKING:
    from jax import Array as JaxArray

    ArrayLike: TypeAlias = np.ndarray | JaxArray
else:
    ArrayLike: TypeAlias = np.ndarray
DirichletLike: TypeAlias = tuple[np.ndarray, np.ndarray]
ExtraTerm: TypeAlias = Callable[
    [np.ndarray],
    tuple[np.ndarray, np.ndarray] | tuple[np.ndarray, np.ndarray, dict[str, Any]] | None,
]


@dataclass
class NonlinearAnalysis:
    """
    Bundle problem data needed for a Newton solve with load scaling.

    Attributes
    ----------
    space : Any
        FE space containing topology/dofs.
    residual_form : Any
        Internal residual form (e.g., neo_hookean_residual_form).
    params : Any
        Parameters forwarded to the residual form.
    base_external_vector : Any | None
        Unscaled external load vector (scaled by load factor in `external_for_load`).
    dirichlet : tuple | None
        (dofs, values) for Dirichlet boundary conditions.
    extra_terms : list[callable] | None
        Optional extra term assemblers returning (K, f[, metrics]).
    jacobian_pattern : Any | None
        Optional sparsity pattern to reuse between load steps.
    dtype : Any
        dtype for the solution vector (defaults to float64).
    """

    space: Any
    residual_form: ResidualForm[Any]
    params: Any
    base_external_vector: ArrayLike | None = None
    dirichlet: DirichletLike | None = None
    extra_terms: list[ExtraTerm] | None = None
    jacobian_pattern: Any | None = None
    dtype: Any = jnp.float64

    def __post_init__(self) -> None:
        if isinstance(self.dirichlet, DirichletBC):
            self.dirichlet = self.dirichlet.as_tuple()

    def external_for_load(self, load_factor: float) -> ArrayLike | None:
        if self.base_external_vector is None:
            return None
        return jnp.asarray(load_factor * self.base_external_vector, dtype=self.dtype)


@dataclass
class NewtonLoopConfig:
    """
    Control parameters for the Newton loop with load stepping.
    """

    tol: float = 1e-8
    atol: float = 0.0
    maxiter: int = 20
    line_search: bool = False
    max_ls: int = 10
    ls_c: float = 1e-4
    linear_solver: str = "spsolve"
    linear_maxiter: int | None = None
    linear_tol: float | None = None
    linear_preconditioner: Any | None = None
    matfree_mode: str = "linearize"
    load_sequence: Sequence[float] | None = None
    n_steps: int = 1

    def schedule(self) -> List[float]:
        """
        Return a monotonically increasing list of load factors (0->1].
        """
        if self.load_sequence is not None:
            return list(self.load_sequence)
        n = max(1, self.n_steps)
        return list(np.linspace(0.0, 1.0, n + 1, endpoint=True)[1:])


class NewtonSolveRunner:
    """
    Run one or more Newton solves across load factors.

    This orchestrates load stepping, assembles external load per step,
    and returns the full (Dirichlet-expanded) solution and per-step history.
    """

    def __init__(self, analysis: NonlinearAnalysis, config: NewtonLoopConfig):
        self.analysis = analysis
        self.config = config
        self._matfree_cache: dict[str, Any] = {}

    def run(
        self,
        u0: ArrayLike | None = None,
        *,
        load_sequence: Sequence[float] | None = None,
        newton_callback: Callable[[dict[str, Any]], None] | None = None,
        step_callback: Callable[[LoadStepResult], None] | None = None,
        timer: "SectionTimer | None" = None,
        report_timing: bool = True
    ) -> tuple[np.ndarray, list[LoadStepResult]]:
        """
        Execute Newton solves over the configured load schedule.

        Parameters
        ----------
        u0 : array-like | None
            Initial guess (defaults to zeros).
        load_sequence : sequence | None
            Optional per-call load schedule; if omitted, uses config.schedule().
        newton_callback : callable | None
            Per-iteration callback passed to `newton_solve`.
        step_callback : callable | None
            Optional hook called after each load step with LoadStepResult.
        Returns:
            u: full solution (Dirichlet expanded), dtype per analysis.dtype
            history: list of LoadStepResult per load factor
        """
        # timer = timer or NullTimer()
        timer = timer or SectionTimer(hierarchical=True)
        with timer.section("run_total"):
            dtype = self.analysis.dtype

            with timer.section("preprocess"):

                if u0 is None:
                    u = jnp.zeros(self.analysis.space.n_dofs, dtype=dtype)
                else:
                    u = jnp.asarray(u0, dtype=dtype)

                schedule_raw = list(load_sequence) if load_sequence is not None else self.config.schedule()
                # enforce monotone increasing 0->1 schedule, warn if dropped
                schedule = []
                dropped = []
                prev = 0.0
                for lf in schedule_raw:
                    lf_clamped = float(lf)
                    if not np.isfinite(lf_clamped):
                        dropped.append(("nonfinite", lf_clamped))
                        continue
                    if lf_clamped < 0.0 or lf_clamped > 1.0:
                        dropped.append(("out_of_range", lf_clamped))
                        continue
                    if lf_clamped < prev:
                        dropped.append(("nonmonotone", lf_clamped))
                        continue
                    schedule.append(lf_clamped)
                    prev = lf_clamped
                history: List[LoadStepResult] = []
                matfree_cache = None
                if self.config.linear_preconditioner == "diag0":
                    n_free = self.analysis.space.n_dofs
                    if self.analysis.dirichlet is not None:
                        n_free -= len(self.analysis.dirichlet[0])
                    cached_free = self._matfree_cache.get("n_free_dofs")
                    if cached_free is not None and cached_free != n_free:
                        self._matfree_cache.clear()
                    self._matfree_cache["n_free_dofs"] = n_free
                    matfree_cache = self._matfree_cache
                for step_i, load_factor in enumerate(schedule, start=1):
                    with timer.section("step"):
                        external = self.analysis.external_for_load(load_factor)
                        iter_log: List[NewtonIterRecord] = []
                        ext_nnz = int(jnp.count_nonzero(external))
                        ext_inf = int(jnp.count_nonzero(external))
                        # assumes u is full dofs (xyz)
                        u_nodes0 = np.asarray(u).reshape(-1, 3)
                        max_u0 = float(np.linalg.norm(u_nodes0, axis=1).max()) if u_nodes0.size else 0.0
                        print(
                            f"[load factor step {step_i}/{len(schedule)}] lf={load_factor:.3f} "
                            f"||F||_inf={ext_inf:.3e} nnz={ext_nnz} max|u|_start={max_u0:.3e}"
                        )

                        def cb(d):
                            # d: {"iter": k, "residual_norm": ..., "rel_residual": ..., "alpha": ..., "step_norm": ...}
                            lin_iters = d.get("linear_iters")
                            lin_res = d.get("linear_residual")
                            lin_conv = d.get("linear_converged")
                            iter_log.append(
                                NewtonIterRecord(
                                    iter=int(d.get("iter", -1)),
                                    res_inf=float(d.get("res_inf", float("nan"))),
                                    res_two=float(d.get("res_two", float("nan"))),
                                    rel_res_inf=float(d.get("rel_residual", float("nan"))),
                                    alpha=float(d.get("alpha", 1.0)),
                                    step_norm=float(d.get("step_norm", float("nan"))),
                                    lin_iters=int(lin_iters) if lin_iters is not None else None,
                                    lin_converged=bool(lin_conv) if lin_conv is not None else None,
                                    lin_residual=float(lin_res) if lin_res is not None else None,
                                    nan_detected=bool(d.get("nan_detected", False)),
                                )
                            )
                            if newton_callback is not None:
                                newton_callback(d)

                    try:
                        u, info = newton_solve(
                            self.analysis.space,
                            self.analysis.residual_form,
                            u,
                            self.analysis.params,
                            tol=self.config.tol,
                            atol=self.config.atol,
                            maxiter=self.config.maxiter,
                            linear_solver=self.config.linear_solver,
                            linear_maxiter=self.config.linear_maxiter,
                            linear_tol=self.config.linear_tol,
                            linear_preconditioner=self.config.linear_preconditioner,
                            matfree_cache=matfree_cache,
                            matfree_mode=self.config.matfree_mode,
                            dirichlet=self.analysis.dirichlet,
                            line_search=self.config.line_search,
                            max_ls=self.config.max_ls,
                            ls_c=self.config.ls_c,
                            external_vector=external,
                            callback=cb,
                            jacobian_pattern=self.analysis.jacobian_pattern,
                            extra_terms=self.analysis.extra_terms,
                        )
                        exception = None
                    except Exception as e:  # pragma: no cover - defensive
                        info = SolverResult(converged=False, iters=0, stop_reason="exception", nan_detected=False)
                        exception = repr(e)

                # ===== [B] OUTER LOOP PRINT (STEP END) =====
                u_nodes1 = np.asarray(u).reshape(-1, 3)
                max_u1 = float(np.linalg.norm(u_nodes1, axis=1).max()) if u_nodes1.size else 0.0
                step_solve_time = timer._records.get("step>newton_solve", [0.0])[-1]
                print(
                    f"  -> converged={getattr(info,'converged',None)} iters={getattr(info,'iters',None)} "
                    f"time={step_solve_time:.3f}s max|u|_end={max_u1:.3e}"
                    + (f" EXC={exception}" if exception else "")
                )

                meta = {
                    "load_factor": load_factor,
                    "linear_solver": self.config.linear_solver,
                    "line_search": self.config.line_search,
                    "maxiter": self.config.maxiter,
                    "n_dofs": self.analysis.space.n_dofs,
                    "dtype": str(self.analysis.dtype),
                    "u_layout": "full",
                    "schedule": schedule,
                    "schedule_dropped": dropped,
                }

                result = LoadStepResult(
                    load_factor=load_factor,
                    info=info,
                    solve_time=step_solve_time,
                    u=u,
                    iter_history=iter_log,
                    exception=exception,
                    meta=meta,
                )
                history.append(result)
                if step_callback is not None:
                    step_callback(result)

        if report_timing:
            timer.report(sort_by="total")
        return u, history


def _condense_flux_dirichlet(
    K: FluxSparseMatrix, F: ArrayLike, dirichlet: DirichletLike
) -> tuple[Any, np.ndarray, np.ndarray | None, np.ndarray, np.ndarray, np.ndarray]:
    dir_dofs, dir_vals = dirichlet
    dir_arr = np.asarray(dir_dofs, dtype=int)
    dir_vals_arr = np.asarray(dir_vals, dtype=float)
    K_csr = K.to_csr()
    mask = np.ones(K_csr.shape[0], dtype=bool)
    mask[dir_arr] = False
    free = np.nonzero(mask)[0]
    F_full = np.asarray(F, dtype=float)
    K_fd = K_csr[free][:, dir_arr] if dir_arr.size > 0 else None
    F_free_base = F_full[free]
    offset = K_fd @ dir_vals_arr if K_fd is not None and dir_arr.size > 0 else None
    K_ff = K_csr[free][:, free]
    return K_ff, F_free_base, offset, free, dir_arr, dir_vals_arr


def solve_nonlinear(
    space,
    residual_form: ResidualForm[Any],
    params: Any,
    *,
    dirichlet: DirichletLike | None = None,
    base_external_vector: ArrayLike | None = None,
    extra_terms: list[ExtraTerm] | None = None,
    dtype=jnp.float64,
    maxiter: int = 20,
    tol: float = 1e-8,
    atol: float = 1e-10,
    linear_solver: str = "spsolve",
    linear_maxiter: int | None = None,
    linear_tol: float | None = None,
    linear_preconditioner=None,
    matfree_mode: str = "linearize",
    line_search: bool = False,
    max_ls: int = 10,
    ls_c: float = 1e-4,
    n_steps: int = 1,
    jacobian_pattern=None,
    u0: ArrayLike | None = None,
) -> tuple[np.ndarray, list[LoadStepResult]]:
    """
    Convenience wrapper: build NonlinearAnalysis and run NewtonSolveRunner.
    """
    analysis = NonlinearAnalysis(
        space=space,
        residual_form=residual_form,
        params=params,
        base_external_vector=base_external_vector,
        dirichlet=dirichlet,
        extra_terms=extra_terms,
        dtype=dtype,
        jacobian_pattern=jacobian_pattern,
    )
    cfg = NewtonLoopConfig(
        maxiter=maxiter,
        tol=tol,
        atol=atol,
        linear_solver=linear_solver,
        linear_maxiter=linear_maxiter,
        linear_tol=linear_tol,
        linear_preconditioner=linear_preconditioner,
        matfree_mode=matfree_mode,
        line_search=line_search,
        max_ls=max_ls,
        ls_c=ls_c,
        n_steps=n_steps,
    )
    runner = NewtonSolveRunner(analysis, cfg)
    u0_use = jnp.zeros(space.n_dofs, dtype=dtype) if u0 is None else u0
    u, history = runner.run(u0=u0_use)
    return u, history


@dataclass
class LinearAnalysis:
    """
    Bundle linear problem data for a single solve or a load-scaled sequence.

    The matrix is assembled once from ``bilinear_form``; the RHS is scaled
    by the load factor.
    """

    space: Any
    bilinear_form: FormKernel[Any]
    params: Any
    base_rhs_vector: ArrayLike
    dirichlet: DirichletLike | None = None
    pattern: Any | None = None
    dtype: Any = jnp.float64

    def assemble_matrix(self):
        return self.space.assemble_bilinear_form(
            self.bilinear_form,
            params=self.params,
            pattern=self.pattern,
        )

    def rhs_for_load(self, load_factor: float) -> ArrayLike:
        return jnp.asarray(load_factor * self.base_rhs_vector, dtype=self.dtype)


@dataclass
class LinearSolveConfig:
    """
    Control parameters for the linear solve with optional load scaling.
    """

    method: str = "spsolve"  # "spsolve" | "spdirect_solve_gpu" | "cg" | "cg_custom" | "petsc_shell"
    tol: float = 1e-8
    maxiter: int | None = None
    preconditioner: Any | None = None
    ksp_type: str | None = None
    pc_type: str | None = None
    ksp_rtol: float | None = None
    ksp_atol: float | None = None
    ksp_max_it: int | None = None
    petsc_ksp_norm_type: str | None = None
    petsc_ksp_monitor_true_residual: bool = False
    petsc_ksp_converged_reason: bool = False
    petsc_ksp_monitor_short: bool = False
    petsc_shell_pmat: bool = False
    petsc_shell_pmat_mode: str = "full"
    petsc_shell_pmat_rebuild_iters: int | None = None
    petsc_shell_fallback: bool = False
    petsc_shell_fallback_ksp_types: tuple[str, ...] = ("bcgs", "gmres")
    petsc_shell_fallback_rebuild_pmat: bool = True

    @classmethod
    def from_preset(cls, name: str) -> "LinearSolveConfig":
        preset = name.lower()
        if preset == "contact":
            return cls(
                method="petsc_shell",
                ksp_type="bcgs",
                pc_type="ilu",
                petsc_shell_pmat=True,
                petsc_shell_pmat_mode="full",
                petsc_ksp_norm_type="unpreconditioned",
                petsc_ksp_monitor_true_residual=True,
                petsc_ksp_converged_reason=True,
                petsc_shell_fallback=True,
            )
        raise ValueError(f"Unknown LinearSolveConfig preset: {name}")


@dataclass
class LinearStepResult:
    """
    Result record for one linear solve step.

    Attributes
    ----------
    info : SolverResult
        Solver status and iteration metadata.
    solve_time : float
        Wall time for the solve section.
    u : Any
        Full solution vector (Dirichlet-expanded).
    """
    info: SolverResult
    solve_time: float
    u: ArrayLike


class LinearSolveRunner:
    """
    Solve linear systems for one or more load factors using a unified interface.
    """

    def __init__(self, analysis: LinearAnalysis, config: LinearSolveConfig):
        self.analysis = analysis
        self.config = config
        self._petsc_shell_pmat = None
        self._petsc_shell_last_iters = None
        self._petsc_shell_pmat_rebuilds = 0

    def run(
        self,
        *,
        step_callback: Callable[[LinearStepResult], None] | None = None,
        timer: "SectionTimer | None" = None,
        report_timing: bool = True
    ) -> tuple[np.ndarray | None, list[LinearStepResult]]:

        timer = timer or SectionTimer(hierarchical=True)
        # timer = timer or NullTimer()
        with timer.section("linear_run_total"):
            with timer.section("assemble_matrix"):
                K = self.analysis.assemble_matrix()

            with timer.section("build_rhs"):
                base_rhs = jnp.asarray(
                    self.analysis.base_rhs_vector, dtype=self.analysis.dtype
                )
                if self.analysis.dirichlet is not None:
                    K_ff, F_free_base, offset, free, dir_arr, dir_vals_arr = _condense_flux_dirichlet(
                        K, base_rhs, self.analysis.dirichlet
                    )
                    n_total = K.shape[0] if hasattr(K, "shape") else self.analysis.space.n_dofs
                else:
                    K_ff = K.to_csr()
                    F_free_base = np.asarray(base_rhs, dtype=float)
                    offset = None
                    free = None
                    dir_arr = dir_vals_arr = None
                    n_total = K_ff.shape[0]

                F_free = np.asarray(F_free_base, dtype=float)
                if offset is not None:
                    F_free = F_free - offset

            with timer.section(f"solve>{self.config.method}"):
                try:
                    if self.config.method == "spsolve":
                        u_free = spdirect_solve_cpu(K_ff, F_free)
                        lin_iters = 1
                        lin_conv = True
                        lin_res = None
                        info = SolverResult(
                            converged=True,
                            iters=lin_iters,
                            linear_iters=lin_iters,
                            linear_converged=lin_conv,
                            linear_residual=lin_res,
                            tol=self.config.tol,
                            stop_reason="converged",
                        )
                    elif self.config.method == "spdirect_solve_gpu":
                        u_free = spdirect_solve_gpu(K_ff, F_free)
                        lin_iters = 1
                        lin_conv = True
                        lin_res = None
                        info = SolverResult(
                            converged=True,
                            iters=lin_iters,
                            linear_iters=lin_iters,
                            linear_converged=lin_conv,
                            linear_residual=lin_res,
                            tol=self.config.tol,
                            stop_reason="converged",
                        )
                    elif self.config.method in ("cg", "cg_custom"):
                        coo = K_ff.tocoo()
                        A_cg = FluxSparseMatrix.from_bilinear(
                            (
                                jnp.asarray(coo.row, dtype=jnp.int64),
                                jnp.asarray(coo.col, dtype=jnp.int64),
                                jnp.asarray(coo.data),
                                K_ff.shape[0],
                            )
                        )
                        cg_solver = cg_solve_jax if self.config.method == "cg" else cg_solve
                        u_free, cg_info = cg_solver(
                            A_cg, jnp.asarray(F_free),
                            tol=self.config.tol,
                            maxiter=self.config.maxiter,
                            preconditioner=self.config.preconditioner
                        )
                        lin_iters = cg_info.get("iters")
                        lin_conv = bool(cg_info.get("converged", True))
                        lin_res = cg_info.get("residual_norm", cg_info.get("residual"))
                        info = SolverResult(
                            converged=lin_conv,
                            iters=int(lin_iters) if lin_iters is not None else 0,
                            linear_iters=int(lin_iters) if lin_iters is not None else None,
                            linear_converged=lin_conv,
                            linear_residual=float(lin_res) if lin_res is not None else None,
                            tol=self.config.tol,
                            stop_reason=("converged" if lin_conv else "linfail"),
                            nan_detected=bool(np.isnan(lin_res)) if lin_res is not None else False,
                        )
                    elif self.config.method == "petsc_shell":
                        base_ksp_type = self.config.ksp_type or "gmres"
                        pc_type = self.config.pc_type if self.config.pc_type is not None else "none"
                        ksp_rtol = self.config.ksp_rtol if self.config.ksp_rtol is not None else self.config.tol
                        ksp_atol = self.config.ksp_atol
                        ksp_max_it = self.config.ksp_max_it if self.config.ksp_max_it is not None else self.config.maxiter
                        petsc_options = {}
                        if self.config.petsc_ksp_norm_type:
                            petsc_options["fluxfem_ksp_norm_type"] = self.config.petsc_ksp_norm_type
                        if self.config.petsc_ksp_monitor_true_residual:
                            petsc_options["fluxfem_ksp_monitor_true_residual"] = ""
                        if self.config.petsc_ksp_converged_reason:
                            petsc_options["fluxfem_ksp_converged_reason"] = ""
                        if self.config.petsc_ksp_monitor_short:
                            petsc_options["fluxfem_ksp_monitor_short"] = ""
                        if not petsc_options:
                            petsc_options = None
                        use_pmat = bool(self.config.petsc_shell_pmat)
                        rebuild_thresh = self.config.petsc_shell_pmat_rebuild_iters
                        if use_pmat:
                            pmat_mode = (self.config.petsc_shell_pmat_mode or "full").lower()
                            if pmat_mode == "none":
                                use_pmat = False
                                pmat = None
                            elif pmat_mode == "full":
                                pmat = K_ff
                            else:
                                warnings.warn(
                                    f"petsc_shell_pmat_mode='{pmat_mode}' is not supported in runner; "
                                    "falling back to 'full'.",
                                    RuntimeWarning,
                                )
                                pmat = K_ff
                            if use_pmat:
                                if self._petsc_shell_pmat is None:
                                    self._petsc_shell_pmat = pmat
                                    self._petsc_shell_pmat_rebuilds += 1
                                elif rebuild_thresh is not None and self._petsc_shell_last_iters is not None:
                                    if self._petsc_shell_last_iters > rebuild_thresh:
                                        self._petsc_shell_pmat = pmat
                                        self._petsc_shell_pmat_rebuilds += 1
                                pmat = self._petsc_shell_pmat
                        if not use_pmat:
                            pmat = None

                        def _attempt_solve(ksp_type: str):
                            return petsc_shell_solve(
                                K_ff,
                                F_free,
                                preconditioner=self.config.preconditioner,
                                ksp_type=ksp_type,
                                pc_type=pc_type,
                                rtol=ksp_rtol,
                                atol=ksp_atol,
                                max_it=ksp_max_it,
                                pmat=pmat,
                                options=petsc_options,
                                return_info=True,
                            )

                        fallback_ksp = [base_ksp_type]
                        if self.config.petsc_shell_fallback:
                            for ksp in self.config.petsc_shell_fallback_ksp_types:
                                if ksp not in fallback_ksp:
                                    fallback_ksp.append(ksp)
                        fallback_attempts = []
                        petsc_info = None
                        u_free = None
                        for ksp in fallback_ksp:
                            fallback_attempts.append(ksp)
                            u_free, petsc_info = _attempt_solve(ksp)
                            lin_conv = petsc_info.get("converged")
                            reason = petsc_info.get("reason")
                            if lin_conv is None and reason is not None:
                                lin_conv = reason > 0
                            if lin_conv:
                                break
                            if self.config.petsc_shell_fallback and use_pmat and self.config.petsc_shell_fallback_rebuild_pmat:
                                self._petsc_shell_pmat = pmat
                                self._petsc_shell_pmat_rebuilds += 1
                        lin_iters = petsc_info.get("iters")
                        lin_res = petsc_info.get("residual_norm")
                        lin_solve_dt = petsc_info.get("solve_time")
                        pc_setup_dt = petsc_info.get("pc_setup_time")
                        pmat_dt = petsc_info.get("pmat_build_time")
                        lin_conv = petsc_info.get("converged")
                        if lin_conv is None and petsc_info.get("reason") is not None:
                            lin_conv = petsc_info.get("reason") > 0
                        if lin_conv is None:
                            lin_conv = True
                        self._petsc_shell_last_iters = lin_iters
                        info = SolverResult(
                            converged=bool(lin_conv),
                            iters=int(lin_iters) if lin_iters is not None else 0,
                            linear_iters=int(lin_iters) if lin_iters is not None else None,
                            linear_converged=bool(lin_conv),
                            linear_residual=float(lin_res) if lin_res is not None else None,
                            linear_solve_time=float(lin_solve_dt) if lin_solve_dt is not None else None,
                            pc_setup_time=float(pc_setup_dt) if pc_setup_dt is not None else None,
                            pmat_build_time=float(pmat_dt) if pmat_dt is not None else None,
                            pmat_rebuilds=self._petsc_shell_pmat_rebuilds if use_pmat else None,
                            pmat_mode=self.config.petsc_shell_pmat_mode if use_pmat else None,
                            tol=self.config.tol,
                            stop_reason=("converged" if lin_conv else "linfail"),
                            nan_detected=bool(np.isnan(lin_res)) if lin_res is not None else False,
                        )
                        if len(fallback_attempts) > 1:
                            info.linear_fallbacks = fallback_attempts
                    else:
                        raise ValueError(f"Unknown linear solve method: {self.config.method}")

                except Exception as e:  # pragma: no cover - defensive
                    exception = repr(e)
                    info = SolverResult(
                        converged=False,
                        iters=0,
                        linear_iters=None,
                        linear_converged=False,
                        linear_residual=None,
                        tol=self.config.tol,
                        stop_reason=f"exception: {exception}",
                        nan_detected=False,
                    )
                    # keep u as None to signal failure; caller should treat as Optional
                    result = LinearStepResult(info=info, solve_time=0.0, u=None)
                    if step_callback is not None:
                        step_callback(result)
                    # Report outside.
                    if report_timing:
                        timer.report(sort_by="total")
                    return None, [result], timer

            with timer.section("expand_dirichlet"):
                if free is None:
                    u_full = np.asarray(u_free, dtype=float)
                else:
                    u_full = expand_dirichlet_solution(
                        u_free, free, dir_arr, dir_vals_arr, n_total
                    )

            solve_key = f"solve>{self.config.method}"
            solve_time = timer._records.get(solve_key, [0.0])[-1]
            result = LinearStepResult(
                info=info, solve_time=solve_time, u=u_full
            )
            if step_callback is not None:
                step_callback(result)

        if report_timing:
            timer.report(sort_by="total")

        return u_full, [result]
