"""Tests for the llm_client module."""

from __future__ import annotations

import base64
from io import BytesIO
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock, patch

import numpy as np
import pytest

from vibe_bot.config import (
    CHAT_MODEL,
    EMBEDDING_ENDPOINT,
    EMBEDDING_ENDPOINT_KEY,
)
from vibe_bot.llm_client import (
    chat_complete,
    chat_completion_instruct,
    embedding,
    image_edit,
    image_generation,
)


@pytest.mark.live
def test_chat_complete_live() -> None:
    """Live call to the chat endpoint via the core async ``chat_complete``.

    Unmocked: requires network access to the configured chat API. Ported from
    the former ``test_chat_completion_think`` (its sync ``chat_completion``
    wrapper was deleted).
    """
    import asyncio

    result = asyncio.run(
        chat_complete(
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "Tell me about Everquest"},
            ],
            model=CHAT_MODEL,
            max_tokens=100,
        )
    )
    assert isinstance(result, str)


@pytest.mark.live
def test_chat_completion_instruct_live() -> None:
    """Live call to the chat endpoint via the async instruct adapter.

    Unmocked: requires network access to the configured chat API.
    """
    import asyncio

    result = asyncio.run(
        chat_completion_instruct(
            system_prompt="You are a helpful assistant.",
            user_prompt="Tell me about Everquest",
            model=CHAT_MODEL,
            max_tokens=100,
        )
    )
    assert isinstance(result, str)


def test_image_generation() -> None:
    """Image generation returns the first b64 payload from the API."""
    import asyncio

    mock_client = MagicMock()
    mock_data = MagicMock()
    mock_data.b64_json = base64.b64encode(b"fake image data").decode()
    mock_response = MagicMock()
    mock_response.data = [mock_data]
    mock_client.images.generate = AsyncMock(return_value=mock_response)

    with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
        result = asyncio.run(
            image_generation(
                prompt="Generate an image of a horse",
                model="test-image-model",
            )
        )
    assert result == base64.b64encode(b"fake image data").decode()


def test_image_generation_api_error_returns_empty() -> None:
    """A 4xx/5xx (APIStatusError) from the image API returns "" without raising."""
    import asyncio

    import openai

    mock_client = MagicMock()
    mock_client.images.generate = AsyncMock(
        side_effect=openai.APIStatusError(
            "boom",
            response=MagicMock(),
            body=None,
        )
    )

    with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
        result = asyncio.run(
            image_generation(
                prompt="Generate an image of a horse",
                model="test-image-model",
            )
        )
    assert result == ""


def test_image_edit() -> None:
    """Image edit returns the first b64 payload from the API."""
    import asyncio

    mock_client = MagicMock()
    mock_data = MagicMock()
    mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
    mock_response = MagicMock()
    mock_response.data = [mock_data]
    mock_client.images.edit = AsyncMock(return_value=mock_response)

    with patch("vibe_bot.llm.images.get_image_edit_client", return_value=mock_client):
        result = asyncio.run(
            image_edit(
                image=BytesIO(b"fake image"),
                prompt="Paint the words 'horse' on the horse.",
                model="test-image-edit-model",
            )
        )
    assert result == base64.b64encode(b"fake edited image data").decode()


