"""Tests for LLM provider abstraction: base, vLLM, Anthropic, OpenAI, registry, router."""

from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

from orxaq.orchestration.providers.anthropic import AnthropicProvider
from orxaq.orchestration.providers.base import LLMProvider, LLMResponse
from orxaq.orchestration.providers.openai import OpenAIProvider
from orxaq.orchestration.providers.registry import ProviderRegistry
from orxaq.orchestration.providers.router import (
    DEFAULT_PRIORITY,
    TASK_ROUTES,
    ProviderRouter,
)
from orxaq.orchestration.providers.vllm import VLLMProvider

# ── LLMResponse ──────────────────────────────────────────────


class TestLLMResponse:
    def test_creation(self):
        r = LLMResponse(text="hello", model="m", provider="p")
        assert r.text == "hello"
        assert r.model == "m"
        assert r.provider == "p"
        assert r.input_tokens == 0
        assert r.output_tokens == 0
        assert r.latency_ms == 0.0
        assert r.metadata == {}

    def test_frozen(self):
        r = LLMResponse(text="x", model="m", provider="p")
        try:
            r.text = "y"  # type: ignore[misc]
            assert False, "should be frozen"
        except AttributeError:
            pass

    def test_with_tokens(self):
        r = LLMResponse(
            text="out",
            model="gpt-4o",
            provider="openai",
            input_tokens=100,
            output_tokens=50,
            latency_ms=123.4,
        )
        assert r.input_tokens == 100
        assert r.output_tokens == 50
        assert r.latency_ms == 123.4

    def test_metadata(self):
        r = LLMResponse(text="", model="m", provider="p", metadata={"error": "fail"})
        assert r.metadata["error"] == "fail"


# ── LLMProvider base ─────────────────────────────────────────


class TestLLMProvider:
    def test_base_generate_raises(self):
        p = LLMProvider()
        try:
            p.generate("hello")
            assert False, "should raise"
        except NotImplementedError:
            pass

    def test_base_health_check(self):
        h = LLMProvider().health_check()
        assert h["healthy"] is False
        assert h["error"] == "not implemented"

    def test_available_models(self):
        p = LLMProvider()
        p.models = ["a", "b"]
        assert p.available_models() == ["a", "b"]


# ── VLLMProvider ─────────────────────────────────────────────


class TestVLLMProvider:
    def test_init_defaults(self):
        p = VLLMProvider()
        assert p.name == "vllm"
        assert p.endpoint == "http://localhost:8000/v1"
        assert p.models == ["default"]

    def test_init_custom(self):
        p = VLLMProvider(endpoint="http://gpu:9000/v1", models=["llama-3"])
        assert p.endpoint == "http://gpu:9000/v1"
        assert p.models == ["llama-3"]

    def test_endpoint_trailing_slash(self):
        p = VLLMProvider(endpoint="http://localhost:8000/v1/")
        assert p.endpoint == "http://localhost:8000/v1"

    def test_generate_success(self):
        mock_response = {
            "choices": [{"message": {"content": "Generated text"}}],
            "usage": {"prompt_tokens": 10, "completion_tokens": 20},
        }
        body = json.dumps(mock_response).encode()

        mock_resp = MagicMock()
        mock_resp.read.return_value = body
        mock_resp.__enter__ = lambda s: s
        mock_resp.__exit__ = MagicMock(return_value=False)

        with patch("urllib.request.urlopen", return_value=mock_resp):
            p = VLLMProvider()
            r = p.generate("test prompt")
            assert r.text == "Generated text"
            assert r.provider == "vllm"
            assert r.input_tokens == 10
            assert r.output_tokens == 20
            assert r.latency_ms > 0

    def test_generate_network_error(self):
        import urllib.error

        with patch(
            "urllib.request.urlopen",
            side_effect=urllib.error.URLError("Connection refused"),
        ):
            p = VLLMProvider()
            r = p.generate("test")
            assert r.text == ""
            assert "error" in r.metadata
            assert "Connection refused" in r.metadata["error"]

    def test_generate_custom_model(self):
        mock_response = {
            "choices": [{"message": {"content": "ok"}}],
            "usage": {},
        }
        body = json.dumps(mock_response).encode()
        mock_resp = MagicMock()
        mock_resp.read.return_value = body
        mock_resp.__enter__ = lambda s: s
        mock_resp.__exit__ = MagicMock(return_value=False)

        with patch("urllib.request.urlopen", return_value=mock_resp):
            p = VLLMProvider()
            r = p.generate("test", model="llama-3")
            assert r.model == "llama-3"

    def test_health_check_success(self):
        mock_response = {"data": [{"id": "model-a"}, {"id": "model-b"}]}
        body = json.dumps(mock_response).encode()
        mock_resp = MagicMock()
        mock_resp.read.return_value = body
        mock_resp.__enter__ = lambda s: s
        mock_resp.__exit__ = MagicMock(return_value=False)

        with patch("urllib.request.urlopen", return_value=mock_resp):
            p = VLLMProvider()
            h = p.health_check()
            assert h["healthy"] is True
            assert "model-a" in h["models"]

    def test_health_check_failure(self):
        import urllib.error

        with patch(
            "urllib.request.urlopen",
            side_effect=urllib.error.URLError("timeout"),
        ):
            p = VLLMProvider()
            h = p.health_check()
            assert h["healthy"] is False
            assert "timeout" in h["error"]


