"""
vita.workout — the live training session

A ``Workout`` wraps a calibrated ``Lift`` and a VPFM ``Session``. It is the main
coaching surface: predict what a set will cost, ask how many reps to do for a
target proximity to failure, log sets, build a session plan, and read an
end-of-session summary.

The result objects (``Prediction``, ``LoggedSet``, ``Plan``, ``WorkoutSummary``)
are coaching readouts, not bare floats. How they display on their own is governed
by the output mode (vita.set_output): in the default ``summary`` mode, echoing a
result in a REPL or notebook shows its coaching block, so log_set() gives you
immediate feedback with no print() needed. print(obj) always shows the summary.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple

from .core import Params, Session
from .calibration import Lift
from .coaching import classify, RIRClass, render, _pct, _lb, _leader, _rule

__all__ = ["Prediction", "LoggedSet", "Plan", "WorkoutSummary", "Workout"]


# --------------------------------------------------------------------------
# result objects (coaching readouts, not bare floats)
# --------------------------------------------------------------------------
@dataclass
class Prediction:
    """What a set is expected to cost, before you do it."""

    lift: str
    weight: Optional[float]
    reps: float
    rir: float  # full precision
    effort: float  # full precision (reps / fatigued max)
    fatigued_max: float
    capacity: float
    _mode: Optional[str] = None  # set by the Workout; None -> global output mode

    @property
    def coaching(self) -> "RIRClass":
        return classify(self.rir)

    @property
    def status(self) -> str:
        return self.coaching.label

    @property
    def rpe(self) -> float:
        """RPE equivalent of the predicted RIR (RPE = 10 - RIR)."""
        return self.coaching.rpe

    @property
    def advice(self) -> str:
        """What to do about this set."""
        return self.coaching.advice

    def summary(self) -> str:
        c = self.coaching
        return "\n".join(
            [
                f"  {self.lift} — prediction",
                _rule(),
                _leader("Weight", _lb(self.weight)),
                _leader("Reps", f"{round(self.reps)}"),
                _leader("Predicted RIR", f"{round(self.rir, 1)}"),
                _leader("RPE", f"{c.rpe:g}"),
                _leader("Effort", _pct(self.effort)),
                _leader("Fatigued max", f"{round(self.fatigued_max, 1)} reps"),
                _leader("Capacity", _pct(self.capacity)),
                _rule(),
                f"  Status: {c.label}",
                f"  Coach:  {c.advice}",
            ]
        )

    def verbose(self) -> str:
        c = self.coaching
        fresh = self.fatigued_max / self.capacity if self.capacity else 0.0
        return "\n".join(
            [
                self.summary(),
                "  " + "- detail " + "-" * 33,
                _leader("Fresh max at this load", f"{round(fresh, 1)} reps"),
                f"  {c.note} (zone: {c.zone})",
            ]
        )

    def _oneline(self) -> str:
        return (
            f"<Prediction {self.lift} · {_lb(self.weight)} × {round(self.reps)}"
            f" → RIR {round(self.rir, 1)} · {self.status}>"
        )

    def __str__(self) -> str:
        return self.summary()

    def __repr__(self) -> str:
        return render(self, self._mode)


@dataclass
class LoggedSet:
    """A set you performed and logged."""

    number: int
    lift: str
    weight: Optional[float]
    reps: float
    rir: float
    effort: float
    capacity_after: float
    fatigued_max: float
    _mode: Optional[str] = None

    @property
    def coaching(self) -> "RIRClass":
        return classify(self.rir)

    @property
    def status(self) -> str:
        return self.coaching.label

    @property
    def rpe(self) -> float:
        """RPE equivalent of the set's RIR (RPE = 10 - RIR)."""
        return self.coaching.rpe

    @property
    def advice(self) -> str:
        """What to do next, given how this set went."""
        return self.coaching.advice

    def summary(self) -> str:
        c = self.coaching
        return "\n".join(
            [
                f"  Set {self.number} logged — {self.lift}",
                _rule(),
                _leader("Weight", _lb(self.weight)),
                _leader("Reps", f"{round(self.reps)}"),
                _leader("RIR", f"{round(self.rir, 1)}"),
                _leader("RPE", f"{c.rpe:g}"),
                _leader("Effort", _pct(self.effort)),
                _leader("Capacity now", _pct(self.capacity_after)),
                _rule(),
                f"  Status: {c.label}",
                f"  Coach:  {c.advice}",
            ]
        )

    def verbose(self) -> str:
        c = self.coaching
        return "\n".join(
            [
                self.summary(),
                "  " + "- detail " + "-" * 33,
                _leader(
                    "Fatigued max (set start)", f"{round(self.fatigued_max, 1)} reps"
                ),
                f"  {c.note} (zone: {c.zone})",
            ]
        )

    def _oneline(self) -> str:
        return (
            f"<LoggedSet #{self.number} {self.lift} · {_lb(self.weight)}"
            f" × {round(self.reps)} → RIR {round(self.rir, 1)} · {self.status}>"
        )

    def __str__(self) -> str:
        return self.summary()

    def __repr__(self) -> str:
        return render(self, self._mode)


