#!/usr/bin/env python3
"""Full-coverage agentic runtime test suite for wee_runtime.py.

Fills gaps not covered by test_wee_runtime_agentic.py and
test_wee_runtime_comprehensive.py:

  - 3-round message chain depth verification (exact sequence)
  - Content + tool call in the same streaming round (mixed delta)
  - Tool arguments split across multiple delta chunks (incremental)
  - Empty choices in a chunk (graceful skip)
  - execute_tool edge cases: empty output, unknown tool, stderr-only,
    subprocess exception, bash/python with special characters
  - Prompt constant validation (anti-hallucination, tool capability)
  - OpenAI client config validation (timeout, max_retries)
  - Temperature=0 passes 0.0 (not omitted), Temperature=None omits key
  - System prompt: without --tools anti-hallucination only; with --tools both
  - Full regression: all 7 core API endpoints after code changes
  - Live Ollama: math-reasoning prompt → Python tool invoked
  - Live Ollama: SSH-aware prompt → StrictHostKeyChecking injected
  - Live OpenRouter: free model simple echo
  - Live OpenRouter: free model tool calling (if supported)
  - Rate limit simulation: 429 response doesn't crash the process
  - Per-round API call count assertions for N-round loops
  - Null-content chunks don't corrupt content accumulation
  - Tool counter increments correctly across rounds
  - Unknown tool falls through gracefully with error message to model

Run all:
    pytest tests/test_wee_agentic_full_coverage.py -v

Run unit only (fast):
    pytest tests/test_wee_agentic_full_coverage.py -v -m "not live"

Run live Ollama:
    pytest tests/test_wee_agentic_full_coverage.py -v -k ollama

Run live OpenRouter:
    pytest tests/test_wee_agentic_full_coverage.py -v -k openrouter
"""

import io
import json
import os
import subprocess
import sys
import time
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock, call, patch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/..")

import wee_runtime  # noqa: E402

# ---------------------------------------------------------------------------
# Test environment configuration
# ---------------------------------------------------------------------------

OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://192.168.1.101:11434")
OLLAMA_MODEL = os.environ.get("OLLAMA_TEST_MODEL", "ollama/qwen3:8b")
OPENROUTER_MODEL = os.environ.get(
    "OPENROUTER_TEST_MODEL", "openrouter/google/gemma-3-12b-it:free"
)
LIVE_TIMEOUT = int(os.environ.get("WEE_LIVE_TIMEOUT", "120"))

WEE_RUNTIME = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "wee_runtime.py"
)


def _has_ollama() -> bool:
    try:
        r = subprocess.run(
            [
                "curl",
                "-s",
                "-o",
                "/dev/null",
                "-w",
                "%{http_code}",
                f"{OLLAMA_HOST}/api/tags",
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )
        return r.stdout.strip() == "200"
    except Exception:
        return False


def _has_openrouter_key() -> bool:
    if os.environ.get("OPENROUTER_API_KEY"):
        return True
    try:
        import keyring

        return bool(keyring.get_password("wee-orchestrator", "OPENROUTER_API_KEY"))
    except Exception:
        return False


HAS_OLLAMA = _has_ollama()
HAS_OPENROUTER = _has_openrouter_key()

skip_ollama = unittest.skipUnless(HAS_OLLAMA, "Ollama not reachable")
skip_openrouter = unittest.skipUnless(HAS_OPENROUTER, "No OPENROUTER_API_KEY")


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _make_chunk(content=None, tool_calls=None, finish_reason=None, choices=True):
    """Build a fake SSE chunk."""
    if not choices:
        return SimpleNamespace(choices=[])
    delta = SimpleNamespace(content=content, tool_calls=tool_calls)
    choice = SimpleNamespace(delta=delta, finish_reason=finish_reason)
    return SimpleNamespace(choices=[choice])


def _make_tc_delta(idx, tc_id=None, name=None, arguments=None):
    """Build a tool-call delta fragment."""
    fn = SimpleNamespace(name=name, arguments=arguments)
    return SimpleNamespace(index=idx, id=tc_id, function=fn)


def _run_main(
    model, prompt, tools=False, system_prompt="", temperature=None, extra_args=None
):
    """Run wee_runtime.main() in-process with captured I/O."""
    argv = ["wee_runtime.py", "--model", model]
    if tools:
        argv.append("--tools")
    if system_prompt:
        argv.extend(["--system-prompt", system_prompt])
    if temperature is not None:
        argv.extend(["--temperature", str(temperature)])
    if extra_args:
        argv.extend(extra_args)
    argv.append(prompt)

    old_argv, old_stdout, old_stderr = sys.argv, sys.stdout, sys.stderr
    sys.argv = argv
    sys.stdout = io.StringIO()
    sys.stderr = io.StringIO()
    exit_code = 0
    try:
        wee_runtime.main()
    except SystemExit as exc:
        exit_code = exc.code if exc.code is not None else 0
    finally:
        stdout_val = sys.stdout.getvalue()
        stderr_val = sys.stderr.getvalue()
        sys.argv, sys.stdout, sys.stderr = old_argv, old_stdout, old_stderr
    return SimpleNamespace(stdout=stdout_val, stderr=stderr_val, returncode=exit_code)


def _run_cli(model, prompt, tools=False, extra_args=None, timeout=LIVE_TIMEOUT):
    """Run wee_runtime.py as a subprocess."""
    cmd = [sys.executable, WEE_RUNTIME, "--model", model, "--timeout", str(timeout)]
    if tools:
        cmd.append("--tools")
    if extra_args:
        cmd.extend(extra_args)
    cmd.append(prompt)
    return subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        timeout=timeout + 30,
        env={**os.environ, "PYTHONUNBUFFERED": "1"},
    )


