"""
Tests for brynq.chain_orchestrator — Chain Orchestrator for multi-step LLM calls.
"""

import asyncio
import os
import pytest
from unittest.mock import MagicMock, patch, call

from brynq.chain_orchestrator import (
    ChainOrchestrator,
    ChainStep,
    ChainResult,
    StepResult,
    _substitute_variables,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _mock_chat(backend, model, messages, **kwargs):
    """Default mock that returns a deterministic response based on backend."""
    prompt = messages[-1]["content"] if messages else ""
    return f"[{backend}/{model}] response to: {prompt[:60]}"


def _mock_chat_by_step(responses: dict):
    """Return a mock chat that returns different responses per call index."""
    call_count = {"n": 0}

    def _chat(backend, model, messages, **kwargs):
        idx = call_count["n"]
        call_count["n"] += 1
        if idx in responses:
            resp = responses[idx]
            if isinstance(resp, Exception):
                raise resp
            return resp
        return f"default_response_{idx}"

    return _chat


# ---------------------------------------------------------------------------
# Tests: variable substitution
# ---------------------------------------------------------------------------


class TestSubstituteVariables:
    def test_replaces_original_input(self):
        result = _substitute_variables(
            "Analyze: {original_input}",
            "some data",
            {},
        )
        assert result == "Analyze: some data"

    def test_replaces_step_output(self):
        result = _substitute_variables(
            "Review: {step_1_output}",
            "",
            {"step_1_output": "analysis result"},
        )
        assert result == "Review: analysis result"

    def test_replaces_multiple_step_outputs(self):
        result = _substitute_variables(
            "A: {step_1_output}\nB: {step_2_output}",
            "",
            {
                "step_1_output": "first",
                "step_2_output": "second",
            },
        )
        assert result == "A: first\nB: second"

    def test_replaces_extra_vars(self):
        result = _substitute_variables(
            "Domain: {domain}, Input: {original_input}",
            "test",
            {},
            extra_vars={"domain": "finance"},
        )
        assert result == "Domain: finance, Input: test"

    def test_unknown_placeholders_left_intact(self):
        result = _substitute_variables(
            "Known: {original_input}, Unknown: {nonexistent}",
            "data",
            {},
        )
        assert result == "Known: data, Unknown: {nonexistent}"

    def test_empty_template(self):
        result = _substitute_variables("", "input", {})
        assert result == ""

    def test_no_placeholders(self):
        result = _substitute_variables("plain text", "input", {})
        assert result == "plain text"

    def test_repeated_placeholder(self):
        result = _substitute_variables(
            "{original_input} and {original_input}",
            "data",
            {},
        )
        assert result == "data and data"


# ---------------------------------------------------------------------------
# Tests: single step execution
# ---------------------------------------------------------------------------


class TestSingleStep:
    @pytest.mark.asyncio
    async def test_single_step_returns_output(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=_mock_chat):
            result = await chain.execute(
                [ChainStep(backend="ollama", model="llama3", prompt="Hello world")],
                original_input="test",
            )

        assert result.success is True
        assert len(result.steps) == 1
        assert result.steps[0].success is True
        assert result.steps[0].output.startswith("[ollama/llama3]")
        assert result.final_output == result.steps[0].output
        assert result.error is None

    @pytest.mark.asyncio
    async def test_single_step_timing(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=_mock_chat):
            result = await chain.execute(
                [ChainStep(backend="ollama", model="llama3", prompt="Hi")],
            )

        assert result.total_elapsed >= 0
        assert result.steps[0].elapsed_seconds >= 0

    @pytest.mark.asyncio
    async def test_single_step_token_estimate(self):
        chain = ChainOrchestrator()
        with patch(
            "brynq.chain_orchestrator.llm_router.chat",
            return_value="This is a response with several words in it",
        ):
            result = await chain.execute(
                [ChainStep(backend="ollama", model="llama3", prompt="Hi")],
            )

        assert result.steps[0].token_estimate > 0
        assert result.total_tokens > 0

    @pytest.mark.asyncio
    async def test_single_step_variable_substitution(self):
        captured_messages = []

        def capture_chat(backend, model, messages, **kwargs):
            captured_messages.extend(messages)
            return "response"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [
                    ChainStep(
                        backend="ollama",
                        model="llama3",
                        prompt="Analyze: {original_input}",
                    ),
                ],
                original_input="Q4 revenue data",
            )

        user_msg = [m for m in captured_messages if m["role"] == "user"][0]
        assert user_msg["content"] == "Analyze: Q4 revenue data"

    @pytest.mark.asyncio
    async def test_empty_chain_returns_empty_result(self):
        chain = ChainOrchestrator()
        result = await chain.execute([], original_input="test")

        assert result.success is True
        assert len(result.steps) == 0
        assert result.final_output == ""
        assert result.original_input == "test"

    @pytest.mark.asyncio
    async def test_system_prompt_included(self):
        captured_messages = []

        def capture_chat(backend, model, messages, **kwargs):
            captured_messages.extend(messages)
            return "response"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [
                    ChainStep(
                        backend="ollama",
                        model="llama3",
                        prompt="Do something",
                        system_prompt="You are an expert analyst.",
                    ),
                ],
            )

        roles = [m["role"] for m in captured_messages]
        assert "system" in roles
        sys_msg = [m for m in captured_messages if m["role"] == "system"][0]
        assert sys_msg["content"] == "You are an expert analyst."