# ── AnthropicProvider ────────────────────────────────────────


class TestAnthropicProvider:
    def test_init_defaults(self):
        with patch.dict("os.environ", {}, clear=True):
            p = AnthropicProvider(auth="test-key")
            assert p.name == "anthropic"
            assert p._auth == "test-key"
            assert "claude" in p.models[0]

    def test_no_api_key(self):
        with patch.dict("os.environ", {}, clear=True):
            p = AnthropicProvider(auth="")
            r = p.generate("test")
            assert r.text == ""
            assert "not set" in r.metadata["error"]

    def test_generate_success(self):
        mock_response = {
            "content": [{"type": "text", "text": "Claude says hi"}],
            "usage": {"input_tokens": 5, "output_tokens": 15},
        }
        body = json.dumps(mock_response).encode()
        mock_resp = MagicMock()
        mock_resp.read.return_value = body
        mock_resp.__enter__ = lambda s: s
        mock_resp.__exit__ = MagicMock(return_value=False)

        with patch("urllib.request.urlopen", return_value=mock_resp):
            p = AnthropicProvider(auth="sk-test")
            r = p.generate("hello")
            assert r.text == "Claude says hi"
            assert r.provider == "anthropic"
            assert r.input_tokens == 5

    def test_generate_error(self):
        import urllib.error

        with patch(
            "urllib.request.urlopen",
            side_effect=urllib.error.URLError("401 Unauthorized"),
        ):
            p = AnthropicProvider(auth="bad-key")
            r = p.generate("test")
            assert r.text == ""
            assert "401" in r.metadata["error"]

    def test_health_check_no_key(self):
        with patch.dict("os.environ", {}, clear=True):
            p = AnthropicProvider(auth="")
            h = p.health_check()
            assert h["healthy"] is False
            assert "not set" in h["error"]

    def test_health_check_with_key(self):
        p = AnthropicProvider(auth="sk-test")
        h = p.health_check()
        assert h["healthy"] is True
        assert h["provider"] == "anthropic"


# ── OpenAIProvider ───────────────────────────────────────────


class TestOpenAIProvider:
    def test_init_defaults(self):
        with patch.dict("os.environ", {}, clear=True):
            p = OpenAIProvider(auth="test-key")
            assert p.name == "openai"
            assert p.models == ["gpt-4o"]

    def test_no_api_key(self):
        with patch.dict("os.environ", {}, clear=True):
            p = OpenAIProvider(auth="")
            r = p.generate("test")
            assert r.text == ""
            assert "not set" in r.metadata["error"]

    def test_generate_success(self):
        mock_response = {
            "choices": [{"message": {"content": "GPT says hi"}}],
            "usage": {"prompt_tokens": 8, "completion_tokens": 12},
        }
        body = json.dumps(mock_response).encode()
        mock_resp = MagicMock()
        mock_resp.read.return_value = body
        mock_resp.__enter__ = lambda s: s
        mock_resp.__exit__ = MagicMock(return_value=False)

        with patch("urllib.request.urlopen", return_value=mock_resp):
            p = OpenAIProvider(auth="sk-test")
            r = p.generate("hello")
            assert r.text == "GPT says hi"
            assert r.provider == "openai"
            assert r.input_tokens == 8
            assert r.output_tokens == 12

    def test_generate_error(self):
        import urllib.error

        with patch(
            "urllib.request.urlopen",
            side_effect=urllib.error.URLError("rate limit"),
        ):
            p = OpenAIProvider(auth="sk-test")
            r = p.generate("test")
            assert r.text == ""
            assert "rate limit" in r.metadata["error"]

    def test_health_check_no_key(self):
        with patch.dict("os.environ", {}, clear=True):
            p = OpenAIProvider(auth="")
            h = p.health_check()
            assert h["healthy"] is False

    def test_health_check_with_key(self):
        p = OpenAIProvider(auth="sk-test")
        h = p.health_check()
        assert h["healthy"] is True