# ===========================================================================
# 1. PROMPT CONSTANTS VALIDATION
# ===========================================================================


class TestPromptConstants(unittest.TestCase):
    """Anti-hallucination and tool capability prompt constants are well-formed."""

    def test_anti_hallucination_prompt_nonempty(self):
        self.assertTrue(len(wee_runtime._ANTI_HALLUCINATION_PROMPT) > 50)

    def test_anti_hallucination_prompt_has_fabricate_warning(self):
        self.assertIn("NEVER fabricate", wee_runtime._ANTI_HALLUCINATION_PROMPT)

    def test_anti_hallucination_prompt_has_stricthostkeychecking(self):
        self.assertIn("StrictHostKeyChecking", wee_runtime._ANTI_HALLUCINATION_PROMPT)

    def test_tool_capability_prompt_nonempty(self):
        self.assertTrue(len(wee_runtime._WEE_TOOL_CAPABILITY_PROMPT) > 50)

    def test_tool_capability_prompt_mentions_bash(self):
        self.assertIn("bash", wee_runtime._WEE_TOOL_CAPABILITY_PROMPT)

    def test_tool_capability_prompt_mentions_python(self):
        self.assertIn("python", wee_runtime._WEE_TOOL_CAPABILITY_PROMPT.lower())

    def test_tool_capability_prompt_has_not_sandboxed(self):
        self.assertIn("NOT sandboxed", wee_runtime._WEE_TOOL_CAPABILITY_PROMPT)

    def test_anti_hallucination_is_string(self):
        self.assertIsInstance(wee_runtime._ANTI_HALLUCINATION_PROMPT, str)

    def test_tool_capability_is_string(self):
        self.assertIsInstance(wee_runtime._WEE_TOOL_CAPABILITY_PROMPT, str)


# ===========================================================================
# 2. SYSTEM PROMPT CONSTRUCTION (mocked client)
# ===========================================================================


class TestSystemPromptConstruction(unittest.TestCase):
    """Verify the effective system prompt injected into the API call."""

    @patch("openai.OpenAI")
    def test_no_tools_only_anti_hallucination(self, mock_cls):
        """Without --tools, system prompt has anti-hallucination but not tool capability."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "hello")

        msgs = mock_client.chat.completions.create.call_args[1]["messages"]
        sys_msgs = [m for m in msgs if m["role"] == "system"]
        self.assertTrue(len(sys_msgs) >= 1, "Expected a system message")
        combined = " ".join(m["content"] for m in sys_msgs)
        self.assertIn("NEVER fabricate", combined)
        self.assertNotIn("NOT sandboxed", combined)

    @patch("openai.OpenAI")
    def test_with_tools_has_both_prompts(self, mock_cls):
        """With --tools, system prompt includes both anti-hallucination and tool capability."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "hello", tools=True)

        msgs = mock_client.chat.completions.create.call_args[1]["messages"]
        sys_msgs = [m for m in msgs if m["role"] == "system"]
        combined = " ".join(m["content"] for m in sys_msgs)
        self.assertIn("NEVER fabricate", combined)
        self.assertIn("NOT sandboxed", combined)

    @patch("openai.OpenAI")
    def test_user_system_prompt_is_preserved(self, mock_cls):
        """Custom --system-prompt text is included in the system message."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "hello", system_prompt="MY_CUSTOM_SYSTEM_TEXT")

        msgs = mock_client.chat.completions.create.call_args[1]["messages"]
        sys_msgs = [m for m in msgs if m["role"] == "system"]
        combined = " ".join(m["content"] for m in sys_msgs)
        self.assertIn("MY_CUSTOM_SYSTEM_TEXT", combined)


# ===========================================================================
# 3. OPENAI CLIENT CONFIGURATION
# ===========================================================================


class TestOpenAIClientConfig(unittest.TestCase):
    """Verify httpx timeout and max_retries are set correctly."""

    @patch("httpx.Timeout")
    @patch("openai.OpenAI")
    def test_max_retries_zero(self, mock_cls, mock_timeout):
        """OpenAI client is created with max_retries=0."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )
        mock_timeout.return_value = MagicMock()

        _run_main("ollama/test", "client config")

        init_kwargs = mock_cls.call_args[1]
        self.assertEqual(init_kwargs.get("max_retries"), 0)

    @patch("httpx.Timeout")
    @patch("openai.OpenAI")
    def test_connect_timeout_15s(self, mock_cls, mock_timeout):
        """httpx.Timeout is called with connect=15.0."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )
        mock_timeout.return_value = MagicMock()

        _run_main("ollama/test", "timeout config")

        mock_timeout.assert_called_once()
        _, timeout_kwargs = mock_timeout.call_args
        self.assertEqual(timeout_kwargs.get("connect"), 15.0)

    @patch("openai.OpenAI")
    def test_api_base_url_set(self, mock_cls):
        """OpenAI client receives resolved base_url."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/my-model", "base url check")

        init_kwargs = mock_cls.call_args[1]
        self.assertIn("base_url", init_kwargs)
        self.assertIn("11434", str(init_kwargs["base_url"]))


# ===========================================================================
# 4. TEMPERATURE FLAG EDGE CASES
# ===========================================================================


