from __future__ import annotations

from typing import Optional, Union, Dict, Any

import sympy as sp


def solve_ode_yprime_equals_fx(
    f: 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:
    """
    Generates an HTML "full steps" solution for the ODE: dy/dx = f(x)  (depends only on x).

    Input:
      f: Python/SymPy expression for f(x), e.g. "x**2", "exp(x)", "sin(x)/x"

    Optional initial condition:
      y(x0) = y0  (provide both x0 and y0)

    Output:
      dict with keys: 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,
    }

    # Helper to LaTeX
    def L(expr: sp.Expr) -> str:
        return sp.latex(expr)

    if locals_map:
        base_locals.update(locals_map)

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

    # -----------------------------
    # 2) Core math: integrate
    # -----------------------------
    antiderivative_raw = sp.integrate(fx, x)
    antiderivative = sp.simplify(antiderivative_raw) if simplify else antiderivative_raw

    y_general = antiderivative + C
    C_value: Optional[sp.Expr] = None

    use_ic = (x0 is not None) and (y0 is not None)
    if use_ic:
        parsed_x0 = sp.sympify(x0, locals=base_locals)
        parsed_y0 = sp.sympify(y0, locals=base_locals)

        C_value_raw = parsed_y0 - antiderivative.subs(x, parsed_x0)
        C_value = sp.simplify(C_value_raw) if simplify else C_value_raw
        y_general = sp.simplify(antiderivative + C_value) if simplify else (antiderivative + C_value)

    # Detect if SymPy left an unevaluated integral
    is_unevaluated = bool(antiderivative.has(sp.Integral))

    # -----------------------------
    # 3) Build step-by-step HTML
    # -----------------------------
    # IMPORTANT: this is an f-string -> use doubled braces {{ }} when writing LaTeX braces.

    ode_ltx = rf"\frac{{d {L(yx)}}}{{d {L(x)}}} = {L(fx)}"
    int_step_ltx = rf"\int \frac{{d {L(yx)}}}{{d {L(x)}}}\, d{L(x)} = \int {L(fx)}\, d{L(x)}"
    left_integral_ltx = rf"{L(yx)} = \int {L(fx)}\, d{L(x)} + {L(C)}"

    if is_unevaluated:
        anti_note = (
            "SymPy kept the antiderivative as an unevaluated integral (this is normal for some functions). "
            "The solution is still correct in integral form."
        )
    else:
        anti_note = ""

    anti_eval_ltx = rf"{L(yx)} = {L(antiderivative)} + {L(C)}"
    final_general_ltx = rf"{L(yx)} = {L(antiderivative)} + {L(C)}"

    soln_stmts = []

    soln_stmt = f"""
    <div class="wrap">
      <div class="card">
        <div class="badge">Note: First-order ODE: depends only on {x_symbol}<br/><br/></div>

        $\\displaystyle
        \\begin{{aligned}}

        & \\displaystyle {ode_ltx} \\\\
        & \\displaystyle \\implies {int_step_ltx} \\\\
        & \\displaystyle \\implies {left_integral_ltx} \\\\
        & \\displaystyle \\implies {anti_eval_ltx} \\qquad \\cdots \\text{{ (i) }} \\\\
        \\end{{aligned}}
        $
        <div class="note">{anti_note}</div>
      </div>
    </div>
    """
    soln_stmts.append(soln_stmt.replace("\n", ""))

    if use_ic:
        # These exist only if use_ic is True
        parsed_x0 = sp.sympify(x0, locals=base_locals)
        parsed_y0 = sp.sympify(y0, locals=base_locals)

        C_ltx = rf"{L(C)} = {L(yx)} - {L(antiderivative)}"

        # ✅ Requested extra line (like your screenshot):
        #    C = y(x) - e^x
        #    C = y(0) - 1   at x=0
        # IMPORTANT: use \qquad and \text (single backslash) inside math.
        C_at_x0_ltx = (
            rf"{L(C)} = {L(yx.subs(x, parsed_x0))} - {L(antiderivative.subs(x, parsed_x0))}"
            rf" \qquad \text{{at }} {L(x)}={L(parsed_x0)}"
        )

        ic_ltx = rf"{L(yx)}\Big|_{{{L(x)}={L(parsed_x0)}}} = {L(parsed_y0)}"
        c_solve_ltx = rf"{L(C)} = {L(parsed_y0)} - {L(antiderivative.subs(x, parsed_x0))} = {L(C_value)}"
        final_ic_ltx = rf"{L(yx)} = {L(y_general)}"

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

            & \\displaystyle \\implies {C_ltx} \\\\
            & \\displaystyle \\implies {C_at_x0_ltx} \\\\
            & \\displaystyle \\implies {c_solve_ltx} \\qquad \\text{{Given }} {ic_ltx} \\\\
            & \\displaystyle \\implies {final_ic_ltx} \\qquad \\cdots \\text{{ from (i) }} \\\\

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

    if use_ic:
        problem_s = f"Solve \\( \\displaystyle {ode_ltx} \\) for \\( \\displaystyle {ic_ltx} \\)"
        answer_s = f"\\( \\displaystyle {final_ic_ltx} \\)"
    else:
        problem_s = f"Solve \\( {ode_ltx} \\)"
        answer_s = f"\\( \\displaystyle {final_general_ltx} \\)"

    soln_s = "".join(soln_stmts)

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

    return result