# ---------------------------------------------------------------------------
# Tests: multi-step chaining
# ---------------------------------------------------------------------------


class TestMultiStepChaining:
    @pytest.mark.asyncio
    async def test_two_step_chain_passes_context(self):
        captured = []

        def tracking_chat(backend, model, messages, **kwargs):
            prompt = messages[-1]["content"]
            captured.append({"backend": backend, "prompt": prompt})
            return f"output_from_{backend}"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=tracking_chat):
            result = await chain.execute(
                [
                    ChainStep(
                        backend="claude",
                        model="sonnet",
                        prompt="Analyze: {original_input}",
                    ),
                    ChainStep(
                        backend="gemini",
                        model="flash",
                        prompt="Review: {step_1_output}",
                    ),
                ],
                original_input="revenue data",
            )

        assert result.success is True
        assert len(result.steps) == 2
        assert result.steps[0].output == "output_from_claude"
        assert result.steps[1].output == "output_from_gemini"
        assert result.final_output == "output_from_gemini"

        # Verify step 2 got step 1's output substituted
        assert "output_from_claude" in captured[1]["prompt"]

    @pytest.mark.asyncio
    async def test_three_step_chain_all_variables(self):
        call_idx = {"n": 0}

        def indexed_chat(backend, model, messages, **kwargs):
            call_idx["n"] += 1
            return f"result_{call_idx['n']}"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=indexed_chat):
            result = await chain.execute(
                [
                    ChainStep(
                        backend="ollama",
                        model="llama3",
                        prompt="Step 1: {original_input}",
                    ),
                    ChainStep(
                        backend="claude",
                        model="sonnet",
                        prompt="Step 2 sees: {step_1_output}",
                    ),
                    ChainStep(
                        backend="gemini",
                        model="flash",
                        prompt="Summary: {step_1_output} and {step_2_output}",
                    ),
                ],
                original_input="raw data",
            )

        assert result.success is True
        assert len(result.steps) == 3
        assert result.final_output == "result_3"
        assert result.total_tokens > 0

        # Step 3 should have had both prior outputs substituted
        step3_prompt = result.steps[2].prompt_rendered
        assert "result_1" in step3_prompt
        assert "result_2" in step3_prompt

    @pytest.mark.asyncio
    async def test_chain_preserves_original_input_in_all_steps(self):
        captured_prompts = []

        def capture_chat(backend, model, messages, **kwargs):
            captured_prompts.append(messages[-1]["content"])
            return "ok"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [
                    ChainStep(prompt="A: {original_input}"),
                    ChainStep(prompt="B: {original_input}"),
                ],
                original_input="shared context",
            )

        assert "shared context" in captured_prompts[0]
        assert "shared context" in captured_prompts[1]

    @pytest.mark.asyncio
    async def test_extra_vars_available_in_all_steps(self):
        captured = []

        def capture_chat(backend, model, messages, **kwargs):
            captured.append(messages[-1]["content"])
            return "ok"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [
                    ChainStep(prompt="Domain: {domain}"),
                    ChainStep(prompt="Also: {domain}"),
                ],
                extra_vars={"domain": "finance"},
            )

        assert "finance" in captured[0]
        assert "finance" in captured[1]


# ---------------------------------------------------------------------------
# Tests: error handling
# ---------------------------------------------------------------------------


