"""Tests for the training monitor."""

import json

import pytest

from orbit.monitor import (
    MonitorState,
    _detect_plateau,
    _loss_trend_icon,
    _parse_csv_logs,
    _parse_jsonl,
    parse_all_logs,
)


class TestLossTrendIcon:
    def test_dropping_normally(self):
        state = MonitorState()
        state.last_loss = 0.5
        state.loss_history = [0.5]
        icon = _loss_trend_icon(0.4, state)
        assert "dropping" in icon

    def test_stable(self):
        state = MonitorState()
        state.last_loss = 0.3
        state.loss_history = [0.3]
        icon = _loss_trend_icon(0.3, state)
        assert "stable" in icon.lower() or "convergence" in icon.lower()

    def test_increasing(self):
        state = MonitorState()
        state.last_loss = 0.3
        state.loss_history = [0.3]
        icon = _loss_trend_icon(0.5, state)
        assert "INCREASING" in icon

    def test_nan_exploded(self):
        state = MonitorState()
        state.last_loss = 0.3
        state.loss_history = [0.3]
        icon = _loss_trend_icon(float("nan"), state)
        assert "EXPLODED" in icon

    def test_empty_history(self):
        state = MonitorState()
        icon = _loss_trend_icon(0.5, state)
        assert "dropping" in icon


class TestPlateauDetection:
    def test_no_plateau_with_few_entries(self):
        state = MonitorState()
        state.loss_history = [0.5, 0.4, 0.3]
        state.step_history = [100, 200, 300]
        assert _detect_plateau(state, 300) is False

    def test_plateau_detected(self):
        state = MonitorState()
        state.loss_history = [0.3] * 15
        state.step_history = list(range(1000, 16000, 1000))
        result = _detect_plateau(state, 15000)
        assert result is True

    def test_no_plateau_when_dropping(self):
        state = MonitorState()
        state.loss_history = [0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.2, 0.15, 0.1, 0.05, 0.04]
        state.step_history = list(range(1000, 12000, 1000))
        assert _detect_plateau(state, 11000) is False

    def test_plateau_start_step_tracked(self):
        state = MonitorState()
        state.loss_history = [0.3] * 12
        state.step_history = list(range(1000, 13000, 1000))
        _detect_plateau(state, 12000)
        assert state.plateau_start_step > 0

    def test_plateau_resets_on_improvement(self):
        state = MonitorState()
        state.loss_history = [0.3] * 12
        state.step_history = list(range(1000, 13000, 1000))
        _detect_plateau(state, 12000)
        assert state.plateau_start_step > 0

        state.loss_history = [0.5, 0.4, 0.3, 0.2, 0.15, 0.1, 0.08, 0.06, 0.04, 0.03, 0.02]
        state.step_history = list(range(13000, 24000, 1000))
        _detect_plateau(state, 23000)
        assert state.plateau_start_step == 0


class TestJSONLParser:
    def test_parses_standard_jsonl(self, tmp_path):
        log_file = tmp_path / "metrics.jsonl"
        lines = [
            json.dumps({"step": 1000, "loss": 0.8, "grad_norm": 5.0}),
            json.dumps({"step": 2000, "loss": 0.5, "grad_norm": 2.0}),
            json.dumps({"step": 3000, "loss": 0.3}),
        ]
        log_file.write_text("\n".join(lines))

        entries = _parse_jsonl(tmp_path)
        assert len(entries) == 3
        assert entries[0]["step"] == 1000
        assert entries[0]["loss"] == 0.8
        assert entries[0]["grad_norm"] == 5.0
        assert entries[2]["step"] == 3000
        assert "grad_norm" not in entries[2]

    def test_handles_alternative_keys(self, tmp_path):
        log_file = tmp_path / "logs.jsonl"
        lines = [
            json.dumps({"global_step": 100, "train/loss": 0.9}),
            json.dumps({"_step": 200, "train_loss": 0.7, "train/grad_norm": 3.0}),
        ]
        log_file.write_text("\n".join(lines))

        entries = _parse_jsonl(tmp_path)
        assert len(entries) == 2
        assert entries[0]["step"] == 100
        assert entries[0]["loss"] == 0.9
        assert entries[1]["grad_norm"] == 3.0

    def test_skips_invalid_json(self, tmp_path):
        log_file = tmp_path / "bad.jsonl"
        log_file.write_text('{"step": 100, "loss": 0.5}\nnot valid json\n{"step": 200, "loss": 0.3}\n')

        entries = _parse_jsonl(tmp_path)
        assert len(entries) == 2

    def test_empty_dir(self, tmp_path):
        entries = _parse_jsonl(tmp_path)
        assert entries == []


class TestCSVParser:
    def test_parses_standard_csv(self, tmp_path):
        csv_file = tmp_path / "metrics.csv"
        csv_file.write_text("step,loss,grad_norm\n1000,0.8,5.0\n2000,0.5,2.0\n3000,0.3,1.0\n")

        entries = _parse_csv_logs(tmp_path)
        assert len(entries) == 3
        assert entries[0]["step"] == 1000
        assert entries[0]["loss"] == 0.8
        assert entries[0]["grad_norm"] == 5.0

    def test_handles_alternative_headers(self, tmp_path):
        csv_file = tmp_path / "train.csv"
        csv_file.write_text("global_step,train_loss\n100,0.9\n200,0.7\n")

        entries = _parse_csv_logs(tmp_path)
        assert len(entries) == 2
        assert entries[0]["step"] == 100

    def test_skips_non_metric_csv(self, tmp_path):
        csv_file = tmp_path / "config.csv"
        csv_file.write_text("name,value\nparam1,0.5\nparam2,0.3\n")

        entries = _parse_csv_logs(tmp_path)
        assert entries == []


class TestParseAllLogs:
    def test_combines_multiple_formats(self, tmp_path):
        # JSONL
        jsonl = tmp_path / "metrics.jsonl"
        jsonl.write_text(json.dumps({"step": 100, "loss": 0.9}) + "\n")

        # CSV
        csv_f = tmp_path / "train.csv"
        csv_f.write_text("step,loss\n200,0.8\n")

        # Text log (lerobot format)
        txt = tmp_path / "train.log"
        txt.write_text("step:300 loss:0.7 grdn:1.5\n")

        entries = parse_all_logs(tmp_path)
        assert len(entries) >= 2  # at least JSONL + CSV (text may also parse)
        steps = [e["step"] for e in entries]
        assert 100 in steps
        assert 200 in steps

    def test_deduplicates_by_step(self, tmp_path):
        # Same step in two formats
        jsonl = tmp_path / "metrics.jsonl"
        jsonl.write_text(json.dumps({"step": 100, "loss": 0.9}) + "\n")

        csv_f = tmp_path / "train.csv"
        csv_f.write_text("step,loss\n100,0.9\n")

        entries = parse_all_logs(tmp_path)
        steps = [e["step"] for e in entries]
        assert steps.count(100) == 1  # deduplicated