# ── ProviderRegistry ─────────────────────────────────────────


class TestProviderRegistry:
    def test_register_and_get(self):
        reg = ProviderRegistry()
        p = VLLMProvider()
        reg.register(p)
        assert reg.get("vllm") is p

    def test_get_missing(self):
        reg = ProviderRegistry()
        assert reg.get("nonexistent") is None

    def test_list_providers(self):
        reg = ProviderRegistry()
        reg.register(VLLMProvider())
        reg.register(OpenAIProvider(auth="k"))
        names = reg.list_providers()
        assert names == ["openai", "vllm"]

    def test_all_providers(self):
        reg = ProviderRegistry()
        reg.register(VLLMProvider())
        assert len(reg.all_providers()) == 1

    def test_len_and_contains(self):
        reg = ProviderRegistry()
        assert len(reg) == 0
        assert "vllm" not in reg
        reg.register(VLLMProvider())
        assert len(reg) == 1
        assert "vllm" in reg

    def test_remove(self):
        reg = ProviderRegistry()
        reg.register(VLLMProvider())
        assert reg.remove("vllm") is True
        assert reg.remove("vllm") is False
        assert len(reg) == 0

    def test_health_check_all(self):
        reg = ProviderRegistry()
        reg.register(VLLMProvider())
        with patch.dict("os.environ", {}, clear=True):
            reg.register(OpenAIProvider(auth=""))
        results = reg.health_check_all()
        assert "vllm" in results
        assert "openai" in results

    def test_available_models(self):
        reg = ProviderRegistry()
        reg.register(VLLMProvider(models=["llama-3"]))
        reg.register(OpenAIProvider(auth="k", models=["gpt-4o"]))
        models = reg.available_models()
        assert models["vllm"] == ["llama-3"]
        assert models["openai"] == ["gpt-4o"]

    def test_default_factory(self):
        with patch.dict("os.environ", {}, clear=True):
            reg = ProviderRegistry.default()
        assert "vllm" in reg
        assert "anthropic" in reg
        assert "openai" in reg
        assert len(reg) == 3


# ── ProviderRouter ───────────────────────────────────────────


def _make_mock_provider(name: str, text: str = "ok", error: str | None = None):
    """Create a mock provider that returns a fixed response."""
    p = MagicMock(spec=LLMProvider)
    p.name = name
    meta = {"error": error} if error else {}
    p.generate.return_value = LLMResponse(
        text=text if not error else "",
        model="test-model",
        provider=name,
        metadata=meta,
    )
    p.health_check.return_value = {"healthy": not error, "provider": name}
    p.available_models.return_value = ["test-model"]
    return p


