"""Tests for the AI-powered training debugger enhancements."""

import json
import math
from pathlib import Path
from unittest.mock import patch

import pytest

from orbit.debugger import (
    _analyze_dynamics,
    _llm_diagnosis,
    _parse_lerobot_logs,
    _rule_based_diagnosis,
)


class TestLLMDiagnosisWithDataset:
    """Test that dataset analysis results are passed to the LLM."""

    @patch("orbit.llm_utils.call_gemini")
    def test_dataset_context_included_in_prompt(self, mock_gemini):
        mock_gemini.return_value = "Root cause: test"

        dynamics = {
            "status": "plateau",
            "message": "Loss plateau",
            "total_steps": 50000,
            "loss_curve_summary": {"start": 0.8, "end": 0.31, "q25": 0.5, "q50": 0.4, "q75": 0.35},
            "findings_data": {"start_loss": 0.8, "end_loss": 0.31, "ratio": 0.39},
        }
        findings = [{"severity": "high", "finding": "plateau", "cause": "data"}]
        analysis_results = {
            "episode_count": 45,
            "consistency": {"overall": 0.62, "path_length_cv": 0.38},
            "signal_health": {"dead_dimensions": [6, 7], "clipping_joints": []},
            "action_divergence": 0.23,
            "outlier_episodes": [3, 17, 28],
        }

        result = _llm_diagnosis(
            dynamics, None, findings, "test/dataset", "fake-key",
            analysis_results=analysis_results,
        )

        assert result == "Root cause: test"
        # Check the prompt passed to Gemini includes dataset context
        call_args = mock_gemini.call_args
        prompt = call_args[0][0]
        assert "DATASET ANALYSIS" in prompt
        assert "45 episodes" in prompt
        assert "Dead joints: [6, 7]" in prompt
        assert "Action divergence: 0.23" in prompt
        assert "Outlier episodes: [3, 17, 28]" in prompt

    @patch("orbit.llm_utils.call_gemini")
    def test_no_dataset_context_when_none(self, mock_gemini):
        mock_gemini.return_value = "Root cause: test"

        dynamics = {
            "status": "converged_well",
            "message": "OK",
            "total_steps": 100000,
            "loss_curve_summary": {},
            "findings_data": {},
        }
        findings = []

        _llm_diagnosis(dynamics, None, findings, "test/dataset", "fake-key", analysis_results=None)

        prompt = mock_gemini.call_args[0][0]
        assert "DATASET ANALYSIS" not in prompt

    @patch("orbit.llm_utils.call_gemini")
    def test_readiness_grade_in_context(self, mock_gemini):
        mock_gemini.return_value = "diagnosis"

        dynamics = {
            "status": "plateau",
            "message": "Plateau",
            "total_steps": 50000,
            "loss_curve_summary": {},
            "findings_data": {},
        }
        analysis_results = {
            "readiness": {"grade": "D", "score": 35},
        }

        _llm_diagnosis(
            dynamics, None, [], "test/data", "key",
            analysis_results=analysis_results,
        )

        prompt = mock_gemini.call_args[0][0]
        assert "Readiness grade: D" in prompt


class TestDebugTrainingWithAI:
    """Test the debug_training function with --ai flag."""

    def test_debug_training_accepts_use_ai(self, tmp_path):
        """Verify the function signature accepts use_ai parameter."""
        from orbit.debugger import debug_training

        # Create a minimal log file
        log_file = tmp_path / "train.log"
        log_file.write_text("step:1000 loss:0.5 grdn:2.0\nstep:2000 loss:0.4 grdn:1.5\n")

        # Should not raise — just verify the parameter is accepted
        with patch("orbit.llm_utils.load_api_key", return_value=None):
            debug_training(str(tmp_path), use_ai=True)

    def test_debug_training_accepts_analysis_results(self, tmp_path):
        """Verify analysis_results parameter is accepted."""
        from orbit.debugger import debug_training

        log_file = tmp_path / "train.log"
        log_file.write_text("step:1000 loss:0.5 grdn:2.0\n")

        analysis = {"episode_count": 50, "consistency": {"overall": 0.8}}

        with patch("orbit.llm_utils.load_api_key", return_value=None):
            debug_training(str(tmp_path), analysis_results=analysis)