class TestErrorHandling:
    @pytest.mark.asyncio
    async def test_step_failure_returns_partial_result(self):
        mock = _mock_chat_by_step({
            0: "step 1 output",
            1: RuntimeError("API rate limited"),
        })

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=mock):
            result = await chain.execute(
                [
                    ChainStep(backend="claude", model="sonnet", prompt="Step 1"),
                    ChainStep(backend="gemini", model="flash", prompt="Step 2: {step_1_output}"),
                    ChainStep(backend="ollama", model="llama3", prompt="Step 3: {step_2_output}"),
                ],
            )

        assert result.success is False
        assert result.error is not None
        assert "Step 2" in result.error
        assert "rate limited" in result.error

        # Step 1 succeeded
        assert result.steps[0].success is True
        assert result.steps[0].output == "step 1 output"

        # Step 2 failed
        assert result.steps[1].success is False
        assert result.steps[1].error is not None

        # Step 3 was never reached (chain breaks on failure)
        assert len(result.steps) == 2

        # final_output is from last successful step
        assert result.final_output == "step 1 output"

    @pytest.mark.asyncio
    async def test_first_step_failure(self):
        mock = _mock_chat_by_step({0: RuntimeError("Connection refused")})

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=mock):
            result = await chain.execute(
                [
                    ChainStep(backend="ollama", model="llama3", prompt="Step 1"),
                    ChainStep(backend="claude", model="sonnet", prompt="Step 2: {step_1_output}"),
                ],
            )

        assert result.success is False
        assert len(result.steps) == 1
        assert result.steps[0].success is False
        assert result.final_output == ""

    @pytest.mark.asyncio
    async def test_error_step_has_timing(self):
        mock = _mock_chat_by_step({0: RuntimeError("boom")})

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=mock):
            result = await chain.execute(
                [ChainStep(backend="ollama", model="llama3", prompt="Hi")],
            )

        assert result.steps[0].elapsed_seconds >= 0
        assert result.total_elapsed >= 0


# ---------------------------------------------------------------------------
# Tests: per-step API key injection (BYOK)
# ---------------------------------------------------------------------------


class TestBYOKApiKeys:
    @pytest.mark.asyncio
    async def test_claude_api_key_injected(self):
        observed_key = {}

        def spy_chat(backend, model, messages, **kwargs):
            observed_key["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY")
            return "response"

        chain = ChainOrchestrator()
        # Clear any existing key
        old = os.environ.pop("ANTHROPIC_API_KEY", None)
        try:
            with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=spy_chat):
                await chain.execute(
                    [
                        ChainStep(
                            backend="claude",
                            model="sonnet",
                            prompt="Hi",
                            api_key="user_provided_key_123",
                        ),
                    ],
                )
        finally:
            if old is not None:
                os.environ["ANTHROPIC_API_KEY"] = old

        assert observed_key["ANTHROPIC_API_KEY"] == "user_provided_key_123"

    @pytest.mark.asyncio
    async def test_gemini_api_key_injected(self):
        observed_key = {}

        def spy_chat(backend, model, messages, **kwargs):
            observed_key["GEMINI_API_KEY"] = os.environ.get("GEMINI_API_KEY")
            return "response"

        chain = ChainOrchestrator()
        old = os.environ.pop("GEMINI_API_KEY", None)
        try:
            with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=spy_chat):
                await chain.execute(
                    [
                        ChainStep(
                            backend="gemini",
                            model="flash",
                            prompt="Hi",
                            api_key="user_gemini_key_456",
                        ),
                    ],
                )
        finally:
            if old is not None:
                os.environ["GEMINI_API_KEY"] = old

        assert observed_key["GEMINI_API_KEY"] == "user_gemini_key_456"

    @pytest.mark.asyncio
    async def test_api_key_restored_after_step(self):
        original_key = "original_system_key"
        os.environ["ANTHROPIC_API_KEY"] = original_key

        try:
            with patch(
                "brynq.chain_orchestrator.llm_router.chat", return_value="ok"
            ):
                chain = ChainOrchestrator()
                await chain.execute(
                    [
                        ChainStep(
                            backend="claude",
                            model="sonnet",
                            prompt="Hi",
                            api_key="temporary_user_key",
                        ),
                    ],
                )

            assert os.environ.get("ANTHROPIC_API_KEY") == original_key
        finally:
            os.environ.pop("ANTHROPIC_API_KEY", None)

    @pytest.mark.asyncio
    async def test_api_key_restored_on_failure(self):
        original_key = "keep_this"
        os.environ["ANTHROPIC_API_KEY"] = original_key

        try:
            with patch(
                "brynq.chain_orchestrator.llm_router.chat",
                side_effect=RuntimeError("fail"),
            ):
                chain = ChainOrchestrator()
                await chain.execute(
                    [
                        ChainStep(
                            backend="claude",
                            model="sonnet",
                            prompt="Hi",
                            api_key="temp_key",
                        ),
                    ],
                )

            assert os.environ.get("ANTHROPIC_API_KEY") == original_key
        finally:
            os.environ.pop("ANTHROPIC_API_KEY", None)

    @pytest.mark.asyncio
    async def test_api_key_removed_when_none_existed(self):
        """If no key was set before, ensure it is removed after."""
        os.environ.pop("ANTHROPIC_API_KEY", None)

        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="ok"):
            chain = ChainOrchestrator()
            await chain.execute(
                [
                    ChainStep(
                        backend="claude",
                        model="sonnet",
                        prompt="Hi",
                        api_key="temp_key",
                    ),
                ],
            )

        assert "ANTHROPIC_API_KEY" not in os.environ

    @pytest.mark.asyncio
    async def test_ollama_api_key_ignored(self):
        """Local backends don't use API keys; the key field should be harmless."""
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="ok"):
            result = await chain.execute(
                [
                    ChainStep(
                        backend="ollama",
                        model="llama3",
                        prompt="Hi",
                        api_key="irrelevant_key",
                    ),
                ],
            )

        assert result.success is True

    @pytest.mark.asyncio
    async def test_different_keys_per_step(self):
        keys_seen = []

        def spy_chat(backend, model, messages, **kwargs):
            if backend == "claude":
                keys_seen.append(os.environ.get("ANTHROPIC_API_KEY"))
            elif backend == "gemini":
                keys_seen.append(os.environ.get("GEMINI_API_KEY"))
            return "ok"

        os.environ.pop("ANTHROPIC_API_KEY", None)
        os.environ.pop("GEMINI_API_KEY", None)

        try:
            chain = ChainOrchestrator()
            with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=spy_chat):
                await chain.execute(
                    [
                        ChainStep(
                            backend="claude",
                            model="sonnet",
                            prompt="Step 1",
                            api_key="claude_key_abc",
                        ),
                        ChainStep(
                            backend="gemini",
                            model="flash",
                            prompt="Step 2: {step_1_output}",
                            api_key="gemini_key_xyz",
                        ),
                    ],
                )
        finally:
            os.environ.pop("ANTHROPIC_API_KEY", None)
            os.environ.pop("GEMINI_API_KEY", None)

        assert keys_seen == ["claude_key_abc", "gemini_key_xyz"]