class TestProviderRouter:
    def test_route_first_provider(self):
        reg = ProviderRegistry()
        p = _make_mock_provider("vllm", text="local response")
        reg.register(p)

        router = ProviderRouter(reg)
        r = router.route("hello")
        assert r.text == "local response"
        assert r.provider == "vllm"

    def test_route_fallback_on_error(self):
        reg = ProviderRegistry()
        p1 = _make_mock_provider("vllm", error="server down")
        p2 = _make_mock_provider("anthropic", text="cloud response")
        reg.register(p1)
        reg.register(p2)

        router = ProviderRouter(reg)
        r = router.route("hello")
        assert r.text == "cloud response"
        assert r.provider == "anthropic"

    def test_route_all_fail(self):
        reg = ProviderRegistry()
        reg.register(_make_mock_provider("vllm", error="down"))
        reg.register(_make_mock_provider("anthropic", error="401"))

        router = ProviderRouter(reg)
        r = router.route("hello")
        assert r.text == ""
        assert "all providers failed" in r.metadata["error"]
        assert len(r.metadata["details"]) == 2

    def test_route_no_providers(self):
        reg = ProviderRegistry()
        router = ProviderRouter(reg)
        r = router.route("hello")
        assert "no providers registered" in r.metadata["error"]

    def test_task_type_routing(self):
        reg = ProviderRegistry()
        vllm = _make_mock_provider("vllm", text="fast")
        anthropic = _make_mock_provider("anthropic", text="smart")
        reg.register(vllm)
        reg.register(anthropic)

        router = ProviderRouter(reg)

        r = router.route("speed test", task_type="fast")
        assert r.provider == "vllm"

        r = router.route("reason about this", task_type="reasoning")
        assert r.provider == "anthropic"

    def test_stats_tracking(self):
        reg = ProviderRegistry()
        reg.register(_make_mock_provider("vllm", text="ok"))
        router = ProviderRouter(reg)

        router.route("a")
        router.route("b")
        stats = router.stats()

        assert stats["priority"] == DEFAULT_PRIORITY
        assert stats["latency"]["vllm"]["calls"] == 2

    def test_failure_count_tracking(self):
        reg = ProviderRegistry()
        p = _make_mock_provider("vllm", error="fail")
        reg.register(p)

        router = ProviderRouter(reg)
        router.route("test")
        stats = router.stats()
        assert stats["failures"]["vllm"] == 1

    def test_health_check_delegates(self):
        reg = ProviderRegistry()
        reg.register(_make_mock_provider("vllm"))
        router = ProviderRouter(reg)
        h = router.health_check()
        assert "vllm" in h

    def test_custom_priority(self):
        reg = ProviderRegistry()
        reg.register(_make_mock_provider("openai", text="gpt"))
        reg.register(_make_mock_provider("vllm", text="local"))

        router = ProviderRouter(reg, priority=["openai", "vllm"])
        r = router.route("test")
        assert r.provider == "openai"

    def test_task_routes_defined(self):
        assert "fast" in TASK_ROUTES
        assert "reasoning" in TASK_ROUTES
        assert "general" in TASK_ROUTES


# ── CLI integration ──────────────────────────────────────────


class TestProvidersCLI:
    def test_providers_check_command(self):
        from orxaq.experiences.cli import main

        with patch.dict("os.environ", {}, clear=True):
            ret = main(["--json", "providers", "check"])
        assert ret == 0

    def test_providers_list_command(self):
        from orxaq.experiences.cli import main

        with patch.dict("os.environ", {}, clear=True):
            ret = main(["--json", "providers", "list"])
        assert ret == 0

    def test_providers_check_text(self, capsys):
        from orxaq.experiences.cli import main

        with patch.dict("os.environ", {}, clear=True):
            ret = main(["providers", "check"])
        assert ret == 0
        out = capsys.readouterr().out
        assert "Provider Health" in out

    def test_providers_list_text(self, capsys):
        from orxaq.experiences.cli import main

        with patch.dict("os.environ", {}, clear=True):
            ret = main(["providers", "list"])
        assert ret == 0
        out = capsys.readouterr().out
        assert "Available Models" in out

    def test_providers_cost_json(self):
        from orxaq.experiences.cli import main

        with patch.dict("os.environ", {}, clear=True):
            ret = main(["--json", "providers", "cost"])
        assert ret == 0

    def test_providers_cost_text(self, capsys):
        from orxaq.experiences.cli import main

        with patch.dict("os.environ", {}, clear=True):
            ret = main(["providers", "cost"])
        assert ret == 0
        out = capsys.readouterr().out
        assert "Pricing" in out


# ── Cost Tracker ────────────────────────────────────────────