@dataclass
class Plan:
    """A prescribed session. The shape depends on the planning mode: 'rir' holds
    a target RIR (reps fall as fatigue builds), 'fixed' holds the reps (RIR falls),
    and 'topset' is a top set plus reduced-load back-offs."""

    lift: str
    mode: str  # "rir" | "fixed" | "topset" | "ramp"
    rest: float
    rows: List[Tuple[int, Optional[float], int, float]]  # set, weight, reps, rir
    header: str
    _mode: Optional[str] = None  # output/display mode (set by the Workout)

    def summary(self) -> str:
        lines = [
            f"  {self.header}",
            _rule(46),
            "  " + f"{'Set':<6}{'Weight':<10}{'Reps':<8}{'RIR':<6}",
        ]
        for s, w, reps, rir in self.rows:
            lines.append("  " + f"{s:<6}{_lb(w):<10}{reps:<8}{rir:<6.1f}")
        lines.append(_rule(46))
        return "\n".join(lines)

    def verbose(self) -> str:
        lowest = min(self.rows, key=lambda row: row[3])
        c = classify(lowest[3])
        return "\n".join(
            [
                self.summary(),
                _leader("Rest between sets", f"{round(self.rest)} s", width=44),
                f"  Hardest set lands at {c.label.lower()} (RPE {c.rpe:g}).",
            ]
        )

    def _oneline(self) -> str:
        return f"<Plan {self.lift} · {self.mode} · {len(self.rows)} sets>"

    def __str__(self) -> str:
        return self.summary()

    def __repr__(self) -> str:
        return render(self, self._mode)

    def to_frame(self):
        """Return the plan as a pandas DataFrame (for .to_csv / .to_excel / .plot)."""
        import pandas as pd

        return pd.DataFrame(self.rows, columns=["set", "weight", "reps", "rir"])


@dataclass
class WorkoutSummary:
    """End-of-session readout."""

    lift: str
    total_sets: int
    total_reps: int
    volume_load: Optional[float]
    avg_rir: float
    lowest_rir: float
    avg_effort: float
    capacity_remaining: float
    hardest_set: int
    _mode: Optional[str] = None

    def summary(self) -> str:
        vol = (
            "—" if self.volume_load is None else f"{round(self.volume_load):,} lb-reps"
        )
        return "\n".join(
            [
                f"  Workout summary — {self.lift}",
                _rule(),
                _leader("Total sets", f"{self.total_sets}"),
                _leader("Total reps", f"{self.total_reps}"),
                _leader("Volume load", vol),
                _leader("Average RIR", f"{round(self.avg_rir, 1)}"),
                _leader("Lowest RIR", f"{round(self.lowest_rir, 1)}"),
                _leader("Average effort", _pct(self.avg_effort)),
                _leader("Capacity remaining", _pct(self.capacity_remaining)),
                _leader("Hardest set", f"#{self.hardest_set}"),
                _rule(),
            ]
        )

    def verbose(self) -> str:
        hardest = classify(self.lowest_rir)
        avg = classify(self.avg_rir)
        return "\n".join(
            [
                self.summary(),
                f"  Hardest set (#{self.hardest_set}) ran at {hardest.label} "
                f"(RPE {hardest.rpe:g})",
                f"  Session averaged a {avg.label.lower()}",
            ]
        )

    def _oneline(self) -> str:
        return (
            f"<WorkoutSummary {self.lift} · {self.total_sets} sets"
            f" · avg RIR {round(self.avg_rir, 1)}>"
        )

    def __str__(self) -> str:
        return self.summary()

    def __repr__(self) -> str:
        return render(self, self._mode)