class TestTemperatureEdgeCases(unittest.TestCase):
    """Temperature=0, None, and non-integer values."""

    @patch("openai.OpenAI")
    def test_temperature_zero_passed(self, mock_cls):
        """Temperature=0.0 is included in create kwargs (not omitted as falsy)."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "temp zero", temperature=0.0)

        kwargs = mock_client.chat.completions.create.call_args[1]
        self.assertIn("temperature", kwargs)
        self.assertEqual(kwargs["temperature"], 0.0)

    @patch("openai.OpenAI")
    def test_temperature_none_omitted(self, mock_cls):
        """Without --temperature, 'temperature' is NOT in create kwargs."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "no temp")

        kwargs = mock_client.chat.completions.create.call_args[1]
        self.assertNotIn("temperature", kwargs)

    @patch("openai.OpenAI")
    def test_temperature_point_seven(self, mock_cls):
        """Temperature=0.7 is passed as a float."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "temp 0.7", temperature=0.7)

        kwargs = mock_client.chat.completions.create.call_args[1]
        self.assertAlmostEqual(kwargs.get("temperature"), 0.7)


# ===========================================================================
# 5. EXECUTE_TOOL EDGE CASES
# ===========================================================================


class TestExecuteToolEdgeCases(unittest.TestCase):
    """Edge cases in execute_tool not covered by other suites."""

    def test_bash_empty_output_returns_no_output_placeholder(self):
        """Bash command producing no stdout returns '(no output)'."""
        result = wee_runtime.execute_tool("bash", {"command": "true"})
        self.assertEqual(result, "(no output)")

    def test_python_empty_output_returns_no_output_placeholder(self):
        """Python code with no print returns '(no output)'."""
        result = wee_runtime.execute_tool("python", {"code": "x = 1 + 1"})
        self.assertEqual(result, "(no output)")

    def test_unknown_tool_name_returns_error_string(self):
        """Unknown tool name returns informative error without raising."""
        result = wee_runtime.execute_tool("nonexistent_tool", {"arg": "val"})
        self.assertIn("Unknown tool", result)
        self.assertIn("nonexistent_tool", result)

    def test_bash_stderr_only_included_on_nonzero_exit(self):
        """Stderr output is included in result only when exit code != 0."""
        result = wee_runtime.execute_tool(
            "bash", {"command": "echo error_msg >&2; exit 1"}
        )
        self.assertIn("error_msg", result)
        self.assertIn("STDERR", result)

    def test_bash_nonzero_exit_with_stdout_and_stderr(self):
        """Both stdout and stderr are present when command exits nonzero."""
        result = wee_runtime.execute_tool(
            "bash",
            {"command": "echo stdout_text; echo stderr_text >&2; exit 2"},
        )
        self.assertIn("stdout_text", result)
        self.assertIn("stderr_text", result)

    def test_python_exception_output_included(self):
        """Python code that raises includes the exception in result."""
        result = wee_runtime.execute_tool(
            "python", {"code": "raise RuntimeError('test_exception_msg')"}
        )
        self.assertIn("test_exception_msg", result)
        self.assertIn("STDERR", result)

    def test_bash_multiline_output_preserved(self):
        """Multi-line bash output is preserved in result."""
        result = wee_runtime.execute_tool(
            "bash", {"command": "printf 'line1\\nline2\\nline3'"}
        )
        self.assertIn("line1", result)
        self.assertIn("line2", result)
        self.assertIn("line3", result)

    def test_bash_with_pipe(self):
        """Bash command with pipes executes correctly."""
        result = wee_runtime.execute_tool(
            "bash", {"command": "echo 'hello_pipe' | tr 'a-z' 'A-Z'"}
        )
        self.assertIn("HELLO_PIPE", result)

    def test_python_multiline_code(self):
        """Multi-line Python code executes correctly."""
        code = "total = sum(range(1, 6))\nprint(f'sum={total}')"
        result = wee_runtime.execute_tool("python", {"code": code})
        self.assertIn("sum=15", result)

    def test_bash_no_command_key_returns_error(self):
        """Bash tool with missing 'command' key returns an error."""
        result = wee_runtime.execute_tool("bash", {})
        self.assertIn("Error", result)

    def test_python_no_code_key_returns_error(self):
        """Python tool with missing 'code' key returns an error."""
        result = wee_runtime.execute_tool("python", {})
        self.assertIn("Error", result)

    def test_restricted_permission_blocks_bash(self):
        """Permission='restricted' blocks bash execution."""
        result = wee_runtime.execute_tool(
            "bash", {"command": "echo blocked"}, permission="restricted"
        )
        self.assertIn("blocked", result.lower())
        self.assertNotIn("echo blocked", result)

    def test_restricted_permission_blocks_python(self):
        """Permission='restricted' blocks python execution."""
        result = wee_runtime.execute_tool(
            "python", {"code": "print('should_not_run')"}, permission="restricted"
        )
        self.assertNotIn("should_not_run", result)

    def test_elevated_permission_same_as_auto(self):
        """Permission='elevated' executes tools (treated identically to 'auto')."""
        result = wee_runtime.execute_tool(
            "bash", {"command": "echo elevated_runs"}, permission="elevated"
        )
        self.assertIn("elevated_runs", result)


# ===========================================================================
# 6. SSH SANITIZATION — SANITIZE_BASH_COMMAND
# ===========================================================================


class TestSanitizeBashCommandFull(unittest.TestCase):
    """Extended sanitization tests."""

    def test_ssh_already_has_flag_unchanged(self):
        """Command with StrictHostKeyChecking already present is not modified."""
        cmd = "ssh -o StrictHostKeyChecking=accept-new user@host"
        self.assertEqual(wee_runtime.sanitize_bash_command(cmd), cmd)

    def test_scp_gets_flag_injected(self):
        """scp command gets StrictHostKeyChecking injected."""
        result = wee_runtime.sanitize_bash_command("scp file user@host:/tmp/")
        self.assertIn("StrictHostKeyChecking", result)

    def test_sftp_gets_flag_injected(self):
        """sftp command gets StrictHostKeyChecking injected."""
        result = wee_runtime.sanitize_bash_command("sftp user@host")
        self.assertIn("StrictHostKeyChecking", result)

    def test_non_ssh_command_unchanged(self):
        """Commands without ssh/scp/sftp are unchanged."""
        cmd = "ls -la /tmp"
        self.assertEqual(wee_runtime.sanitize_bash_command(cmd), cmd)

    def test_empty_command_unchanged(self):
        """Empty string returns unchanged."""
        self.assertEqual(wee_runtime.sanitize_bash_command(""), "")

    def test_word_boundary_no_match_sshd(self):
        """'sshd' should not match the ssh pattern (word boundary)."""
        cmd = "systemctl status sshd"
        # StrictHostKeyChecking should NOT be injected for 'sshd'
        result = wee_runtime.sanitize_bash_command(cmd)
        self.assertNotIn("StrictHostKeyChecking", result)

    def test_command_with_pipe_ssh_sanitized(self):
        """SSH in a pipe still gets sanitized."""
        cmd = "cat file | ssh user@host 'cat > /tmp/out'"
        result = wee_runtime.sanitize_bash_command(cmd)
        self.assertIn("StrictHostKeyChecking", result)

    def test_multiple_ssh_commands_all_sanitized(self):
        """Multiple ssh invocations in one command all get sanitized."""
        cmd = "ssh host1 echo hello; ssh host2 echo world"
        result = wee_runtime.sanitize_bash_command(cmd)
        count = result.count("StrictHostKeyChecking")
        self.assertGreaterEqual(count, 2)


# ===========================================================================
# 7. MOCKED LOOP — FULL MESSAGE CHAIN VERIFICATION
# ===========================================================================


class TestFullMessageChain(unittest.TestCase):
    """Verify the exact message sequence after multi-round tool calling."""

    @patch("openai.OpenAI")
    def test_three_round_message_sequence(self, mock_cls):
        """Three-round loop builds: sys→user→asst(tool)→tool→asst(tool)→tool→asst."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        # Round 1: bash call
        td1 = _make_tc_delta(
            0, tc_id="tc_r1", name="bash", arguments='{"command": "echo round1"}'
        )
        r1 = [
            _make_chunk(content="", tool_calls=[td1]),
            _make_chunk(finish_reason="tool_calls"),
        ]

        # Round 2: python call
        td2 = _make_tc_delta(
            0, tc_id="tc_r2", name="python", arguments='{"code": "print(2)"}'
        )
        r2 = [
            _make_chunk(content="", tool_calls=[td2]),
            _make_chunk(finish_reason="tool_calls"),
        ]

        # Round 3: final answer
        r3 = [_make_chunk(content="Final answer"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2), iter(r3)]

        _run_main("ollama/test", "multi-round", tools=True)

        # Should have been called 3 times
        self.assertEqual(mock_client.chat.completions.create.call_count, 3)

        # Round 3 call receives all prior messages
        r3_msgs = mock_client.chat.completions.create.call_args_list[2][1]["messages"]

        roles = [m["role"] for m in r3_msgs]

        # system + user + assistant(round1) + tool(round1) + assistant(round2) + tool(round2)
        self.assertIn("system", roles)
        self.assertIn("user", roles)
        self.assertEqual(roles.count("assistant"), 2)
        self.assertEqual(roles.count("tool"), 2)

    @patch("openai.OpenAI")
    def test_tool_messages_contain_real_output(self, mock_cls):
        """Tool messages in round N+1 contain the actual execute_tool output."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        # Use echo with a distinctive string
        td = _make_tc_delta(
            0,
            tc_id="tc_echo",
            name="bash",
            arguments='{"command": "echo UNIQUE_OUTPUT_12345"}',
        )
        r1 = [
            _make_chunk(content="", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="Done"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        _run_main("ollama/test", "echo test", tools=True)

        r2_msgs = mock_client.chat.completions.create.call_args_list[1][1]["messages"]
        tool_msgs = [m for m in r2_msgs if m["role"] == "tool"]
        self.assertEqual(len(tool_msgs), 1)
        self.assertIn("UNIQUE_OUTPUT_12345", tool_msgs[0]["content"])

    @patch("openai.OpenAI")
    def test_assistant_message_includes_tool_call_list(self, mock_cls):
        """After a tool round, assistant message has 'tool_calls' list in messages."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        td = _make_tc_delta(
            0,
            tc_id="tc_asst_check",
            name="bash",
            arguments='{"command": "echo asst_check"}',
        )
        r1 = [
            _make_chunk(content="thinking", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="Done"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        _run_main("ollama/test", "check asst msg", tools=True)

        r2_msgs = mock_client.chat.completions.create.call_args_list[1][1]["messages"]
        asst_msgs = [m for m in r2_msgs if m["role"] == "assistant"]
        self.assertTrue(len(asst_msgs) >= 1)
        asst_with_tools = [m for m in asst_msgs if m.get("tool_calls")]
        self.assertEqual(len(asst_with_tools), 1)
        tc_list = asst_with_tools[0]["tool_calls"]
        self.assertIsInstance(tc_list, list)
        self.assertEqual(len(tc_list), 1)
        self.assertEqual(tc_list[0]["id"], "tc_asst_check")


# ===========================================================================
# 8. STREAMING EDGE CASES
# ===========================================================================


class TestStreamingEdgeCasesFull(unittest.TestCase):
    """Edge cases in streaming: empty choices, null content, mixed deltas."""

    @patch("openai.OpenAI")
    def test_empty_choices_chunk_skipped_gracefully(self, mock_cls):
        """Chunk with empty choices list does not crash."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        chunks = [
            _make_chunk(choices=False),  # empty choices
            _make_chunk(content="Hello "),
            _make_chunk(content="world"),
            _make_chunk(finish_reason="stop"),
        ]
        mock_client.chat.completions.create.return_value = iter(chunks)

        result = _run_main("ollama/test", "skip empty choices")
        self.assertEqual(result.returncode, 0)
        self.assertIn("Hello world", result.stdout)

    @patch("openai.OpenAI")
    def test_null_content_chunks_ignored_in_output(self, mock_cls):
        """Chunks with content=None are not appended to output."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        chunks = [
            _make_chunk(content=None),
            _make_chunk(content="Good"),
            _make_chunk(content=None),
            _make_chunk(content=" output"),
            _make_chunk(finish_reason="stop"),
        ]
        mock_client.chat.completions.create.return_value = iter(chunks)

        result = _run_main("ollama/test", "null content")
        self.assertEqual(result.returncode, 0)
        self.assertIn("Good output", result.stdout)
        # Should not have "None" literally in output
        self.assertNotIn("None", result.stdout)

    @patch("openai.OpenAI")
    def test_content_and_tool_call_in_same_chunk(self, mock_cls):
        """When content + tool_call arrive in same delta, both are handled."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        # Single chunk has both content and a tool call
        td = _make_tc_delta(
            0, tc_id="tc_mixed", name="bash", arguments='{"command": "echo mixed"}'
        )
        r1 = [
            _make_chunk(content="Let me check that:", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="Done"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        result = _run_main("ollama/test", "mixed delta", tools=True)
        self.assertEqual(result.returncode, 0)

    @patch("openai.OpenAI")
    def test_tool_arguments_split_across_deltas(self, mock_cls):
        """Tool arguments sent in multiple incremental chunks are concatenated."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        # Arguments delivered in 3 pieces
        td_part1 = _make_tc_delta(
            0, tc_id="tc_split", name="bash", arguments='{"command": '
        )
        td_part2 = _make_tc_delta(0, tc_id=None, name=None, arguments='"echo ')
        td_part3 = _make_tc_delta(0, tc_id=None, name=None, arguments='split_args"}')

        r1 = [
            _make_chunk(content="", tool_calls=[td_part1]),
            _make_chunk(content="", tool_calls=[td_part2]),
            _make_chunk(content="", tool_calls=[td_part3]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="Done"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        result = _run_main("ollama/test", "split args test", tools=True)
        self.assertEqual(result.returncode, 0)

        # Tool message in round 2 should include 'split_args'
        r2_msgs = mock_client.chat.completions.create.call_args_list[1][1]["messages"]
        tool_msgs = [m for m in r2_msgs if m["role"] == "tool"]
        self.assertEqual(len(tool_msgs), 1)
        self.assertIn("split_args", tool_msgs[0]["content"])

    @patch("openai.OpenAI")
    def test_streaming_output_flushed_incrementally(self, mock_cls):
        """Streaming produces output with correct final newline."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        chunks = [
            _make_chunk(content="Tok1 "),
            _make_chunk(content="Tok2 "),
            _make_chunk(content="Tok3"),
            _make_chunk(finish_reason="stop"),
        ]
        mock_client.chat.completions.create.return_value = iter(chunks)

        result = _run_main("ollama/test", "streaming flush")
        # Final output should include all tokens and end with newline
        self.assertIn("Tok1 Tok2 Tok3", result.stdout)
        self.assertTrue(result.stdout.endswith("\n"))


# ===========================================================================
# 9. MOCKED LOOP — API CALL COUNT ASSERTIONS
# ===========================================================================


class TestAPICallCounts(unittest.TestCase):
    """Assert exact number of API calls for N-round agentic loops."""

    @patch("openai.OpenAI")
    def test_no_tool_calls_one_api_call(self, mock_cls):
        """Simple text response: exactly 1 API call."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.return_value = iter(
            [
                _make_chunk(content="ok"),
                _make_chunk(finish_reason="stop"),
            ]
        )

        _run_main("ollama/test", "simple")
        self.assertEqual(mock_client.chat.completions.create.call_count, 1)

    @patch("openai.OpenAI")
    def test_one_tool_call_two_api_calls(self, mock_cls):
        """Single tool round + final text = exactly 2 API calls."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        td = _make_tc_delta(
            0, tc_id="tc_1", name="bash", arguments='{"command": "echo hi"}'
        )
        r1 = [
            _make_chunk(content="", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="Done"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        _run_main("ollama/test", "one tool", tools=True)
        self.assertEqual(mock_client.chat.completions.create.call_count, 2)

    @patch("openai.OpenAI")
    def test_five_tool_rounds_six_api_calls(self, mock_cls):
        """Five tool rounds + final text = 6 API calls."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        def make_tool_round(n):
            td = _make_tc_delta(
                0,
                tc_id=f"tc_{n}",
                name="bash",
                arguments=f'{{"command": "echo round{n}"}}',
            )
            return iter(
                [
                    _make_chunk(content="", tool_calls=[td]),
                    _make_chunk(finish_reason="tool_calls"),
                ]
            )

        rounds = [make_tool_round(i) for i in range(5)]
        final = iter(
            [_make_chunk(content="All done"), _make_chunk(finish_reason="stop")]
        )
        mock_client.chat.completions.create.side_effect = rounds + [final]

        _run_main("ollama/test", "five rounds", tools=True)
        self.assertEqual(mock_client.chat.completions.create.call_count, 6)

    @patch("openai.OpenAI")
    def test_tools_not_in_final_round_call(self, mock_cls):
        """The MAX_TOOL_ROUNDS+1 call does not include 'tools' kwarg."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        max_rounds = wee_runtime.MAX_TOOL_ROUNDS

        def make_tool_round(n):
            td = _make_tc_delta(
                0, tc_id=f"tc_{n}", name="bash", arguments=f'{{"command": "echo {n}"}}'
            )
            return iter(
                [
                    _make_chunk(content="", tool_calls=[td]),
                    _make_chunk(finish_reason="tool_calls"),
                ]
            )

        # Exactly MAX_TOOL_ROUNDS tool calls, then a final text response
        rounds = [make_tool_round(i) for i in range(max_rounds)]
        final = iter([_make_chunk(content="Final"), _make_chunk(finish_reason="stop")])
        mock_client.chat.completions.create.side_effect = rounds + [final]

        _run_main("ollama/test", "max rounds exact", tools=True)

        # Final call at index max_rounds should have no 'tools' key
        final_kwargs = mock_client.chat.completions.create.call_args_list[max_rounds][1]
        self.assertNotIn(
            "tools",
            final_kwargs,
            f"Round {max_rounds+1} (final) should not include 'tools'",
        )


# ===========================================================================
# 10. UNKNOWN TOOL HANDLING IN LOOP
# ===========================================================================


class TestUnknownToolInLoop(unittest.TestCase):
    """Model returns an unknown tool name — loop should not crash."""

    @patch("openai.OpenAI")
    def test_unknown_tool_error_sent_to_model(self, mock_cls):
        """Unknown tool produces error message sent back to model as tool result."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        td = _make_tc_delta(
            0, tc_id="tc_unk", name="unknown_tool_xyz", arguments='{"arg": "val"}'
        )
        r1 = [
            _make_chunk(content="", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [
            _make_chunk(content="I see the error."),
            _make_chunk(finish_reason="stop"),
        ]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        result = _run_main("ollama/test", "unknown tool", tools=True)
        self.assertEqual(result.returncode, 0)

        r2_msgs = mock_client.chat.completions.create.call_args_list[1][1]["messages"]
        tool_msgs = [m for m in r2_msgs if m["role"] == "tool"]
        self.assertEqual(len(tool_msgs), 1)
        self.assertIn("Unknown tool", tool_msgs[0]["content"])


# ===========================================================================
# 11. MALFORMED JSON TOOL ARGUMENTS
# ===========================================================================


class TestMalformedToolArguments(unittest.TestCase):
    """Model streams broken JSON for tool arguments."""

    @patch("openai.OpenAI")
    def test_invalid_json_args_does_not_crash(self, mock_cls):
        """Invalid JSON in tool args falls back gracefully."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        td = _make_tc_delta(
            0, tc_id="tc_bad_json", name="bash", arguments="{not valid json{{{{"
        )
        r1 = [
            _make_chunk(content="", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="Handled"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        result = _run_main("ollama/test", "bad json", tools=True)
        self.assertEqual(result.returncode, 0)

    @patch("openai.OpenAI")
    def test_empty_json_args(self, mock_cls):
        """Empty JSON object '{}' for bash command returns error from execute_tool."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        td = _make_tc_delta(0, tc_id="tc_empty_args", name="bash", arguments="{}")
        r1 = [
            _make_chunk(content="", tool_calls=[td]),
            _make_chunk(finish_reason="tool_calls"),
        ]
        r2 = [_make_chunk(content="I see"), _make_chunk(finish_reason="stop")]

        mock_client.chat.completions.create.side_effect = [iter(r1), iter(r2)]

        result = _run_main("ollama/test", "empty args", tools=True)
        self.assertEqual(result.returncode, 0)

        r2_msgs = mock_client.chat.completions.create.call_args_list[1][1]["messages"]
        tool_msgs = [m for m in r2_msgs if m["role"] == "tool"]
        self.assertTrue(len(tool_msgs) >= 1)
        self.assertIn("Error", tool_msgs[0]["content"])


# ===========================================================================
# 12. MAX TOOL ROUNDS — FALLBACK CONTENT
# ===========================================================================


class TestMaxRoundsFallback(unittest.TestCase):
    """When MAX_TOOL_ROUNDS is exhausted, fallback text appears in output."""

    @patch("openai.OpenAI")
    def test_fallback_content_not_empty(self, mock_cls):
        """After max rounds exhaustion, stdout is non-empty."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        def make_loop():
            td = _make_tc_delta(
                0,
                tc_id=f"tc_{time.time_ns()}",
                name="bash",
                arguments='{"command": "echo loop"}',
            )
            return iter(
                [
                    _make_chunk(content="", tool_calls=[td]),
                    _make_chunk(finish_reason="tool_calls"),
                ]
            )

        side_effects = [make_loop() for _ in range(wee_runtime.MAX_TOOL_ROUNDS + 2)]
        mock_client.chat.completions.create.side_effect = side_effects

        result = _run_main("ollama/test", "infinite loop", tools=True)
        self.assertEqual(result.returncode, 0)
        self.assertTrue(
            result.stdout.strip(), "Stdout must not be empty after max rounds"
        )

    @patch("openai.OpenAI")
    def test_fallback_mentions_tool_execution(self, mock_cls):
        """Fallback text indicates tool execution completed or max rounds hit."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client

        def make_loop():
            td = _make_tc_delta(
                0,
                tc_id=f"tc_{time.time_ns()}",
                name="bash",
                arguments='{"command": "echo x"}',
            )
            return iter(
                [
                    _make_chunk(content="", tool_calls=[td]),
                    _make_chunk(finish_reason="tool_calls"),
                ]
            )

        side_effects = [make_loop() for _ in range(wee_runtime.MAX_TOOL_ROUNDS + 2)]
        mock_client.chat.completions.create.side_effect = side_effects

        result = _run_main("ollama/test", "max rounds fallback", tools=True)
        # Fallback is either "Tool execution completed" or "Max tool rounds"
        fallback_detected = (
            "Tool execution completed" in result.stdout
            or "Max tool rounds" in result.stdout
            or "echo x" in result.stdout  # last tool result included
        )
        self.assertTrue(fallback_detected)


# ===========================================================================
# 13. API EXCEPTION HANDLING
# ===========================================================================


class TestAPIExceptionHandling(unittest.TestCase):
    """Runtime handles API errors gracefully."""

    @patch("openai.OpenAI")
    def test_api_exception_nonzero_exit(self, mock_cls):
        """OpenAI API exception causes exit code != 0."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.side_effect = Exception(
            "connection refused"
        )

        result = _run_main("ollama/test", "api failure")
        self.assertNotEqual(result.returncode, 0)

    @patch("openai.OpenAI")
    def test_api_exception_error_in_stderr(self, mock_cls):
        """API exception produces 'Error' message in stderr."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.side_effect = RuntimeError("boom")

        result = _run_main("ollama/test", "stderr error")
        self.assertIn("Error", result.stderr)

    @patch("openai.OpenAI")
    def test_tools_api_exception_nonzero_exit(self, mock_cls):
        """API exception in tools mode causes graceful non-zero exit."""
        mock_client = MagicMock()
        mock_cls.return_value = mock_client
        mock_client.chat.completions.create.side_effect = ConnectionError("unreachable")

        result = _run_main("ollama/test", "tools api fail", tools=True)
        self.assertNotEqual(result.returncode, 0)


# ===========================================================================
# 14. LIVE TESTS — OLLAMA
# ===========================================================================


@skip_ollama
class TestOllamaLiveMath(unittest.TestCase):
    """Live Ollama tests: math reasoning with Python tool."""

    def test_math_answer_correct_no_tools(self):
        """Ollama returns correct answer to simple arithmetic."""
        result = _run_cli(
            OLLAMA_MODEL,
            "What is 17 + 25? Reply with ONLY the number.",
            timeout=60,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("42", result.stdout)

    def test_math_with_python_tool(self):
        """Ollama uses Python tool to compute math and returns result."""
        result = _run_cli(
            OLLAMA_MODEL,
            "Use the python tool to compute: sum(range(1, 101)). "
            "Reply with the number only.",
            tools=True,
            timeout=LIVE_TIMEOUT,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("5050", result.stdout)

    def test_bash_tool_date_year(self):
        """Ollama uses bash tool to get current year (4-digit)."""
        import datetime

        current_year = str(datetime.datetime.now().year)
        result = _run_cli(
            OLLAMA_MODEL,
            "Use the bash tool to run: date +%Y. Reply with the year only.",
            tools=True,
            timeout=LIVE_TIMEOUT,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn(current_year, result.stdout)

    def test_text_generation_nonempty(self):
        """Ollama returns non-empty text for a creative prompt."""
        result = _run_cli(
            OLLAMA_MODEL,
            "In exactly three words, describe what color the sky is.",
            timeout=60,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertTrue(len(result.stdout.strip()) > 0)


@skip_ollama
class TestOllamaLiveSSH(unittest.TestCase):
    """Live Ollama tests: SSH sanitization applied during tool execution."""

    def test_ssh_command_sanitization_in_bash_tool(self):
        """SSH version query via bash tool completes without crash."""
        result = _run_cli(
            OLLAMA_MODEL,
            "Use the bash tool to run this exact command: ssh -V 2>&1",
            tools=True,
            timeout=LIVE_TIMEOUT,
        )
        # wee_runtime should not crash (ssh -V exits fast, no connection needed)
        self.assertNotIn("Traceback", result.stderr)
        self.assertEqual(result.returncode, 0, result.stderr)


@skip_ollama
class TestOllamaLiveMultiStep(unittest.TestCase):
    """Live Ollama tests: multi-step reasoning requiring 2+ tool calls."""

    def test_two_step_bash_then_python(self):
        """Ollama uses bash to echo a value, then reports the result."""
        result = _run_cli(
            OLLAMA_MODEL,
            "Use the bash tool to run: echo MULTISTEP_OK. Then report what the output was.",
            tools=True,
            timeout=LIVE_TIMEOUT,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("MULTISTEP_OK", result.stdout)


# ===========================================================================
# 15. LIVE TESTS — OPENROUTER
# ===========================================================================


@skip_openrouter
class TestOpenRouterLiveFree(unittest.TestCase):
    """Live OpenRouter tests using free tier models."""

    def test_simple_text_response(self):
        """OpenRouter free model returns non-empty text response."""
        result = _run_cli(
            OPENROUTER_MODEL,
            "Reply with exactly: OPENROUTER_TEST_OK",
            timeout=LIVE_TIMEOUT,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("OPENROUTER_TEST_OK", result.stdout)

    def test_math_response(self):
        """OpenRouter free model answers simple math correctly."""
        result = _run_cli(
            OPENROUTER_MODEL,
            "What is 8 times 9? Reply with ONLY the number.",
            timeout=LIVE_TIMEOUT,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn("72", result.stdout)

    def test_tool_calling_if_supported(self):
        """OpenRouter free model with tools: uses bash if supported, or falls back."""
        result = _run_cli(
            OPENROUTER_MODEL,
            "Use the bash tool to run: echo OPENROUTER_TOOL_OK. "
            "If tools are unavailable, reply: TOOLS_UNSUPPORTED.",
            tools=True,
            timeout=LIVE_TIMEOUT,
        )
        self.assertEqual(result.returncode, 0, result.stderr)
        success = (
            "OPENROUTER_TOOL_OK" in result.stdout
            or "TOOLS_UNSUPPORTED" in result.stdout
            or len(result.stdout.strip()) > 0
        )
        self.assertTrue(success, f"Expected response, got: {result.stdout[:200]}")


@skip_openrouter
class TestOpenRouterKeyRequired(unittest.TestCase):
    """Without API key, OpenRouter should fail with clear error."""

    def test_missing_key_exits_nonzero(self):
        """openrouter/ model without API key causes sys.exit(1)."""
        env = {k: v for k, v in os.environ.items() if k not in ("OPENROUTER_API_KEY",)}
        # Patch keyring to return nothing
        with patch("keyring.get_password", return_value=None):
            with patch.dict(os.environ, {"OPENROUTER_API_KEY": ""}, clear=False):
                try:
                    _, base, _ = wee_runtime.resolve_model_and_endpoint(
                        "openrouter/google/gemma"
                    )
                except SystemExit as e:
                    self.assertEqual(e.code, 1)
                    return
        # If no exception, the keyring might have a real key — skip
        self.skipTest("OPENROUTER_API_KEY found in keyring")


# ===========================================================================
# 16. WEE_TOOLS CONSTANT INTEGRITY
# ===========================================================================


class TestWeeTooDefinitions(unittest.TestCase):
    """_WEE_TOOLS constant is well-formed OpenAI tool spec."""

    def test_tools_is_list(self):
        self.assertIsInstance(wee_runtime._WEE_TOOLS, list)

    def test_two_tools_defined(self):
        self.assertEqual(len(wee_runtime._WEE_TOOLS), 2)

    def test_bash_tool_present(self):
        names = [t["function"]["name"] for t in wee_runtime._WEE_TOOLS]
        self.assertIn("bash", names)

    def test_python_tool_present(self):
        names = [t["function"]["name"] for t in wee_runtime._WEE_TOOLS]
        self.assertIn("python", names)

    def test_each_tool_has_type_function(self):
        for tool in wee_runtime._WEE_TOOLS:
            self.assertEqual(tool.get("type"), "function")

    def test_bash_requires_command_parameter(self):
        bash = next(
            t for t in wee_runtime._WEE_TOOLS if t["function"]["name"] == "bash"
        )
        params = bash["function"]["parameters"]
        self.assertIn("command", params["properties"])
        self.assertIn("command", params["required"])

    def test_python_requires_code_parameter(self):
        python = next(
            t for t in wee_runtime._WEE_TOOLS if t["function"]["name"] == "python"
        )
        params = python["function"]["parameters"]
        self.assertIn("code", params["properties"])
        self.assertIn("code", params["required"])

    def test_tools_serializable_to_json(self):
        """Tool definitions must be JSON-serializable for API."""
        serialized = json.dumps(wee_runtime._WEE_TOOLS)
        restored = json.loads(serialized)
        self.assertEqual(len(restored), 2)


# ===========================================================================
# 17. RUNTIME CONSTANTS SANITY
# ===========================================================================


class TestRuntimeConstants(unittest.TestCase):
    """MAX_TOOL_ROUNDS and TOOL_TIMEOUT are within safe operating ranges."""

    def test_max_tool_rounds_gte_5(self):
        """MAX_TOOL_ROUNDS is at least 5 to support complex agentic workflows."""
        self.assertGreaterEqual(wee_runtime.MAX_TOOL_ROUNDS, 5)

    def test_max_tool_rounds_lte_20(self):
        """MAX_TOOL_ROUNDS is capped to prevent infinite loops."""
        self.assertLessEqual(wee_runtime.MAX_TOOL_ROUNDS, 20)

    def test_tool_timeout_gte_30(self):
        """TOOL_TIMEOUT gives tools at least 30 seconds to complete."""
        self.assertGreaterEqual(wee_runtime.TOOL_TIMEOUT, 30)

    def test_tool_timeout_lte_300(self):
        """TOOL_TIMEOUT prevents runaway tools from blocking indefinitely."""
        self.assertLessEqual(wee_runtime.TOOL_TIMEOUT, 300)

    def test_ssh_bin_re_is_compiled_regex(self):
        """_SSH_BIN_RE is a compiled regex object."""
        import re

        self.assertIsInstance(wee_runtime._SSH_BIN_RE, type(re.compile("")))


if __name__ == "__main__":
    unittest.main(verbosity=2)
