import os
import urllib.error
import pytest
from unittest.mock import MagicMock, patch

from brynq import llm_router

@pytest.fixture
def mock_probes():
    with patch("brynq.llm_router._probe_ollama") as p_ollama, \
         patch("brynq.llm_router._probe_lmstudio") as p_lmstudio, \
         patch("brynq.llm_router._probe_claude") as p_claude, \
         patch("brynq.llm_router._probe_gemini") as p_gemini:
        yield p_ollama, p_lmstudio, p_claude, p_gemini

@pytest.fixture
def ollama_running(mock_probes):
    p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
    p_ollama.return_value = {"name": "ollama", "status": "running", "models": ["llama3"]}
    p_lmstudio.return_value = {"name": "lmstudio", "status": "unavailable", "models": []}
    p_claude.return_value = {"name": "claude", "status": "not_configured", "models": []}
    p_gemini.return_value = {"name": "gemini", "status": "not_configured", "models": []}
    return mock_probes

@pytest.fixture
def lmstudio_running(mock_probes):
    p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
    p_ollama.return_value = {"name": "ollama", "status": "unavailable", "models": []}
    p_lmstudio.return_value = {"name": "lmstudio", "status": "running", "models": ["local-model"]}
    p_claude.return_value = {"name": "claude", "status": "not_configured", "models": []}
    p_gemini.return_value = {"name": "gemini", "status": "not_configured", "models": []}
    return mock_probes

@pytest.fixture
def claude_configured(mock_probes):
    p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
    p_ollama.return_value = {"name": "ollama", "status": "unavailable", "models": []}
    p_lmstudio.return_value = {"name": "lmstudio", "status": "unavailable", "models": []}
    p_claude.return_value = {"name": "claude", "status": "configured", "models": ["claude-sonnet-4-20250514"]}
    p_gemini.return_value = {"name": "gemini", "status": "not_configured", "models": []}
    return mock_probes

@pytest.fixture
def nothing_available(mock_probes):
    p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
    p_ollama.return_value = {"name": "ollama", "status": "unavailable", "models": []}
    p_lmstudio.return_value = {"name": "lmstudio", "status": "unavailable", "models": []}
    p_claude.return_value = {"name": "claude", "status": "not_configured", "models": []}
    p_gemini.return_value = {"name": "gemini", "status": "not_configured", "models": []}
    return mock_probes

@pytest.fixture
def mock_chat_handlers():
    with patch("brynq.llm_router._chat_ollama") as c_ollama, \
         patch("brynq.llm_router._chat_lmstudio") as c_lmstudio, \
         patch("brynq.llm_router._chat_claude") as c_claude, \
         patch("brynq.llm_router._chat_gemini") as c_gemini:
        
        c_ollama.return_value = "ollama reply"
        c_lmstudio.return_value = "lmstudio reply"
        c_claude.return_value = "claude reply"
        c_gemini.return_value = "gemini reply"

        # Update the dispatch table for tests so it uses our mocked functions
        original_dispatch = llm_router._BACKEND_DISPATCH.copy()
        llm_router._BACKEND_DISPATCH.update({
            "ollama": c_ollama,
            "lmstudio": c_lmstudio,
            "claude": c_claude,
            "gemini": c_gemini,
        })
        
        yield c_ollama, c_lmstudio, c_claude, c_gemini
        
        # Restore original dispatch table
        llm_router._BACKEND_DISPATCH.update(original_dispatch)


class TestDiscoverBackends:
    def test_returns_four_backends(self, mock_probes):
        res = llm_router.discover_backends()
        assert len(res) == 4

    def test_ollama_running_detected(self, ollama_running):
        res = llm_router.discover_backends()
        assert res[0]["name"] == "ollama"
        assert res[0]["status"] == "running"

    def test_lmstudio_running_detected(self, lmstudio_running):
        res = llm_router.discover_backends()
        assert res[1]["name"] == "lmstudio"
        assert res[1]["status"] == "running"

    def test_claude_configured_via_env_var(self, claude_configured):
        res = llm_router.discover_backends()
        assert res[2]["name"] == "claude"
        assert res[2]["status"] == "configured"

    def test_gemini_configured_via_env_var(self, mock_probes):
        p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
        p_gemini.return_value = {"name": "gemini", "status": "configured", "models": ["gemini-2.0-flash"]}
        res = llm_router.discover_backends()
        assert res[3]["name"] == "gemini"
        assert res[3]["status"] == "configured"

    def test_missing_api_key_is_not_configured(self, mock_probes):
        p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
        p_claude.return_value = {"name": "claude", "status": "not_configured", "models": []}
        res = llm_router.discover_backends()
        assert res[2]["status"] == "not_configured"

    def test_unreachable_servers_are_unavailable(self, nothing_available):
        res = llm_router.discover_backends()
        assert all(b["status"] in ("unavailable", "not_configured") for b in res)

