"""Tests for the interactive AI assistant (agentic tool-using loop)."""

import json
import os
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest


# ---------------------------------------------------------------------------
# System prompt tests
# ---------------------------------------------------------------------------


class TestSystemPrompt:
    def test_contains_all_orbit_commands(self):
        from orbit.assist.system_prompt import build_system_prompt

        prompt = build_system_prompt()
        assert "orbit_analyze" in prompt
        assert "orbit_gate" in prompt
        assert "orbit_suggest" in prompt
        assert "orbit_debug" in prompt
        assert "orbit_verify" in prompt
        assert "orbit_convert" in prompt
        assert "orbit_clean" in prompt
        assert "orbit_curate" in prompt
        assert "scan_directory" in prompt
        assert "run_shell_command" in prompt
        assert "generate_config_file" in prompt
        assert "search_orbit_docs" in prompt

    def test_mentions_all_frameworks(self):
        from orbit.assist.system_prompt import build_system_prompt

        prompt = build_system_prompt()
        assert "LeRobot" in prompt
        assert "OpenVLA" in prompt
        assert "OpenPI" in prompt or "pi0" in prompt
        assert "GR00T" in prompt
        assert "OSMO" in prompt

    def test_includes_workspace_context(self):
        from orbit.assist.system_prompt import build_system_prompt

        ctx = "Detected datasets:\n  - ./data/ [lerobot (v3)]"
        prompt = build_system_prompt(ctx)
        assert "./data/ [lerobot (v3)]" in prompt

    def test_prompt_not_too_large(self):
        from orbit.assist.system_prompt import build_system_prompt

        prompt = build_system_prompt()
        # Should be under ~20K chars (well within any model's context)
        assert len(prompt) < 20000


# ---------------------------------------------------------------------------
# Context / directory scan tests
# ---------------------------------------------------------------------------


