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

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

import pytest

from fibinaci._testing.mock_client import MockFibinaci


_DEFAULT_USAGE = {"input_tokens": 50, "output_tokens": 30}
_DEFAULT_META = {"model_name": "gpt-4o", "finish_reason": "stop"}

_SENTINEL = object()


@dataclass
class FakeAIMessage:
    content: str = "Hello"
    usage_metadata: Any = _SENTINEL
    response_metadata: Any = _SENTINEL

    def __post_init__(self):
        if self.usage_metadata is _SENTINEL:
            self.usage_metadata = dict(_DEFAULT_USAGE)
        if self.response_metadata is _SENTINEL:
            self.response_metadata = dict(_DEFAULT_META)


class TestLangChainInstrumentation:
    @pytest.fixture
    def setup_langchain(self):
        mod = MagicMock()
        base_chat_cls = MagicMock()
        original_invoke = MagicMock(return_value=FakeAIMessage())
        original_ainvoke = MagicMock(return_value=FakeAIMessage())
        base_chat_cls.invoke = original_invoke
        base_chat_cls.ainvoke = original_ainvoke
        mod.language_models.chat_models.BaseChatModel = base_chat_cls
        return mod, base_chat_cls, original_invoke

    def test_instrument_patches_invoke(self, setup_langchain, mock_fibinaci: MockFibinaci):
        mod, base_chat_cls, original_invoke = setup_langchain
        with patch.dict("sys.modules", {
            "langchain_core": mod,
            "langchain_core.language_models": mod.language_models,
            "langchain_core.language_models.chat_models": mod.language_models.chat_models,
        }):
            from fibinaci.instrumentation.langchain import LangChainInstrumentor
            inst = LangChainInstrumentor()
            inst.instrument()
            assert base_chat_cls.invoke is not original_invoke
            inst.uninstrument()

    def test_extract_langchain_usage(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.langchain import _extract_langchain_usage

        msg = FakeAIMessage()
        inp, out, model = _extract_langchain_usage(msg)
        assert inp == 50
        assert out == 30
        assert model == "gpt-4o"

    def test_extract_langchain_usage_v1_fallback(self, mock_fibinaci: MockFibinaci):
        """Test fallback to response_metadata.token_usage."""
        from fibinaci.instrumentation.langchain import _extract_langchain_usage

        msg = FakeAIMessage(
            usage_metadata=None,
            response_metadata={"token_usage": {"prompt_tokens": 100, "completion_tokens": 60}, "model": "gpt-3.5-turbo"},
        )
        inp, out, model = _extract_langchain_usage(msg)
        assert inp == 100
        assert out == 60
        assert model == "gpt-3.5-turbo"
