"""Tests for LlamaIndex auto-instrumentation."""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from fibinaci._testing.mock_client import MockFibinaci


@dataclass
class FakeUsage:
    prompt_tokens: int = 120
    completion_tokens: int = 45


@dataclass
class FakeChatResponse:
    message: str = "test response"
    raw: dict = None
    additional_kwargs: dict = None

    def __post_init__(self):
        if self.raw is None:
            self.raw = {"usage": FakeUsage()}
        if self.additional_kwargs is None:
            self.additional_kwargs = {}


@dataclass
class FakeSourceNode:
    text: str = "source text"
    score: float = 0.9


@dataclass
class FakeQueryResponse:
    response: str = "answer"
    source_nodes: list = None

    def __post_init__(self):
        if self.source_nodes is None:
            self.source_nodes = [FakeSourceNode(), FakeSourceNode()]


class TestLlamaIndexInstrumentation:
    @pytest.fixture
    def setup_llamaindex(self):
        mod = MagicMock()
        llm_cls = MagicMock()
        original_chat = MagicMock(return_value=FakeChatResponse())
        original_achat = MagicMock(return_value=FakeChatResponse())
        llm_cls.chat = original_chat
        llm_cls.achat = original_achat
        mod.llms.LLM = llm_cls
        return mod, llm_cls, original_chat

    def test_instrument_patches_chat(self, setup_llamaindex, mock_fibinaci: MockFibinaci):
        mod, llm_cls, original_chat = setup_llamaindex
        with patch.dict("sys.modules", {
            "llama_index": mod,
            "llama_index.core": mod,
            "llama_index.core.llms": mod.llms,
        }):
            from fibinaci.instrumentation.llamaindex import LlamaIndexInstrumentor
            inst = LlamaIndexInstrumentor()
            inst.instrument()
            assert llm_cls.chat is not original_chat
            inst.uninstrument()

    def test_extract_llm_usage(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.llamaindex import _extract_llm_usage

        self_ref = MagicMock()
        self_ref.model = "gpt-4o"
        response = FakeChatResponse()
        inp, out, model = _extract_llm_usage(response, self_ref)
        assert inp == 120
        assert out == 45
        assert model == "gpt-4o"

    def test_extract_llm_usage_from_additional_kwargs(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.llamaindex import _extract_llm_usage

        self_ref = MagicMock()
        self_ref.model = "gpt-4o"
        self_ref.model_name = None
        response = FakeChatResponse(
            raw={},
            additional_kwargs={"prompt_tokens": 200, "completion_tokens": 80},
        )
        inp, out, model = _extract_llm_usage(response, self_ref)
        assert inp == 200
        assert out == 80
