from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_first_order_linear_nonhomogeneous(
    P: Union[str, sp.Expr],
    Q: Union[str, sp.Expr],
    *,
    x_symbol: str = "x",
    y_symbol: str = "y",
    C_symbol: str = "C",
    x0: Optional[Union[int, float, str, sp.Expr]] = None,
    y0: Optional[Union[int, float, str, sp.Expr]] = None,
    locals_map: Optional[Dict[str, Any]] = None,
    simplify: bool = True,
) -> dict:
    """
    Full-steps HTML solution for the first-order linear nonhomogeneous ODE:
        y' + P(x) y = Q(x)

    Optional initial condition:
        y(x0) = y0

    Returns:
        { "problem_s": "...", "soln_s": "...", "answer_s": "..." }
    """

    # -----------------------------
    # 1) Setup symbols + parsing
    # -----------------------------
    x = sp.Symbol(x_symbol, real=True)
    yx = sp.Function(y_symbol)(x)
    C = sp.Symbol(C_symbol)

    base_locals = {
        x_symbol: x,
        "pi": sp.pi,
        "E": sp.E,
        "exp": sp.exp,
        "log": sp.log,
        "ln": sp.log,
        "sqrt": sp.sqrt,
        "sin": sp.sin,
        "cos": sp.cos,
        "tan": sp.tan,
        "asin": sp.asin,
        "acos": sp.acos,
        "atan": sp.atan,
        "sinh": sp.sinh,
        "cosh": sp.cosh,
        "tanh": sp.tanh,
        "Abs": sp.Abs,
        "sec": sp.sec,
        "csc": sp.csc,
        "cot": sp.cot,
    }
    if locals_map:
        base_locals.update(locals_map)

    def L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    try:
        P_ = sp.sympify(P, locals=base_locals) if isinstance(P, str) else sp.sympify(P)
        Q_ = sp.sympify(Q, locals=base_locals) if isinstance(Q, str) else sp.sympify(Q)
    except Exception as e:
        raise ValueError(f"Could not parse P,Q from inputs: P={P!r}, Q={Q!r}. Error: {e}") from e

    if simplify:
        P_ = sp.simplify(P_)
        Q_ = sp.simplify(Q_)

    # -----------------------------
    # 2) Integrating factor method
    # -----------------------------
    intP_raw = sp.integrate(P_, x)
    intP = sp.simplify(intP_raw) if simplify else intP_raw

    mu_raw = sp.exp(intP)
    mu = sp.simplify(mu_raw) if simplify else mu_raw

    muQ_raw = sp.simplify(mu * Q_) if simplify else (mu * Q_)
    I_raw = sp.integrate(muQ_raw, x)
    I = sp.simplify(I_raw) if simplify else I_raw

    # General solution: mu*y = I + C  => y = (I + C)/mu
    y_general = (I + C) / mu
    if simplify:
        y_general = sp.simplify(y_general)

    use_ic = (x0 is not None) and (y0 is not None)
    C_value: Optional[sp.Expr] = None
    y_particular: Optional[sp.Expr] = None

    if use_ic:
        x0_ = sp.sympify(x0, locals=base_locals)
        y0_ = sp.sympify(y0, locals=base_locals)

        # C = mu*y - I, then substitute x=x0, y=y0
        C_value_raw = (mu * yx - I).subs({x: x0_, yx: y0_})
        C_value = sp.simplify(C_value_raw) if simplify else C_value_raw

        y_particular_raw = (I + C_value) / mu
        y_particular = sp.simplify(y_particular_raw) if simplify else y_particular_raw

    # -----------------------------
    # 3) Build step-by-step HTML
    #    IMPORTANT: Avoid \text, \cdots, \implies (these break in many pipelines)
    #    Use \mathrm, "...", \Rightarrow instead.
    # -----------------------------
    ode_ltx = rf"\frac{{d {L(yx)}}}{{d {L(x)}}} + \left({L(P_)}\right){L(yx)} = {L(Q_)}"

    mu_def_ltx = rf"\mu({L(x)}) = e^{{\int {L(P_)}\, d{L(x)}}}"
    mu_val_ltx = rf"\mu({L(x)}) = e^{{{L(intP)}}} = {L(mu)}"

    mult_ltx = rf"{L(mu)}\frac{{d {L(yx)}}}{{d {L(x)}}} + {L(mu)}\left({L(P_)}\right){L(yx)} = {L(mu)}\left({L(Q_)}\right)"
    deriv_ltx = rf"\frac{{d}}{{d{L(x)}}}\left({L(mu)}\,{L(yx)}\right) = {L(mu)}\left({L(Q_)}\right)"

    integrate_ltx = rf"\int \frac{{d}}{{d{L(x)}}}\left({L(mu)}\,{L(yx)}\right)\, d{L(x)} = \int {L(mu)}\left({L(Q_)}\right)\, d{L(x)}"
    mu_y_eq_ltx = rf"{L(mu)}\,{L(yx)} = \int {L(mu)}\left({L(Q_)}\right)\, d{L(x)} + {L(C)}"

    I_ltx = rf"\int {L(mu)}\left({L(Q_)}\right)\, d{L(x)} = {L(I)}"
    gen_sol_ltx = rf"{L(yx)} = \frac{{{L(I)} + {L(C)}}}{{{L(mu)}}}"

    # Note for unevaluated integrals (no \text)
    uneval_parts = []
    if bool(intP.has(sp.Integral)):
        uneval_parts.append("We could not solve integral of P.")
    if bool(I.has(sp.Integral)):
        uneval_parts.append("")
    anti_note = " ".join(uneval_parts)

    soln_stmts = []

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: First-order linear nonhomogeneous ODE (Integrating Factor)<br/></div>

        \\(\\displaystyle
        \\begin{{aligned}}

        & \\displaystyle {ode_ltx} \\\\[6pt]

        & \\displaystyle \\mathrm{{Integrating\\ factor:}}\\; {mu_def_ltx} \\\\
        & \\displaystyle \\Rightarrow {mu_val_ltx} \\\\[6pt]

        & \\displaystyle \\mathrm{{Multiply\\ the\\ ODE\\ by}}\\; \\mu({L(x)}): \\\\
        & \\displaystyle \\Rightarrow {mult_ltx} \\\\
        & \\displaystyle \\Rightarrow {deriv_ltx} \\\\[6pt]

        & \\displaystyle \\mathrm{{Integrate\\ both\\ sides:}} \\\\
        & \\displaystyle \\Rightarrow {integrate_ltx} \\\\
        & \\displaystyle \\Rightarrow {mu_y_eq_ltx} \\\\[6pt]

        & \\displaystyle \\mathrm{{Compute\\ the\\ integral:}} \\\\
        & \\displaystyle \\Rightarrow {I_ltx} \\\\[6pt]

        & \\displaystyle \\Rightarrow {gen_sol_ltx}  \\qquad \\cdots (i) \\\\

        \\end{{aligned}}
        \\)
        <div class="note">{anti_note}</div>
      </div>
    </div>
    """
    soln_stmts.append(soln_stmt.replace("\n", ""))

    ic_ltx = ""
    final_ic_ltx = ""

    if use_ic and (y_particular is not None) and (C_value is not None):
        x0_ = sp.sympify(x0, locals=base_locals)
        y0_ = sp.sympify(y0, locals=base_locals)

        ic_ltx = rf"{L(yx)}\Big|_{{{L(x)}={L(x0_)}}} = {L(y0_)}"

        # C = mu*y - I
        C_expr_ltx = rf"{L(C)} = {L(mu)}\,{L(yx)} - {L(I)}"
        # Extra substitution line, like your screenshot
        C_at_x0_ltx = rf"{L(C)} = {L(mu.subs(x, x0_))}\,{L(y0_)} - {L(I.subs(x, x0_))}"
        C_val_ltx = rf"{L(C)} = {L(C_value)}"

        final_ic_ltx = rf"{L(yx)} = {L(y_particular)}"

        ic_block = f"""
            <br/>
            \\(\\displaystyle
            \\begin{{aligned}}

            & \\displaystyle \\Rightarrow {C_expr_ltx} \\\\
            & \\displaystyle \\Rightarrow {C_at_x0_ltx}  \\qquad \\mathrm{{at}}\\; {L(x)}={L(x0_)} \\\\
            & \\displaystyle \\Rightarrow {C_val_ltx}  \\qquad \\mathrm{{Given}}\\; {ic_ltx} \\\\
            & \\displaystyle \\Rightarrow {final_ic_ltx}  \\qquad \\cdots  \\mathrm{{from}} (i) \\\\

            \\end{{aligned}}
            \\)
        """
        soln_stmts.append(ic_block.replace("\n", ""))

    # -----------------------------
    # 4) Final problem/answer strings
    # -----------------------------
    if use_ic:
        problem_s = f"Solve \\( {ode_ltx} \\) for \\( {ic_ltx} \\)"
        answer_s = f"\\( {final_ic_ltx} \\)" if final_ic_ltx else f"\\( {L(yx)} = {L(y_general)} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( {L(yx)} = {L(y_general)} \\)"

    soln_s = "".join(soln_stmts)

    return {
        "problem_s": problem_s.strip(),
        "soln_s": soln_s.strip(),
        "answer_s": answer_s.strip(),
    }