# ---------------------------------------------------------------------------
# Tests: step naming and metadata
# ---------------------------------------------------------------------------


class TestStepMetadata:
    @pytest.mark.asyncio
    async def test_custom_step_names(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="ok"):
            result = await chain.execute(
                [
                    ChainStep(prompt="A", name="analyze"),
                    ChainStep(prompt="B", name="review"),
                ],
            )

        assert result.steps[0].step_name == "analyze"
        assert result.steps[1].step_name == "review"

    @pytest.mark.asyncio
    async def test_default_step_names(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="ok"):
            result = await chain.execute(
                [
                    ChainStep(prompt="A"),
                    ChainStep(prompt="B"),
                    ChainStep(prompt="C"),
                ],
            )

        assert result.steps[0].step_name == "step_1"
        assert result.steps[1].step_name == "step_2"
        assert result.steps[2].step_name == "step_3"

    @pytest.mark.asyncio
    async def test_prompt_rendered_stored(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="ok"):
            result = await chain.execute(
                [ChainStep(prompt="Analyze: {original_input}")],
                original_input="the data",
            )

        assert result.steps[0].prompt_rendered == "Analyze: the data"

    @pytest.mark.asyncio
    async def test_original_input_stored_in_result(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="ok"):
            result = await chain.execute(
                [ChainStep(prompt="Hi")],
                original_input="my input",
            )

        assert result.original_input == "my input"


# ---------------------------------------------------------------------------
# Tests: kwargs forwarding (temperature, max_tokens)
# ---------------------------------------------------------------------------


class TestKwargsForwarding:
    @pytest.mark.asyncio
    async def test_temperature_forwarded(self):
        captured_kwargs = {}

        def capture_chat(backend, model, messages, **kwargs):
            captured_kwargs.update(kwargs)
            return "ok"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [ChainStep(prompt="Hi", temperature=0.2)],
            )

        assert captured_kwargs.get("temperature") == 0.2

    @pytest.mark.asyncio
    async def test_max_tokens_forwarded(self):
        captured_kwargs = {}

        def capture_chat(backend, model, messages, **kwargs):
            captured_kwargs.update(kwargs)
            return "ok"

        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [ChainStep(prompt="Hi", max_tokens=500)],
            )

        assert captured_kwargs.get("max_tokens") == 500

    @pytest.mark.asyncio
    async def test_default_max_tokens_used(self):
        captured_kwargs = {}

        def capture_chat(backend, model, messages, **kwargs):
            captured_kwargs.update(kwargs)
            return "ok"

        chain = ChainOrchestrator(default_max_tokens=4096)
        with patch("brynq.chain_orchestrator.llm_router.chat", side_effect=capture_chat):
            await chain.execute(
                [ChainStep(prompt="Hi")],
            )

        assert captured_kwargs.get("max_tokens") == 4096


# ---------------------------------------------------------------------------
# Tests: sync wrapper
# ---------------------------------------------------------------------------


class TestSyncExecution:
    def test_execute_sync_works(self):
        chain = ChainOrchestrator()
        with patch("brynq.chain_orchestrator.llm_router.chat", return_value="sync_ok"):
            result = chain.execute_sync(
                [ChainStep(prompt="Hi")],
                original_input="test",
            )

        assert result.success is True
        assert result.steps[0].output == "sync_ok"
