"""Training analytics computed from parsed Sessions.

Pure functions over the normalized data from `parser.py`. Output is JSON-serializable
so it can be returned straight to the frontend and also fed to the chat model as context.
"""

from __future__ import annotations

from collections import defaultdict
from datetime import datetime, timedelta
from statistics import mean
from typing import Any

from .parser import Session


def estimated_1rm(weight_kg: float, reps: int) -> float:
    """Epley formula. Caps reps so very high-rep sets don't explode the estimate."""
    if weight_kg <= 0 or reps <= 0:
        return 0.0
    reps = min(reps, 20)
    return round(weight_kg * (1 + reps / 30.0), 1)


def _iso_week(dt: datetime) -> str:
    y, w, _ = dt.isocalendar()
    return f"{y}-W{w:02d}"


def summarize(sessions: list[Session]) -> dict[str, Any]:
    """Top-level totals for the dashboard header."""
    if not sessions:
        return {
            "total_sessions": 0,
            "total_sets": 0,
            "total_volume_kg": 0.0,
            "total_reps": 0,
            "first_session": None,
            "last_session": None,
            "avg_session_minutes": None,
            "unique_exercises": 0,
        }

    total_sets = sum(s.total_sets for s in sessions)
    total_volume = sum(s.volume_kg for s in sessions)
    total_reps = sum(
        st.reps or 0 for s in sessions for e in s.exercises for st in e.sets
    )
    durations = [s.duration_minutes for s in sessions if s.duration_minutes]
    dated = [s for s in sessions if s.start_time]
    unique_ex = {e.name for s in sessions for e in s.exercises}

    return {
        "total_sessions": len(sessions),
        "total_sets": total_sets,
        "total_volume_kg": round(total_volume, 1),
        "total_reps": total_reps,
        "first_session": dated[0].start_time.isoformat() if dated else None,
        "last_session": dated[-1].start_time.isoformat() if dated else None,
        "avg_session_minutes": round(mean(durations), 1) if durations else None,
        "unique_exercises": len(unique_ex),
    }


def session_volume_series(sessions: list[Session]) -> list[dict[str, Any]]:
    """Per-session volume / sets, in chronological order (for line charts)."""
    series = []
    for s in sessions:
        series.append({
            "date": s.start_time.isoformat() if s.start_time else None,
            "title": s.title,
            "volume_kg": round(s.volume_kg, 1),
            "sets": s.total_sets,
            "duration_minutes": s.duration_minutes,
        })
    return series


def weekly_volume(sessions: list[Session]) -> list[dict[str, Any]]:
    """Volume aggregated by ISO week."""
    buckets: dict[str, dict[str, float]] = defaultdict(lambda: {"volume_kg": 0.0, "sessions": 0, "sets": 0})
    for s in sessions:
        if not s.start_time:
            continue
        key = _iso_week(s.start_time)
        buckets[key]["volume_kg"] += s.volume_kg
        buckets[key]["sessions"] += 1
        buckets[key]["sets"] += s.total_sets
    return [
        {"week": k, "volume_kg": round(v["volume_kg"], 1),
         "sessions": int(v["sessions"]), "sets": int(v["sets"])}
        for k, v in sorted(buckets.items())
    ]


def muscle_group_balance(sessions: list[Session]) -> list[dict[str, Any]]:
    """Volume and set count grouped by muscle group, with share of total."""
    by_group: dict[str, dict[str, float]] = defaultdict(lambda: {"volume_kg": 0.0, "sets": 0})
    for s in sessions:
        for e in s.exercises:
            by_group[e.muscle_group]["volume_kg"] += e.volume_kg
            by_group[e.muscle_group]["sets"] += len(e.sets)
    total_sets = sum(v["sets"] for v in by_group.values()) or 1
    rows = [
        {
            "muscle_group": g,
            "volume_kg": round(v["volume_kg"], 1),
            "sets": int(v["sets"]),
            "set_share_pct": round(100 * v["sets"] / total_sets, 1),
        }
        for g, v in by_group.items()
    ]
    rows.sort(key=lambda r: r["sets"], reverse=True)
    return rows


def exercise_progress(sessions: list[Session]) -> list[dict[str, Any]]:
    """Per-exercise progression: best est-1RM, top set, trend, and PR history."""
    history: dict[str, list[dict[str, Any]]] = defaultdict(list)

    for s in sessions:
        if not s.start_time:
            continue
        for e in s.exercises:
            best_1rm = 0.0
            best_set = None
            for st in e.sets:
                if st.weight_kg and st.reps:
                    e1 = estimated_1rm(st.weight_kg, st.reps)
                    if e1 > best_1rm:
                        best_1rm = e1
                        best_set = st
            if best_set is not None:
                history[e.name].append({
                    "date": s.start_time.isoformat(),
                    "est_1rm_kg": best_1rm,
                    "top_weight_kg": best_set.weight_kg,
                    "top_reps": best_set.reps,
                    "muscle_group": e.muscle_group,
                })

    results = []
    for name, points in history.items():
        points.sort(key=lambda p: p["date"])
        first = points[0]["est_1rm_kg"]
        last = points[-1]["est_1rm_kg"]
        best = max(points, key=lambda p: p["est_1rm_kg"])
        delta = round(last - first, 1)
        results.append({
            "exercise": name,
            "muscle_group": points[0]["muscle_group"],
            "sessions": len(points),
            "first_est_1rm_kg": first,
            "latest_est_1rm_kg": last,
            "best_est_1rm_kg": best["est_1rm_kg"],
            "best_weight_kg": best["top_weight_kg"],
            "best_reps": best["top_reps"],
            "change_kg": delta,
            "change_pct": round(100 * delta / first, 1) if first else 0.0,
            "history": points,
        })
    results.sort(key=lambda r: r["sessions"], reverse=True)
    return results