def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Calculate cosine similarity between two arrays.

    Returns a value close to 1 for similar vectors,
    close to 0 for orthogonal vectors,
    and close to -1 for opposite vectors.
    """
    a_arr, b_arr = np.array(a), np.array(b)
    return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))


EMBEDDING_SIMILARITY_HIGH = 0.9
EMBEDDING_SIMILARITY_LOW = 0.5


def _mock_embedding_session(
    post: MagicMock,
) -> MagicMock:
    """Build a mock requests.Session whose .post is ``post``."""
    session = MagicMock()
    session.post = post
    return session


def test_embeddings() -> None:
    """Embedding similarity for similar and different texts."""
    mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
    mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
    mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024

    def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
        json_data = kwargs.get("json", {})
        text = json_data["input"][0]
        if "horse" in text and "donkey" not in text and "also" not in text:
            embedding_data = mock_horse_vec
        elif "also" in text:
            embedding_data = mock_horse_also_vec
        else:
            embedding_data = mock_donkey_vec
        mock_resp = MagicMock()
        mock_resp.json.return_value = {"data": [{"embedding": embedding_data}]}
        return mock_resp

    session = _mock_embedding_session(MagicMock(side_effect=mock_post))
    with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
        result1 = embedding(
            "this is a horse",
            url=EMBEDDING_ENDPOINT,
            api_key=EMBEDDING_ENDPOINT_KEY,
            model="embed",
        )
        result2 = embedding(
            "this is a horse also",
            url=EMBEDDING_ENDPOINT,
            api_key=EMBEDDING_ENDPOINT_KEY,
            model="embed",
        )
        result3 = embedding(
            "this is a donkey",
            url=EMBEDDING_ENDPOINT,
            api_key=EMBEDDING_ENDPOINT_KEY,
            model="embed",
        )
        similarity_1 = _cosine_similarity(np.array(result1), np.array(result2))
        assert similarity_1 > EMBEDDING_SIMILARITY_HIGH

        similarity_2 = _cosine_similarity(np.array(result1), np.array(result3))
        assert similarity_2 < EMBEDDING_SIMILARITY_LOW


def test_embedding_non_json_2xx_returns_empty() -> None:
    """A 2xx response with a non-JSON body must return [] without raising.

    Regression test for ``resp.json()`` sitting outside the try block, so an
    HTML error page (or any non-JSON 2xx body) raised JSONDecodeError out of
    ``embedding`` and, through it, out of ``get_conversation_context``.
    """
    mock_resp = MagicMock()
    mock_resp.raise_for_status.return_value = None
    mock_resp.json.side_effect = ValueError("<html>rate limited</html>")

    session = _mock_embedding_session(MagicMock(return_value=mock_resp))
    with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
        result = embedding(
            "this is a horse",
            url=EMBEDDING_ENDPOINT,
            api_key=EMBEDDING_ENDPOINT_KEY,
            model="embed",
        )

    assert result == []


def test_chat_client_singleton_identity() -> None:
    """The shared chat client is built once and reused across calls."""
    from vibe_bot import llm_client

    client1 = llm_client.get_chat_client()
    client2 = llm_client.get_chat_client()
    assert client1 is client2


def test_image_gen_client_singleton_identity() -> None:
    """The shared image-generation client is built once, from a cold start."""
    import vibe_bot.llm.images as images_mod

    saved = images_mod._image_gen_client
    images_mod._image_gen_client = None
    try:
        client1 = images_mod.get_image_gen_client()
        client2 = images_mod.get_image_gen_client()
        assert client1 is client2
    finally:
        images_mod._image_gen_client = saved


def test_image_edit_client_singleton_identity() -> None:
    """The shared image-edit client is built once, from a cold start."""
    import vibe_bot.llm.images as images_mod

    saved = images_mod._image_edit_client
    images_mod._image_edit_client = None
    try:
        client1 = images_mod.get_image_edit_client()
        client2 = images_mod.get_image_edit_client()
        assert client1 is client2
    finally:
        images_mod._image_edit_client = saved


def test_flows_build_no_new_clients_or_sessions(
    mock_ctx: MagicMock,
    temp_db_path: str,
) -> None:
    """A full !doodlebob + chat turn constructs no new clients or sessions.

    Every shared client and the embedding session are built once at
    "startup"; running the whole image flow and a chat turn through the real
    singletons (with only HTTP mocked) must not construct another
    AsyncOpenAI client or requests.Session.
    """
    import asyncio

    import openai
    import requests

    from vibe_bot import llm_client
    from vibe_bot.database import ChatDatabase
    from vibe_bot.services.chat_service import ChatService
    from vibe_bot.services.image_service import ImageService

    # Startup: build every shared client and the embedding session once.
    chat_client = llm_client.get_chat_client()
    gen_client = llm_client.get_image_gen_client()
    edit_client = llm_client.get_image_edit_client()
    llm_client.get_embedding_session()

    counts = {"async_openai": 0, "session": 0}
    real_session = requests.Session

    class CountingAsyncOpenAI(openai.AsyncOpenAI):
        def __init__(self, **kwargs: Any) -> None:
            counts["async_openai"] += 1
            super().__init__(**kwargs)

    def counting_session() -> requests.Session:
        counts["session"] += 1
        return real_session()

    # layout, image prompt, verify verdict, chat reply — in call order.
    completions_create = AsyncMock(
        side_effect=[
            _make_response("square", None),
            _make_response("a detailed prompt", None),
            _make_response("PASS", None),
            _make_response("a chat reply", None),
        ]
    )
    image_response = MagicMock()
    image_response.data = [MagicMock(b64_json=base64.b64encode(b"fake image").decode())]
    images_generate = AsyncMock(return_value=image_response)

    registry = MagicMock()
    registry.to_openai_tools.return_value = []

    db = ChatDatabase(db_path=temp_db_path)

    with (
        patch.object(openai, "AsyncOpenAI", CountingAsyncOpenAI),
        patch.object(requests, "Session", counting_session),
        patch.object(chat_client.chat.completions, "create", completions_create),
        patch.object(gen_client.images, "generate", images_generate),
        patch("vibe_bot.llm_client.embedding", return_value=[0.25] * 32),
    ):
        asyncio.run(
            ImageService(db, MagicMock()).generate(mock_ctx, message="a centaur")
        )
        asyncio.run(
            ChatService(db, registry).handle(
                mock_ctx,
                bot_name="alfred",
                message="hello",
                system_prompt="you are a butler",
                response_prefix="alfred response",
            )
        )

    assert counts["async_openai"] == 0
    assert counts["session"] == 0
    assert llm_client.get_chat_client() is chat_client
    assert llm_client.get_image_gen_client() is gen_client
    assert llm_client.get_image_edit_client() is edit_client


def _make_response(content: str | None, tool_calls: list[object] | None) -> MagicMock:
    """Build a mock chat completion response with the given message fields."""
    message = MagicMock()
    message.content = content
    message.tool_calls = tool_calls
    return MagicMock(choices=[MagicMock(message=message)])


def test_chat_complete_skips_non_function_tool_call() -> None:
    """A tool call that is not of type 'function' is skipped, not executed."""
    import asyncio

    from vibe_bot.llm_client import chat_complete

    called = {"n": 0}

    def tool_executor(name: str, args: dict[str, str]) -> str:
        called["n"] += 1
        return f"executed:{name}"

    custom_tool_call = MagicMock()
    custom_tool_call.type = "custom"

    mock_client = MagicMock()
    mock_client.chat.completions.create = AsyncMock(
        side_effect=[
            _make_response(content=None, tool_calls=[custom_tool_call]),
            _make_response(content="final answer", tool_calls=None),
        ]
    )

    with patch("vibe_bot.llm.chat.get_chat_client", return_value=mock_client):
        result = asyncio.run(
            chat_complete(
                [{"role": "user", "content": "hi"}],
                model="m",
                max_tokens=10,
                tool_executor=tool_executor,
            )
        )

    assert result == "final answer"
    assert called["n"] == 0


def test_tool_registry_dispatch_and_unknown_tool() -> None:
    """The registry renders schemas, dispatches known tools, and names unknowns."""
    from vibe_bot.llm_client import ToolRegistry

    def echo_tool(name: str, args: dict[str, str], **kwargs: object) -> str:
        return f"echo:{args.get('text', '')}"

    registry = ToolRegistry()
    registry.register(
        "echo",
        "Echoes the text argument back.",
        {"type": "object", "properties": {"text": {"type": "string"}}},
        echo_tool,
    )

    tools = registry.to_openai_tools()
    first = tools[0]
    assert first["type"] == "function"
    function_def = cast("dict[str, object]", first["function"])
    assert function_def["name"] == "echo"
    assert function_def["description"] == "Echoes the text argument back."

    assert registry.execute("echo", {"text": "hi"}) == "echo:hi"
    assert registry.execute("does_not_exist", {}) == "Unknown tool: does_not_exist"