class TestDirectoryScan:
    def test_empty_directory(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        context = build_workspace_context(str(tmp_path))
        assert context["datasets"] == []
        assert context["training_outputs"] == []
        assert context["config_files"] == []
        assert "environment" in context

    def test_detects_lerobot_dataset(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        meta = tmp_path / "my_dataset" / "meta"
        meta.mkdir(parents=True)
        (meta / "info.json").write_text('{"codebase_version": "v3.0", "features": {}}')
        context = build_workspace_context(str(tmp_path))
        assert len(context["datasets"]) == 1
        assert context["datasets"][0]["format"] == "lerobot"
        assert context["datasets"][0]["version"] == "v3"

    def test_detects_lerobot_v2(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        meta = tmp_path / "ds" / "meta"
        meta.mkdir(parents=True)
        (meta / "info.json").write_text('{"codebase_version": "v2.0", "video_keys": []}')
        context = build_workspace_context(str(tmp_path))
        assert context["datasets"][0]["version"] == "v2"

    def test_detects_hdf5_files(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        (tmp_path / "data.hdf5").touch()
        context = build_workspace_context(str(tmp_path))
        assert any(d["format"] == "hdf5" for d in context["datasets"])

    def test_detects_mcap_files(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        (tmp_path / "recording.mcap").touch()
        context = build_workspace_context(str(tmp_path))
        assert any(d["format"] == "mcap" for d in context["datasets"])

    def test_detects_rlds_files(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        (tmp_path / "data.tfrecord").touch()
        context = build_workspace_context(str(tmp_path))
        assert any(d["format"] == "rlds" for d in context["datasets"])

    def test_detects_rosbag_files(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        (tmp_path / "recording.bag").touch()
        context = build_workspace_context(str(tmp_path))
        assert any(d["format"] == "rosbag" for d in context["datasets"])

    def test_detects_training_outputs(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        out_dir = tmp_path / "outputs" / "train"
        out_dir.mkdir(parents=True)
        (out_dir / "events.out.tfevents.12345").touch()
        context = build_workspace_context(str(tmp_path))
        assert len(context["training_outputs"]) == 1
        assert context["training_outputs"][0]["type"] == "tensorboard"

    def test_detects_config_files(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        (tmp_path / "train_config.yaml").touch()
        context = build_workspace_context(str(tmp_path))
        assert any("train_config.yaml" in f for f in context["config_files"])

    def test_environment_has_python_version(self, tmp_path):
        from orbit.assist.context import build_workspace_context

        context = build_workspace_context(str(tmp_path))
        assert "python_version" in context["environment"]
        assert "." in context["environment"]["python_version"]

    def test_format_context_for_prompt(self, tmp_path):
        from orbit.assist.context import build_workspace_context, format_context_for_prompt

        (tmp_path / "data.hdf5").touch()
        context = build_workspace_context(str(tmp_path))
        text = format_context_for_prompt(context)
        assert "hdf5" in text
        assert "Python:" in text


# ---------------------------------------------------------------------------
# Tool definitions tests
# ---------------------------------------------------------------------------


class TestToolDefinitions:
    def test_all_tools_have_declarations(self):
        from orbit.assist.tools import TOOL_DECLARATIONS, _EXECUTORS

        decl_names = {d["name"] for d in TOOL_DECLARATIONS}
        exec_names = set(_EXECUTORS.keys())
        assert decl_names == exec_names, f"Mismatch: {decl_names ^ exec_names}"

    def test_declarations_have_required_fields(self):
        from orbit.assist.tools import TOOL_DECLARATIONS

        for decl in TOOL_DECLARATIONS:
            assert "name" in decl, f"Missing name in {decl}"
            assert "description" in decl, f"Missing description in {decl.get('name')}"
            assert "parameters" in decl, f"Missing parameters in {decl['name']}"
            assert decl["parameters"]["type"] == "object"

    def test_tool_count(self):
        from orbit.assist.tools import TOOL_DECLARATIONS

        assert len(TOOL_DECLARATIONS) == 12

    def test_safety_classification_complete(self):
        from orbit.assist.tools import DESTRUCTIVE_TOOLS, SAFE_TOOLS, TOOL_DECLARATIONS

        all_tools = {d["name"] for d in TOOL_DECLARATIONS}
        classified = SAFE_TOOLS | DESTRUCTIVE_TOOLS
        assert all_tools == classified, f"Unclassified tools: {all_tools - classified}"


class TestToolExecutors:
    def test_unknown_tool_returns_error(self):
        from orbit.assist.tools import execute_tool

        result = execute_tool("nonexistent_tool", {})
        assert "error" in result

    def test_scan_directory_executor(self, tmp_path):
        from orbit.assist.tools import execute_tool

        result = execute_tool("scan_directory", {"path": str(tmp_path)})
        assert "datasets" in result
        assert "environment" in result

    def test_search_orbit_docs_executor(self):
        from orbit.assist.tools import execute_tool

        result = execute_tool("search_orbit_docs", {"query": "loss is NaN"})
        assert "relevant_entries" in result
        assert "related_commands" in result

    def test_search_finds_nan_issue(self):
        from orbit.assist.tools import execute_tool

        result = execute_tool("search_orbit_docs", {"query": "loss is NaN training"})
        entries = result.get("relevant_entries", [])
        assert len(entries) > 0
        # Should find the NaN troubleshooting entry
        assert any("NaN" in e["symptom"] or "nan" in e["symptom"].lower() for e in entries)

    def test_shell_command_blocks_dangerous(self):
        from orbit.assist.tools import execute_tool

        result = execute_tool("run_shell_command", {"command": "rm -rf /"})
        assert "error" in result
        assert "Blocked" in result["error"]

    def test_shell_command_runs_safe(self):
        from orbit.assist.tools import execute_tool

        result = execute_tool("run_shell_command", {"command": "echo hello"})
        assert result.get("exit_code") == 0
        assert "hello" in result.get("stdout", "")


# ---------------------------------------------------------------------------
# Knowledge base tests
# ---------------------------------------------------------------------------


class TestKnowledgeBase:
    def test_all_frameworks_present(self):
        from orbit.assist.knowledge.frameworks import FRAMEWORK_KNOWLEDGE

        assert "lerobot" in FRAMEWORK_KNOWLEDGE
        assert "openvla" in FRAMEWORK_KNOWLEDGE
        assert "openpi" in FRAMEWORK_KNOWLEDGE
        assert "groot" in FRAMEWORK_KNOWLEDGE
        assert "osmo" in FRAMEWORK_KNOWLEDGE

    def test_frameworks_have_required_fields(self):
        from orbit.assist.knowledge.frameworks import FRAMEWORK_KNOWLEDGE

        for key, fw in FRAMEWORK_KNOWLEDGE.items():
            assert "name" in fw, f"{key} missing name"
            assert "command_template" in fw, f"{key} missing command_template"
            assert "format_required" in fw, f"{key} missing format_required"
            assert "gpu_minimum" in fw, f"{key} missing gpu_minimum"
            assert "tips" in fw, f"{key} missing tips"
            assert len(fw["tips"]) > 0, f"{key} has no tips"

    def test_framework_summary_generation(self):
        from orbit.assist.knowledge.frameworks import get_framework_summary

        summary = get_framework_summary()
        assert "LeRobot" in summary
        assert "GR00T" in summary
        assert len(summary) > 100

    def test_all_formats_present(self):
        from orbit.assist.knowledge.formats import FORMAT_KNOWLEDGE

        assert "lerobot-v3" in FORMAT_KNOWLEDGE
        assert "lerobot-v2" in FORMAT_KNOWLEDGE
        assert "rlds" in FORMAT_KNOWLEDGE
        assert "hdf5" in FORMAT_KNOWLEDGE
        assert "rosbag" in FORMAT_KNOWLEDGE

    def test_troubleshooting_entries_complete(self):
        from orbit.assist.knowledge.troubleshoot import TROUBLESHOOTING_DB

        assert len(TROUBLESHOOTING_DB) >= 10
        for entry in TROUBLESHOOTING_DB:
            assert entry.symptom, "Missing symptom"
            assert entry.cause, "Missing cause"
            assert len(entry.fix) > 0, f"No fixes for: {entry.symptom}"
            assert len(entry.keywords) > 0, f"No keywords for: {entry.symptom}"

    def test_troubleshooting_search(self):
        from orbit.assist.knowledge.troubleshoot import search_troubleshooting

        results = search_troubleshooting("loss is NaN after a few steps")
        assert len(results) > 0
        assert "NaN" in results[0].symptom or "nan" in results[0].symptom.lower()

    def test_troubleshooting_search_groot(self):
        from orbit.assist.knowledge.troubleshoot import search_troubleshooting

        results = search_troubleshooting("groot v3 not supported")
        assert len(results) > 0

    def test_cloud_providers_present(self):
        from orbit.assist.knowledge.cloud import CLOUD_KNOWLEDGE

        assert "gcp" in CLOUD_KNOWLEDGE
        assert "aws" in CLOUD_KNOWLEDGE
        assert "runpod" in CLOUD_KNOWLEDGE
        assert "lambda" in CLOUD_KNOWLEDGE
        assert "osmo" in CLOUD_KNOWLEDGE

    def test_cloud_gpu_recommendations(self):
        from orbit.assist.knowledge.cloud import recommend_cloud_gpu

        recs = recommend_cloud_gpu(vram_needed_gb=40)
        assert len(recs) > 0
        # Should include A100-80GB options
        assert any("a100" in r["gpu"] for r in recs)
        # Should be sorted by cost
        costs = [r["cost_value"] for r in recs]
        assert costs == sorted(costs)

    def test_format_conversion_advice(self):
        from orbit.assist.knowledge.formats import get_conversion_advice

        # lerobot-v3 needs conversion for groot
        advice = get_conversion_advice("lerobot-v3", "groot")
        assert advice is not None
        assert "lerobot-v2" in advice

    def test_no_conversion_needed(self):
        from orbit.assist.knowledge.formats import get_conversion_advice

        # lerobot-v2 is compatible with groot
        advice = get_conversion_advice("lerobot-v2", "groot")
        assert advice is None


# ---------------------------------------------------------------------------
# Gemini backend tests (mocked API)
# ---------------------------------------------------------------------------


class TestGeminiBackend:
    def _make_backend(self):
        from orbit.assist.backends.gemini import GeminiBackend

        return GeminiBackend(api_key="test-key")

    def test_text_only_response(self):
        """Model returns text without tool calls."""
        backend = self._make_backend()

        mock_response = {
            "candidates": [{
                "content": {
                    "parts": [{"text": "Hello! I can help with that."}],
                },
            }],
        }

        with patch.object(backend, "_call_gemini", return_value=mock_response):
            result = backend.run_turn("hi", "system prompt")
            assert result == "Hello! I can help with that."

    def test_single_tool_call(self):
        """Model calls one tool then returns text."""
        backend = self._make_backend()

        # First call: model wants to use a tool
        tool_response = {
            "candidates": [{
                "content": {
                    "parts": [{
                        "functionCall": {
                            "name": "scan_directory",
                            "args": {"path": "."},
                        }
                    }],
                },
            }],
        }
        # Second call: model returns final text
        text_response = {
            "candidates": [{
                "content": {
                    "parts": [{"text": "I found 2 datasets in your directory."}],
                },
            }],
        }

        with patch.object(backend, "_call_gemini", side_effect=[tool_response, text_response]):
            with patch("orbit.assist.tools.execute_tool", return_value={"datasets": []}):
                result = backend.run_turn("what's in my directory?", "system prompt")
                assert "datasets" in result.lower() or "directory" in result.lower()

    def test_multi_step_tool_calls(self):
        """Model chains multiple tool calls."""
        backend = self._make_backend()

        call1 = {
            "candidates": [{
                "content": {
                    "parts": [{
                        "functionCall": {
                            "name": "orbit_analyze",
                            "args": {"dataset": "./data"},
                        }
                    }],
                },
            }],
        }
        call2 = {
            "candidates": [{
                "content": {
                    "parts": [{
                        "functionCall": {
                            "name": "orbit_suggest",
                            "args": {"dataset": "./data"},
                        }
                    }],
                },
            }],
        }
        final = {
            "candidates": [{
                "content": {
                    "parts": [{"text": "Your dataset scores a B. Here's the training command."}],
                },
            }],
        }

        with patch.object(backend, "_call_gemini", side_effect=[call1, call2, final]):
            with patch("orbit.assist.tools.execute_tool", return_value={"grade": "B"}):
                result = backend.run_turn("analyze and suggest", "system prompt")
                assert "B" in result

    def test_api_failure_returns_error(self):
        """API failure is handled gracefully."""
        backend = self._make_backend()

        with patch.object(backend, "_call_gemini", return_value=None):
            result = backend.run_turn("hello", "system prompt")
            assert "failed" in result.lower() or "error" in result.lower()

    def test_max_iterations_safety(self):
        """Doesn't loop forever if model keeps calling tools."""
        backend = self._make_backend()

        infinite_tool = {
            "candidates": [{
                "content": {
                    "parts": [{
                        "functionCall": {
                            "name": "scan_directory",
                            "args": {},
                        }
                    }],
                },
            }],
        }

        # Return tool calls forever — should hit max iterations
        with patch.object(backend, "_call_gemini", return_value=infinite_tool):
            with patch("orbit.assist.tools.execute_tool", return_value={}):
                result = backend.run_turn("loop forever", "system prompt")
                assert "maximum" in result.lower()

    def test_history_accumulates(self):
        """History grows across turns."""
        backend = self._make_backend()

        text_response = {
            "candidates": [{
                "content": {
                    "parts": [{"text": "Response"}],
                },
            }],
        }

        with patch.object(backend, "_call_gemini", return_value=text_response):
            backend.run_turn("first message", "sys")
            assert len(backend.history) == 2  # user + model

            backend.run_turn("second message", "sys")
            assert len(backend.history) == 4  # 2 more

    def test_reset_clears_history(self):
        backend = self._make_backend()
        backend.history = [{"role": "user", "parts": [{"text": "test"}]}]
        backend.reset()
        assert backend.history == []

    def test_on_tool_call_callback(self):
        """Tool call callback is invoked."""
        backend = self._make_backend()

        tool_response = {
            "candidates": [{
                "content": {
                    "parts": [{
                        "functionCall": {
                            "name": "scan_directory",
                            "args": {"path": "."},
                        }
                    }],
                },
            }],
        }
        text_response = {
            "candidates": [{
                "content": {
                    "parts": [{"text": "Done."}],
                },
            }],
        }

        callback = MagicMock()
        with patch.object(backend, "_call_gemini", side_effect=[tool_response, text_response]):
            with patch("orbit.assist.tools.execute_tool", return_value={"ok": True}):
                backend.run_turn("scan", "sys", on_tool_call=callback)
                callback.assert_called_once()
                assert callback.call_args[0][0] == "scan_directory"


# ---------------------------------------------------------------------------
# Agent entry point tests
# ---------------------------------------------------------------------------


class TestRunAssist:
    def test_requires_api_key(self):
        """Should exit cleanly when no API key is set."""
        with patch("orbit.llm_utils.load_api_key", return_value=None):
            with patch("orbit.assist.agent._run_repl") as mock_repl:
                from orbit.assist import run_assist

                run_assist()
                mock_repl.assert_not_called()

    def test_function_exists(self):
        from orbit.assist import run_assist

        assert callable(run_assist)

    def test_backward_compatible_import(self):
        """Old import path still works."""
        from orbit.assist import run_assist

        assert callable(run_assist)


# ---------------------------------------------------------------------------
# Claude backend tests
# ---------------------------------------------------------------------------


class TestClaudeBackend:
    def test_not_available_without_key(self):
        from orbit.assist.backends.claude import is_available

        with patch.dict(os.environ, {}, clear=True):
            # Remove ANTHROPIC_API_KEY if set
            env = os.environ.copy()
            env.pop("ANTHROPIC_API_KEY", None)
            with patch.dict(os.environ, env, clear=True):
                assert not is_available()


# ---------------------------------------------------------------------------
# Fix 1: Gemini response shape variants (camelCase vs snake_case)
# ---------------------------------------------------------------------------


class TestGeminiFunctionCallParsing:
    """Verify that _extract_function_calls handles all known Gemini response shapes."""

    def test_camel_case_function_call(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"functionCall": {"name": "orbit_analyze", "args": {"dataset": "./data"}}}]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0] == ("orbit_analyze", {"dataset": "./data"})

    def test_snake_case_function_call(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"function_call": {"name": "orbit_gate", "args": {"dataset": "./d", "policy": "act"}}}]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0][0] == "orbit_gate"
        assert calls[0][1]["policy"] == "act"

    def test_args_as_json_string(self):
        """Gemini sometimes returns args as a JSON string instead of dict."""
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"functionCall": {"name": "scan_directory", "args": '{"path": "."}'}}]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0][1] == {"path": "."}

    def test_args_as_invalid_string(self):
        """Malformed args string should not crash."""
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"functionCall": {"name": "scan_directory", "args": "not json"}}]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0][1] == {"raw_args": "not json"}

    def test_empty_args(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"functionCall": {"name": "scan_directory"}}]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0][1] == {}

    def test_null_args(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"functionCall": {"name": "scan_directory", "args": None}}]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0][1] == {}

    def test_mixed_text_and_function_calls(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [
            {"text": "Let me check that for you."},
            {"functionCall": {"name": "orbit_analyze", "args": {"dataset": "x"}}},
        ]
        calls = _extract_function_calls(parts)
        assert len(calls) == 1
        assert calls[0][0] == "orbit_analyze"

    def test_multiple_function_calls_in_one_response(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [
            {"functionCall": {"name": "orbit_analyze", "args": {"dataset": "a"}}},
            {"functionCall": {"name": "orbit_suggest", "args": {"dataset": "a"}}},
        ]
        calls = _extract_function_calls(parts)
        assert len(calls) == 2

    def test_no_function_calls(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"text": "Just a text response."}]
        calls = _extract_function_calls(parts)
        assert calls == []

    def test_malformed_part_ignored(self):
        from orbit.assist.backends.gemini import _extract_function_calls

        parts = [{"functionCall": "not a dict"}, {"text": "hello"}]
        calls = _extract_function_calls(parts)
        assert calls == []


# ---------------------------------------------------------------------------
# Fix 2: Large payload truncation
# ---------------------------------------------------------------------------


class TestPayloadTruncation:
    def test_small_payload_unchanged(self):
        from orbit.assist.backends.gemini import _truncate_result

        result = {"grade": "B", "score": 78}
        truncated = _truncate_result(result, "orbit_analyze")
        assert truncated == result

    def test_large_analyze_payload_summarized(self):
        from orbit.assist.backends.gemini import _MAX_TOOL_RESULT_CHARS, _truncate_result

        # Build a payload larger than the limit
        result = {
            "grade": "B",
            "score": 78,
            "episode_count": 100,
            "huge_field": "x" * 20000,
            "per_episode_data": [{"ep": i, "metrics": "y" * 100} for i in range(500)],
        }
        truncated = _truncate_result(result, "orbit_analyze")
        serialized = json.dumps(truncated, default=str)
        # Should be significantly smaller than original
        assert len(serialized) < len(json.dumps(result, default=str))
        # Should keep grade and score
        assert truncated.get("grade") == "B"
        assert truncated.get("score") == 78

    def test_large_scan_result_capped(self):
        from orbit.assist.backends.gemini import _summarize_scan_result

        # Test the summarizer directly (it's what _truncate_result calls
        # once the size check triggers)
        result = {
            "datasets": [{"path": f"/data/ds_{i}", "format": "lerobot"} for i in range(100)],
            "training_outputs": [{"path": f"/out/{i}", "type": "tb"} for i in range(50)],
            "config_files": [],
            "environment": {"python_version": "3.12"},
        }
        truncated = _summarize_scan_result(result)
        assert len(truncated["datasets"]) <= 15
        assert truncated["datasets_total"] == 100
        assert len(truncated["training_outputs"]) <= 15
        assert truncated["training_outputs_total"] == 50

    def test_deep_truncate_long_strings(self):
        from orbit.assist.backends.gemini import _deep_truncate

        result = {"key": "x" * 5000, "nested": {"deep": "y" * 5000}}
        truncated = _deep_truncate(result)
        assert len(truncated["key"]) < 5000
        assert len(truncated["nested"]["deep"]) < 5000

    def test_deep_truncate_long_lists(self):
        from orbit.assist.backends.gemini import _deep_truncate

        result = {"items": list(range(100))}
        truncated = _deep_truncate(result)
        assert len(truncated["items"]) <= 11  # 10 items + "...(N more)"


# ---------------------------------------------------------------------------
# Fix 3: Subprocess environment
# ---------------------------------------------------------------------------


class TestSubprocessEnvironment:
    def test_resolve_python_returns_executable(self):
        from orbit.assist.tools import _resolve_python

        python = _resolve_python()
        assert python  # non-empty
        assert "python" in python.lower() or "Python" in python

    def test_build_env_has_no_color(self):
        from orbit.assist.tools import _build_subprocess_env

        env = _build_subprocess_env()
        assert env["NO_COLOR"] == "1"

    def test_build_env_preserves_path(self):
        from orbit.assist.tools import _build_subprocess_env

        env = _build_subprocess_env()
        assert "PATH" in env

    def test_build_env_preserves_virtual_env(self):
        from orbit.assist.tools import _build_subprocess_env

        with patch.dict(os.environ, {"VIRTUAL_ENV": "/some/venv"}):
            env = _build_subprocess_env()
            assert env.get("VIRTUAL_ENV") == "/some/venv"

    def test_json_parse_with_ansi_prefix(self):
        """Verify _try_parse_json strips ANSI/log noise before JSON."""
        from orbit.assist.tools import _try_parse_json

        # Simulates a command that prints a warning before JSON
        raw = 'WARNING: something\n\n{"grade": "B", "score": 78}'
        parsed = _try_parse_json(raw)
        assert parsed is not None
        assert parsed["grade"] == "B"

    def test_json_parse_with_ansi_codes(self):
        from orbit.assist.tools import _try_parse_json

        raw = '\x1b[1m\x1b[31m{"error": "bad"}\x1b[0m'
        parsed = _try_parse_json(raw)
        assert parsed is not None
        assert parsed["error"] == "bad"

    def test_json_parse_clean(self):
        from orbit.assist.tools import _try_parse_json

        raw = '{"passed": true, "grade": "A"}'
        parsed = _try_parse_json(raw)
        assert parsed == {"passed": True, "grade": "A"}

    def test_json_parse_returns_none_for_garbage(self):
        from orbit.assist.tools import _try_parse_json

        assert _try_parse_json("not json at all") is None
        assert _try_parse_json("") is None
        assert _try_parse_json(None) is None


# ---------------------------------------------------------------------------
# Fix 4: CI/CD exit code determination
# ---------------------------------------------------------------------------


class TestExitCodeDetermination:
    def test_success_returns_0(self):
        from orbit.assist.agent import _determine_exit_code

        result = _determine_exit_code([], "Everything looks good.")
        assert result == 0

    def test_tool_error_returns_1(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{"name": "orbit_analyze", "result": {"error": "Dataset not found"}}]
        assert _determine_exit_code(results, "I encountered an error.") == 1

    def test_gate_fail_returns_1(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{"name": "orbit_gate", "result": {"passed": False, "grade": "D"}}]
        assert _determine_exit_code(results, "The dataset did not pass.") == 1

    def test_gate_pass_returns_0(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{"name": "orbit_gate", "result": {"passed": True, "grade": "B"}}]
        assert _determine_exit_code(results, "Dataset passes quality gate.") == 0

    def test_debug_critical_finding_returns_1(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{
            "name": "orbit_debug",
            "result": {
                "status": "nan_loss",
                "findings": [{"severity": "critical", "finding": "NaN loss"}],
            },
        }]
        assert _determine_exit_code(results, "Training failed with NaN.") == 1

    def test_debug_info_finding_returns_0(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{
            "name": "orbit_debug",
            "result": {
                "status": "converged_well",
                "findings": [{"severity": "info", "finding": "Converged"}],
            },
        }]
        assert _determine_exit_code(results, "Training converged normally.") == 0

    def test_nonzero_exit_code_returns_1(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{"name": "orbit_analyze", "result": {"exit_code": 1, "stdout": "error"}}]
        assert _determine_exit_code(results, "Something went wrong.") == 1

    def test_prose_failure_phrases(self):
        from orbit.assist.agent import _determine_exit_code

        assert _determine_exit_code([], "Failed to load the dataset.") == 1
        assert _determine_exit_code([], "Could not connect to the API.") == 1
        assert _determine_exit_code([], "Unable to find any data.") == 1

    def test_no_false_positive_on_error_word(self):
        """The word 'error' in an explanation shouldn't trigger exit 1."""
        from orbit.assist.agent import _determine_exit_code

        # "error" appearing in context, not as a failure signal
        assert _determine_exit_code([], "The error handling looks correct.") == 0
        assert _determine_exit_code([], "There was no error in the training.") == 0

    def test_gate_exit_code_1_means_fail(self):
        from orbit.assist.agent import _determine_exit_code

        results = [{"name": "orbit_gate", "result": {"exit_code": 1, "passed": False}}]
        assert _determine_exit_code(results, "Gate failed.") == 1

    def test_mixed_tools_one_failure(self):
        from orbit.assist.agent import _determine_exit_code

        results = [
            {"name": "orbit_analyze", "result": {"grade": "B", "score": 78}},
            {"name": "orbit_gate", "result": {"passed": False, "grade": "D"}},
        ]
        assert _determine_exit_code(results, "Analysis complete but gate failed.") == 1

    def test_mixed_tools_all_success(self):
        from orbit.assist.agent import _determine_exit_code

        results = [
            {"name": "orbit_analyze", "result": {"grade": "B", "score": 78}},
            {"name": "orbit_gate", "result": {"passed": True, "grade": "B"}},
            {"name": "orbit_suggest", "result": {"train_command": "python train.py"}},
        ]
        assert _determine_exit_code(results, "All done, here's the command.") == 0