class TestCostTracker:
    def test_default_pricing_exists(self):
        from orxaq.orchestration.providers.cost import DEFAULT_PRICING

        assert "vllm" in DEFAULT_PRICING
        assert "anthropic" in DEFAULT_PRICING
        assert "openai" in DEFAULT_PRICING

    def test_vllm_is_free(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        cost = tracker.compute_cost("vllm", "default", 1000, 1000)
        assert cost == 0.0

    def test_anthropic_pricing(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        cost = tracker.compute_cost("anthropic", "claude-3-sonnet", 1000, 1000)
        assert cost > 0.0

    def test_record_creates_entry(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        rec = tracker.record("anthropic", "claude-3-sonnet", 500, 200, latency_ms=100.0)
        assert rec.provider == "anthropic"
        assert rec.model == "claude-3-sonnet"
        assert rec.input_tokens == 500
        assert rec.output_tokens == 200
        assert rec.cost_usd > 0.0
        assert rec.latency_ms == 100.0

    def test_summary_empty(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        s = tracker.summary()
        assert s["total_cost_usd"] == 0.0
        assert s["total_calls"] == 0

    def test_summary_with_records(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        tracker.record("vllm", "default", 100, 100)
        tracker.record("anthropic", "default", 100, 100)
        s = tracker.summary()
        assert s["total_calls"] == 2
        assert "vllm" in s["by_provider"]
        assert "anthropic" in s["by_provider"]
        assert s["by_provider"]["anthropic"]["cost_usd"] > 0.0

    def test_set_custom_pricing(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        tracker.set_pricing("custom_provider", "model-v1", 0.01, 0.02)
        cost = tracker.compute_cost("custom_provider", "model-v1", 1000, 1000)
        assert cost == (0.01 * 1000 + 0.02 * 1000) / 1000.0

    def test_unknown_provider_free(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        cost = tracker.compute_cost("unknown", "unknown", 1000, 1000)
        assert cost == 0.0

    def test_records_property(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        tracker.record("vllm", "default", 100, 100)
        recs = tracker.records
        assert len(recs) == 1
        # Returns a copy
        recs.clear()
        assert len(tracker.records) == 1

    def test_reset(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        tracker.record("vllm", "default", 100, 100)
        assert len(tracker.records) == 1
        tracker.reset()
        assert len(tracker.records) == 0


class TestBudgetEnforcement:
    def test_unlimited_budget(self):
        from orxaq.orchestration.providers.cost import CostTracker

        tracker = CostTracker()
        assert tracker.is_within_budget()

    def test_budget_check_structure(self):
        from orxaq.orchestration.providers.cost import BudgetConfig, CostTracker

        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=1.0, soft_limit_usd=0.5))
        check = tracker.check_budget()
        assert "spent_usd" in check
        assert "within_budget" in check
        assert "warning" in check
        assert check["within_budget"] is True

    def test_hard_limit_exceeded(self):
        from orxaq.orchestration.providers.cost import BudgetConfig, CostTracker

        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=0.001))
        tracker.record("anthropic", "default", 10000, 10000)
        assert not tracker.is_within_budget()
        check = tracker.check_budget()
        assert check["within_budget"] is False

    def test_soft_limit_warning(self):
        from orxaq.orchestration.providers.cost import BudgetConfig, CostTracker

        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=10.0, soft_limit_usd=0.001))
        tracker.record("anthropic", "default", 10000, 10000)
        check = tracker.check_budget()
        assert check["warning"] is True
        assert check["within_budget"] is True


class TestCostAwareRouter:
    def test_rank_by_cost(self):
        from orxaq.orchestration.providers.cost import CostAwareRouter, CostTracker

        tracker = CostTracker()
        router = CostAwareRouter(tracker)
        ranked = router.rank_by_cost(["openai", "anthropic", "vllm"])
        assert ranked[0] == "vllm"
        assert ranked[-1] == "openai"

    def test_select_filters_expensive(self):
        from orxaq.orchestration.providers.cost import (
            BudgetConfig,
            CostAwareRouter,
            CostTracker,
        )

        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=0.001))
        tracker.record("anthropic", "default", 10000, 10000)
        router = CostAwareRouter(tracker)
        selected = router.select(["vllm", "anthropic", "openai"])
        # vllm is free, should be selected
        assert "vllm" in selected

    def test_should_skip_free_provider(self):
        from orxaq.orchestration.providers.cost import CostAwareRouter, CostTracker

        tracker = CostTracker()
        router = CostAwareRouter(tracker)
        assert not router.should_skip("vllm")

    def test_select_returns_at_least_one(self):
        from orxaq.orchestration.providers.cost import (
            BudgetConfig,
            CostAwareRouter,
            CostTracker,
        )

        tracker = CostTracker(budget=BudgetConfig(hard_limit_usd=0.0000001))
        tracker.record("anthropic", "default", 100000, 100000)
        router = CostAwareRouter(tracker)
        selected = router.select(["anthropic"])
        # Even when all over budget, returns at least one
        assert len(selected) >= 1

    def test_no_budget_no_skip(self):
        from orxaq.orchestration.providers.cost import CostAwareRouter, CostTracker

        tracker = CostTracker()
        router = CostAwareRouter(tracker)
        assert not router.should_skip("openai", estimated_tokens=100000)
