# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Regression tests for plot_results.py and the training log format.

Guards against:
  Bug 2: plot_results.py reading "mean_reward" instead of "mean_return"
         (training logs use "mean_return"; plots were a flat line at 0.0)
  Bug 3: train_grpo.py not writing a local JSONL log
         (plot_results.py had nothing to read without W&B export)
  Bug 3b: training log missing "difficulty" key
          (tier-change markers never appeared on the plot)
"""

import json
import os
import sys
import tempfile
from pathlib import Path

import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


# ── Helpers ───────────────────────────────────────────────────────────────────

def _write_jsonl(path: str, records: list[dict]) -> None:
    with open(path, "w") as f:
        for r in records:
            f.write(json.dumps(r) + "\n")


def _write_eval_json(path: str, per_episode: list[dict]) -> None:
    """Write a baseline_eval-style JSON file."""
    summary = {
        "mean_reward": sum(e.get("reward", 0) for e in per_episode) / max(len(per_episode), 1),
        "mean_ndcg": sum(e.get("ndcg", 0.5) for e in per_episode) / max(len(per_episode), 1),
        "mean_fakes_in_top5": sum(e.get("fakes_in_top5", 1) for e in per_episode) / max(len(per_episode), 1),
    }
    with open(path, "w") as f:
        json.dump({"per_episode": per_episode, "summary": summary}, f)


# ── Test: training log key names ──────────────────────────────────────────────

class TestTrainingLogKeys:
    """Regression for Bug 2: wrong key name in plot_results.py."""

    def test_mean_return_key_is_read(self, tmp_path):
        """Reward data using 'mean_return' key (train_grpo.py format) must load correctly."""
        from recruiter.training.plot_results import _load_jsonl

        log_path = str(tmp_path / "training_log.jsonl")
        records = [
            {"update": 0, "episode": 4, "mean_return": -3.2, "difficulty": 1},
            {"update": 1, "episode": 8, "mean_return": -2.8, "difficulty": 1},
            {"update": 2, "episode": 12, "mean_return": -1.5, "difficulty": 2},
        ]
        _write_jsonl(log_path, records)

        loaded = _load_jsonl(log_path)
        rewards = [r.get("mean_return", r.get("mean_reward", 0.0)) for r in loaded]
        assert rewards == [-3.2, -2.8, -1.5], (
            f"Training log must be readable via 'mean_return' key. Got: {rewards}"
        )

    def test_mean_reward_key_fallback(self, tmp_path):
        """Legacy logs using 'mean_reward' key must still work (backward compat)."""
        from recruiter.training.plot_results import _load_jsonl

        log_path = str(tmp_path / "training_log.jsonl")
        records = [{"update": 0, "mean_reward": -4.1}]
        _write_jsonl(log_path, records)

        loaded = _load_jsonl(log_path)
        rewards = [r.get("mean_return", r.get("mean_reward", 0.0)) for r in loaded]
        assert rewards == [-4.1], (
            f"Fallback to 'mean_reward' key must work. Got: {rewards}"
        )

    def test_missing_both_keys_defaults_to_zero(self, tmp_path):
        """Missing reward key must default to 0.0, not crash."""
        from recruiter.training.plot_results import _load_jsonl

        log_path = str(tmp_path / "training_log.jsonl")
        records = [{"update": 0, "loss": 2.1}]
        _write_jsonl(log_path, records)

        loaded = _load_jsonl(log_path)
        rewards = [r.get("mean_return", r.get("mean_reward", 0.0)) for r in loaded]
        assert rewards == [0.0], "Missing reward key must default to 0.0"

    def test_difficulty_key_present_in_log(self, tmp_path):
        """Training log records must include 'difficulty' for tier-change markers."""
        from recruiter.training.plot_results import _load_jsonl

        log_path = str(tmp_path / "training_log.jsonl")
        records = [
            {"update": 0, "mean_return": -3.0, "difficulty": 1},
            {"update": 10, "mean_return": -1.0, "difficulty": 2},
        ]
        _write_jsonl(log_path, records)

        loaded = _load_jsonl(log_path)
        difficulties = [r.get("difficulty", -1) for r in loaded]
        assert difficulties == [1, 2], (
            f"'difficulty' key must be present in training log for tier markers. Got: {difficulties}"
        )

    def test_difficulty_key_missing_returns_sentinel(self, tmp_path):
        """Records without 'difficulty' default gracefully (no KeyError)."""
        from recruiter.training.plot_results import _load_jsonl

        log_path = str(tmp_path / "training_log.jsonl")
        records = [{"update": 0, "mean_return": -2.0}]  # no difficulty key
        _write_jsonl(log_path, records)

        loaded = _load_jsonl(log_path)
        # should default to 1 (or any fallback) — not crash
        difficulties = [r.get("difficulty", 1) for r in loaded]
        assert difficulties == [1], "Missing difficulty should default to 1"


# ── Test: plot() function doesn't crash with valid input ──────────────────────

class TestPlotFunction:
    """Smoke tests: plot() must not crash on valid and missing inputs."""

    def test_plot_training_curve_only(self, tmp_path):
        """plot() with training log only (no before/after eval) must not crash."""
        pytest.importorskip("matplotlib")
        from recruiter.training.plot_results import plot

        log_path = str(tmp_path / "training_log.jsonl")
        records = [
            {"update": i, "episode": i * 4, "mean_return": -3.0 + i * 0.1, "difficulty": 1}
            for i in range(20)
        ]
        _write_jsonl(log_path, records)

        out_path = str(tmp_path / "results.png")
        plot(training_log=log_path, output=out_path)
        assert Path(out_path).exists(), "plot() must write PNG output"

    def test_plot_with_before_after_eval(self, tmp_path):
        """plot() with baseline + trained eval must produce 3-subplot figure."""
        pytest.importorskip("matplotlib")
        from recruiter.training.plot_results import plot

        log_path = str(tmp_path / "training_log.jsonl")
        records = [
            {"update": i, "episode": i * 4, "mean_return": -3.0 + i * 0.05, "difficulty": 1}
            for i in range(10)
        ]
        _write_jsonl(log_path, records)

        per_ep = [
            {"episode": i, "reward": -3.5, "ndcg": 0.55, "fakes_in_top5": 2, "difficulty": 1}
            for i in range(20)
        ]
        baseline_path = str(tmp_path / "baseline.json")
        trained_path = str(tmp_path / "trained.json")
        _write_eval_json(baseline_path, per_ep)
        _write_eval_json(trained_path, [
            {"episode": i, "reward": -1.2, "ndcg": 0.85, "fakes_in_top5": 0, "difficulty": 1}
            for i in range(20)
        ])

        out_path = str(tmp_path / "results.png")
        plot(
            training_log=log_path,
            baseline_path=baseline_path,
            trained_path=trained_path,
            output=out_path,
        )
        assert Path(out_path).exists(), "plot() must write PNG with before/after data"

    def test_plot_no_inputs_doesnt_crash(self, tmp_path):
        """plot() with no valid inputs must write a placeholder PNG, not crash."""
        pytest.importorskip("matplotlib")
        from recruiter.training.plot_results import plot

        out_path = str(tmp_path / "results.png")
        plot(output=out_path)  # all optional inputs absent
        assert Path(out_path).exists(), "plot() must always write output even with no data"

    def test_tier_change_markers_detected(self, tmp_path):
        """difficulty changes in training log must be detectable for tier markers."""
        from recruiter.training.plot_results import _load_jsonl

        log_path = str(tmp_path / "training_log.jsonl")
        records = [
            {"update": i, "mean_return": float(i), "difficulty": 1 if i < 5 else 2}
            for i in range(10)
        ]
        _write_jsonl(log_path, records)

        loaded = _load_jsonl(log_path)
        diffs = [r.get("difficulty", 1) for r in loaded]
        changes = [i for i in range(1, len(diffs)) if diffs[i] != diffs[i - 1]]
        assert changes == [5], f"Tier change should be detected at index 5. Got changes at: {changes}"