def detect_plateaus(progress: list[dict[str, Any]], min_sessions: int = 3) -> list[dict[str, Any]]:
    """Flag exercises trained >= min_sessions whose est-1RM has stalled or dropped."""
    flags = []
    for p in progress:
        if p["sessions"] < min_sessions:
            continue
        recent = p["history"][-min_sessions:]
        values = [pt["est_1rm_kg"] for pt in recent]
        if max(values) - min(values) <= max(values) * 0.02:  # within 2%
            flags.append({
                "exercise": p["exercise"],
                "muscle_group": p["muscle_group"],
                "status": "plateau",
                "latest_est_1rm_kg": p["latest_est_1rm_kg"],
                "note": f"No meaningful change over the last {min_sessions} sessions.",
            })
        elif p["change_kg"] < 0:
            flags.append({
                "exercise": p["exercise"],
                "muscle_group": p["muscle_group"],
                "status": "regression",
                "latest_est_1rm_kg": p["latest_est_1rm_kg"],
                "note": f"Estimated 1RM down {abs(p['change_kg'])}kg vs first session.",
            })
    return flags


def training_frequency(sessions: list[Session]) -> dict[str, Any]:
    """Rough cadence: sessions per week over the logged span."""
    dated = [s.start_time for s in sessions if s.start_time]
    if len(dated) < 2:
        return {"sessions_per_week": None, "span_days": None, "days_since_last": None}
    span_days = (max(dated) - min(dated)).days or 1
    per_week = round(len(dated) / (span_days / 7.0), 2)
    days_since_last = (datetime.now() - max(dated)).days
    return {
        "sessions_per_week": per_week,
        "span_days": span_days,
        "days_since_last": max(days_since_last, 0),
    }


def build_report(sessions: list[Session]) -> dict[str, Any]:
    """Full analytics payload returned to the frontend."""
    progress = exercise_progress(sessions)
    return {
        "summary": summarize(sessions),
        "session_volume": session_volume_series(sessions),
        "weekly_volume": weekly_volume(sessions),
        "muscle_balance": muscle_group_balance(sessions),
        "exercise_progress": progress,
        "plateaus": detect_plateaus(progress),
        "frequency": training_frequency(sessions),
    }


def build_coach_context(sessions: list[Session], max_exercises: int = 12) -> str:
    """Compact, token-friendly text summary of the athlete for the chat model."""
    if not sessions:
        return "No workout data has been imported yet."

    report = build_report(sessions)
    s = report["summary"]
    freq = report["frequency"]
    lines: list[str] = ["ATHLETE TRAINING SUMMARY", ""]

    span = ""
    if s["first_session"] and s["last_session"]:
        span = f" between {s['first_session'][:10]} and {s['last_session'][:10]}"
    lines.append(
        f"- {s['total_sessions']} sessions{span}, {s['total_sets']} sets, "
        f"{s['total_volume_kg']:.0f} kg total volume, {s['unique_exercises']} distinct exercises."
    )
    if freq["sessions_per_week"] is not None:
        lines.append(
            f"- Frequency: ~{freq['sessions_per_week']} sessions/week; "
            f"{freq['days_since_last']} day(s) since last session."
        )

    lines.append("")
    lines.append("MUSCLE-GROUP BALANCE (by set share):")
    for row in report["muscle_balance"]:
        lines.append(
            f"  - {row['muscle_group']}: {row['sets']} sets ({row['set_share_pct']}%), "
            f"{row['volume_kg']:.0f} kg"
        )

    lines.append("")
    lines.append("KEY LIFT PROGRESSION (estimated 1RM):")
    for p in report["exercise_progress"][:max_exercises]:
        trend = "+" if p["change_kg"] >= 0 else ""
        lines.append(
            f"  - {p['exercise']} [{p['muscle_group']}]: now ~{p['latest_est_1rm_kg']}kg "
            f"(best {p['best_est_1rm_kg']}kg, {trend}{p['change_kg']}kg over {p['sessions']} sessions)"
        )

    if report["plateaus"]:
        lines.append("")
        lines.append("FLAGS:")
        for f in report["plateaus"]:
            lines.append(f"  - {f['exercise']} [{f['status']}]: {f['note']}")

    return "\n".join(lines)