class TestChat:
    def test_routes_to_ollama(self, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        res = llm_router.chat("ollama", "llama3", [])
        assert res == "ollama reply"
        c_ollama.assert_called_once_with("llama3", [])

    def test_routes_to_lmstudio(self, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        res = llm_router.chat("lmstudio", "local-model", [])
        assert res == "lmstudio reply"
        c_lmstudio.assert_called_once_with("local-model", [])

    def test_routes_to_claude(self, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        res = llm_router.chat("claude", "claude-sonnet", [])
        assert res == "claude reply"
        c_claude.assert_called_once_with("claude-sonnet", [])

    def test_routes_to_gemini(self, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        res = llm_router.chat("gemini", "gemini-2", [])
        assert res == "gemini reply"
        c_gemini.assert_called_once_with("gemini-2", [])

    def test_auto_backend_picks_first_available(self, ollama_running, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        res = llm_router.chat("auto", "auto", [])
        assert res == "ollama reply"
        c_ollama.assert_called_once_with("llama3", [])

    def test_auto_model_uses_default(self, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        llm_router.chat("claude", "auto", [])
        c_claude.assert_called_once_with(llm_router._CLAUDE_DEFAULT_MODEL, [])

    def test_unknown_backend_raises_runtime_error(self):
        with pytest.raises(RuntimeError, match="Unknown backend"):
            llm_router.chat("unknown", "model", [])

    def test_forwards_kwargs_to_handler(self, mock_chat_handlers):
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        llm_router.chat("ollama", "llama3", [], temperature=0.5)
        c_ollama.assert_called_once_with("llama3", [], temperature=0.5)

class TestQuickChat:
    @patch("brynq.llm_router.chat", return_value="quick reply")
    def test_wraps_prompt_in_messages(self, mock_chat):
        res = llm_router.quick_chat("hello")
        assert res == "quick reply"
        mock_chat.assert_called_once_with(backend="auto", model="auto", messages=[{"role": "user", "content": "hello"}])

    @patch("brynq.llm_router.chat", return_value="quick reply")
    def test_passes_backend_and_model(self, mock_chat):
        llm_router.quick_chat("hello", backend="claude", model="sonnet")
        mock_chat.assert_called_once_with(backend="claude", model="sonnet", messages=[{"role": "user", "content": "hello"}])

class TestGetBestBackend:
    def test_code_prefers_claude(self, claude_configured):
        backend, model = llm_router.get_best_backend("code")
        assert backend == "claude"

    def test_fast_prefers_local(self, ollama_running):
        backend, model = llm_router.get_best_backend("fast")
        assert backend == "ollama"

    def test_quality_prefers_claude_then_gemini(self, mock_probes):
        p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
        p_gemini.return_value = {"name": "gemini", "status": "configured", "models": ["gemini-2.0-flash"]}
        p_claude.return_value = {"name": "claude", "status": "not_configured", "models": []}
        backend, model = llm_router.get_best_backend("quality")
        assert backend == "gemini"

    def test_general_returns_first_available(self, lmstudio_running):
        backend, model = llm_router.get_best_backend("general")
        assert backend == "lmstudio"

    def test_raises_when_nothing_available(self, nothing_available):
        with pytest.raises(RuntimeError, match="No LLM backends available"):
            llm_router.get_best_backend("code")

    def test_fast_picks_shortest_model_name(self, mock_probes):
        p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
        p_ollama.return_value = {"name": "ollama", "status": "running", "models": ["llama-3-8b-instruct", "qwen", "mistral-nemo"]}
        backend, model = llm_router.get_best_backend("fast")
        assert backend == "ollama"
        assert model == "qwen"

class TestIsAvailable:
    @patch("brynq.llm_router._http_json")
    def test_ollama_available(self, mock_http):
        assert llm_router.is_available("ollama") is True

    @patch("os.environ.get", return_value="key")
    def test_claude_available_via_env_var(self, mock_env):
        assert llm_router.is_available("claude") is True

    @patch("brynq.llm_router._http_json", side_effect=RuntimeError)
    def test_unavailable_returns_false(self, mock_http):
        assert llm_router.is_available("ollama") is False

class TestFallbackChain:
    def test_ollama_down_falls_to_lmstudio(self, mock_probes, mock_chat_handlers):
        p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
        p_ollama.return_value = {"name": "ollama", "status": "unavailable", "models": []}
        p_lmstudio.return_value = {"name": "lmstudio", "status": "running", "models": ["local"]}
        
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        llm_router.chat("auto", "auto", [])
        c_lmstudio.assert_called_once()

    def test_all_local_down_falls_to_claude(self, mock_probes, mock_chat_handlers):
        p_ollama, p_lmstudio, p_claude, p_gemini = mock_probes
        p_ollama.return_value = {"name": "ollama", "status": "unavailable", "models": []}
        p_lmstudio.return_value = {"name": "lmstudio", "status": "unavailable", "models": []}
        p_claude.return_value = {"name": "claude", "status": "configured", "models": ["claude"]}
        
        c_ollama, c_lmstudio, c_claude, c_gemini = mock_chat_handlers
        llm_router.chat("auto", "auto", [])
        c_claude.assert_called_once()

    def test_all_down_raises_runtime_error(self, nothing_available):
        with pytest.raises(RuntimeError, match="No LLM backends available"):
            llm_router.chat("auto", "auto", [])


class TestHttpJson:
    @patch("urllib.request.urlopen")
    def test_http_json_success(self, mock_urlopen):
        mock_urlopen.return_value.__enter__.return_value.read.return_value = b'{"a": 1}'
        res = llm_router._http_json("GET", "http://test")
        assert res == {"a": 1}

    @patch("urllib.request.urlopen")
    def test_http_json_http_error(self, mock_urlopen):
        err = urllib.error.HTTPError("url", 500, "err", {}, None)
        err.read = MagicMock(return_value=b"detail")
        mock_urlopen.side_effect = err
        import pytest
        with pytest.raises(RuntimeError):
            llm_router._http_json("GET", "http://test")

    @patch("urllib.request.urlopen")
    def test_http_json_url_error(self, mock_urlopen):
        mock_urlopen.side_effect = urllib.error.URLError("refused")
        import pytest
        with pytest.raises(RuntimeError):
            llm_router._http_json("GET", "http://test")

class TestProbes:
    @patch("brynq.llm_router._probe_url", return_value=True)
    @patch("brynq.llm_router._http_json")
    def test_probe_ollama_running(self, mock_http, mock_probe_url):
        mock_http.return_value = {"models": [{"name": "m1"}]}
        info = llm_router._probe_ollama()
        assert info["status"] == "running"
        assert info["models"] == ["m1"]

    @patch("brynq.llm_router._probe_url", return_value=True)
    @patch("brynq.llm_router._http_json")
    def test_probe_lmstudio_running(self, mock_http, mock_probe_url):
        mock_http.return_value = {"data": [{"id": "m1"}]}
        info = llm_router._probe_lmstudio()
        assert info["status"] == "running"
        assert info["models"] == ["m1"]

class TestChatImpls:
    @patch("brynq.llm_router._http_json", return_value={"message": {"content": "ok"}})
    def test_chat_ollama(self, mock_http):
        res = llm_router._chat_ollama("m", [{"role": "user", "content": "hi"}], temperature=0)
        assert res == "ok"
        mock_http.assert_called_once()
        body = mock_http.call_args[1]["body"]
        assert body["model"] == "m"

    @patch("brynq.llm_router._http_json", return_value={"choices": [{"message": {"content": "ok"}}]})
    def test_chat_lmstudio(self, mock_http):
        res = llm_router._chat_lmstudio("m", [{"role": "user", "content": "hi"}], max_tokens=10)
        assert res == "ok"
        mock_http.assert_called_once()
        body = mock_http.call_args[1]["body"]
        assert body["max_tokens"] == 10

    @patch("os.environ.get", return_value="key")
    @patch("brynq.llm_router._http_json", return_value={"content": [{"type": "text", "text": "ok"}]})
    def test_chat_claude(self, mock_http, mock_env):
        res = llm_router._chat_claude("m", [{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}])
        assert res == "ok"
        body = mock_http.call_args[1]["body"]
        assert body["system"] == "sys"
        assert body["messages"][0]["role"] == "user"

    @patch("os.environ.get", return_value="key")
    @patch("brynq.llm_router._http_json", return_value={"candidates": [{"content": {"parts": [{"text": "ok"}]}}]})
    def test_chat_gemini(self, mock_http, mock_env):
        res = llm_router._chat_gemini("m", [{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}])
        assert res == "ok"

class TestPickFirstAvailable:
    @patch("brynq.llm_router._probe_ollama", return_value={"name": "ollama", "status": "running", "models": ["m"]})
    @patch("brynq.llm_router._probe_lmstudio", return_value={"status": "unavailable"})
    @patch("brynq.llm_router._probe_claude", return_value={"status": "not_configured"})
    @patch("brynq.llm_router._probe_gemini", return_value={"status": "not_configured"})
    def test_pick_first_available(self, p1, p2, p3, p4):
        backend, model = llm_router._pick_first_available()
        assert backend == "ollama"
        assert model == "m"