# --------------------------------------------------------------------------
# Workout = a live training session
# --------------------------------------------------------------------------
class Workout:
    """One training session for a calibrated lift."""

    def __init__(
        self,
        lift: Lift,
        rest: float = 180.0,
        params: Optional["Params"] = None,
        output: Optional[str] = None,
    ):
        self.lift = lift
        self.default_rest = rest
        self.output = output  # None -> use the global output mode
        self._s = Session(lift.curve, params or Params(), capacity=1.0)
        self._effort_log: List[float] = []
        self._last: Optional[LoggedSet] = None

    def _stamp(self, result):
        """Tag a result object with this workout's output mode."""
        result._mode = self.output
        return result

    # --- how tired am I? ---
    @property
    def capacity(self) -> float:
        return self._s.capacity

    # --- log my set ---
    def log_set(
        self,
        weight: Optional[float] = None,
        reps: Optional[float] = None,
        rest: Optional[float] = None,
    ) -> LoggedSet:
        rest = self.default_rest if rest is None else rest
        r = self._s.do_set(reps=reps, load=weight, rest=rest)
        self._effort_log.append(r.effort)
        ls = LoggedSet(
            number=len(self._s.history),
            lift=self.lift.name,
            weight=weight,
            reps=reps,
            rir=r.rir,
            effort=r.effort,
            capacity_after=r.capacity_after_rest,
            fatigued_max=r.fatigued_max,
        )
        self._last = ls
        return self._stamp(ls)

    # --- how many reps should I do? ---
    def target_reps(self, weight: Optional[float] = None, rir: float = 0) -> int:
        return int(round(self._s.reps_for_rir(target_rir=rir, load=weight)))

    # --- what will this set cost? ---
    def predict_set(
        self, weight: Optional[float] = None, reps: Optional[float] = None
    ) -> Prediction:
        rir = self._s.predicted_rir(reps=reps, load=weight)
        fmax = self._s.capacity * self.lift.fresh_reps(weight)
        effort = (reps / fmax) if fmax > 0 else float("inf")
        return self._stamp(
            Prediction(
                lift=self.lift.name,
                weight=weight,
                reps=reps,
                rir=rir,
                effort=effort,
                fatigued_max=fmax,
                capacity=self._s.capacity,
            )
        )

    # --- build today's session ---
    def build_plan(
        self,
        weight: Optional[float] = None,
        target_rir: float = 2,
        sets: int = 5,
        rest: Optional[float] = None,
        percentage: Optional[float] = None,
        mode: str = "rir",
        reps: Optional[int] = None,
        backoff_pct: float = 10.0,
        weights: Optional[List[float]] = None,
    ) -> Plan:
        """Prescribe your sets, simulating fatigue set by set.

        Single weight (the default):
          mode="rir"     hold target_rir each set; reps fall as fatigue builds
          mode="fixed"   hold `reps` each set; the predicted RIR falls across sets
          mode="topset"  one top set to target_rir at `weight`, then back-off sets
                         at a load reduced by backoff_pct, each also to target_rir

        A weight per set (a ramp) -- pass weights=[...]:
          weights=[135, 185, 225, 275]                  vita fills reps to target_rir
          weights=[135, 185, 225, 275], reps=[8,6,4,3]  you set the reps; vita reads
                                                         the effort of each set

        The fatigue simulation runs on real-valued reps and rounds only for the
        printed plan, so the rep targets never tick back up between sets.
        """
        if weights is not None:
            return self._build_ramp(list(weights), reps, target_rir, rest)
        if mode not in ("rir", "fixed", "topset"):
            raise ValueError("mode must be 'rir', 'fixed', or 'topset'")
        if percentage is not None and weight is None and not self.lift.is_bodyweight:
            weight = round(self.lift.one_rep_max * percentage / 100.0)
        rest = self.default_rest if rest is None else rest

        # simulate on a copy from the CURRENT state; the live workout is untouched
        sim = Session(self.lift.curve, self._s.p, capacity=self._s.capacity)
        rows: List[Tuple[int, Optional[float], int, float]] = []

        if mode == "rir":
            for i in range(1, sets + 1):
                exact = sim.reps_for_rir(target_rir=target_rir, load=weight)
                sim.do_set(reps=exact, load=weight, rest=rest)  # real value (the fix)
                rows.append(
                    (i, weight, round(exact), target_rir)
                )  # round for display only
            header = (
                f"Session plan: {self.lift.name} @ {_lb(weight)}, "
                f"target {round(target_rir)} RIR"
            )

        elif mode == "fixed":
            if reps is None:
                raise ValueError(
                    "fixed mode needs reps= (the rep count to hold each set)"
                )
            for i in range(1, sets + 1):
                rir_pred = sim.predicted_rir(reps=reps, load=weight)
                sim.do_set(reps=reps, load=weight, rest=rest)
                rows.append((i, weight, int(reps), rir_pred))
            header = (
                f"Session plan: {self.lift.name} @ {_lb(weight)}, "
                f"{int(reps)}×{sets} fixed reps"
            )

        else:  # topset
            if self.lift.is_bodyweight:
                raise ValueError(
                    "topset mode needs a weighted lift (can't back off bodyweight)"
                )
            if weight is None:
                raise ValueError("topset mode needs a top-set weight= (or percentage=)")
            backoff_w = round(weight * (1.0 - backoff_pct / 100.0))
            for i in range(1, sets + 1):
                w_i = weight if i == 1 else backoff_w
                exact = sim.reps_for_rir(target_rir=target_rir, load=w_i)
                sim.do_set(reps=exact, load=w_i, rest=rest)
                rows.append((i, w_i, round(exact), target_rir))
            header = (
                f"Session plan: {self.lift.name}, top set @ {_lb(weight)} + "
                f"{sets - 1} back-off @ {_lb(backoff_w)}, target {round(target_rir)} RIR"
            )

        return self._stamp(
            Plan(lift=self.lift.name, mode=mode, rest=rest, rows=rows, header=header)
        )

    def _build_ramp(self, weights, reps, target_rir, rest) -> Plan:
        """A plan with an explicit weight per set. `reps` may be a list (one per
        weight), a single number (same for every set), or None (fill each set to
        target_rir). Fatigue carries set to set, same as the other modes."""
        if not weights:
            raise ValueError(
                "weights must be a non-empty list, e.g. weights=[135, 185, 225]"
            )
        rest = self.default_rest if rest is None else rest

        if reps is None:
            reps_list = [None] * len(weights)
        elif isinstance(reps, (list, tuple)):
            if len(reps) != len(weights):
                raise ValueError(
                    f"reps has {len(reps)} entries but weights has {len(weights)}; give one "
                    "rep count per weight, a single number for all, or leave reps out"
                )
            reps_list = list(reps)
        else:
            reps_list = [int(reps)] * len(weights)

        sim = Session(self.lift.curve, self._s.p, capacity=self._s.capacity)
        rows: List[Tuple[int, Optional[float], int, float]] = []
        for i, (w_i, r_i) in enumerate(zip(weights, reps_list), start=1):
            if r_i is None:
                exact = sim.reps_for_rir(target_rir=target_rir, load=w_i)
                n, rir_i = round(exact), float(target_rir)
                sim.do_set(reps=exact, load=w_i, rest=rest)
            else:
                n = int(r_i)
                rir_i = sim.predicted_rir(reps=n, load=w_i)
                sim.do_set(reps=n, load=w_i, rest=rest)
            rows.append((i, w_i, n, rir_i))

        header = f"Session plan: {self.lift.name}, {len(weights)} sets"
        return self._stamp(
            Plan(lift=self.lift.name, mode="ramp", rest=rest, rows=rows, header=header)
        )

    # --- how did the workout go? ---
    def summary(self) -> WorkoutSummary:
        hist = self._s.history
        if not hist:
            raise ValueError("no sets logged yet")
        rirs = [h.rir for h in hist]
        total_reps = int(sum(h.reps for h in hist))
        weights = [h.load for h in hist]
        if any(w is None for w in weights):
            vol = None
        else:
            vol = sum(h.load * h.reps for h in hist)
        hardest = min(range(len(hist)), key=lambda i: hist[i].rir) + 1
        return self._stamp(
            WorkoutSummary(
                lift=self.lift.name,
                total_sets=len(hist),
                total_reps=total_reps,
                volume_load=vol,
                avg_rir=sum(rirs) / len(rirs),
                lowest_rir=min(rirs),
                avg_effort=sum(self._effort_log) / len(self._effort_log),
                capacity_remaining=self._s.capacity,
                hardest_set=hardest,
            )
        )

    def finish(self) -> WorkoutSummary:
        """End the session and return the summary (same as summary())."""
        return self.summary()

    # --- quick access to the most recent set ---
    @property
    def last_set(self) -> Optional[LoggedSet]:
        return self._last

    @property
    def last_recommendation(self) -> Optional[str]:
        return self._last.advice if self._last is not None else None

    # --- backwards-compatible aliases (old vpfm names still work) ---
    def do_set(self, reps=None, load=None, rest=None) -> LoggedSet:
        return self.log_set(weight=load, reps=reps, rest=rest)

    def reps_for_rir(self, target_rir, load=None) -> int:
        return self.target_reps(weight=load, rir=target_rir)

    def predicted_rir(self, reps, load=None) -> float:
        return self.predict_set(weight=load, reps=reps).rir
