import asyncio
import importlib.metadata
import langchain_openai
import pytest
from langchain_core.prompts import PromptTemplate

from opik.integrations.langchain import OpikTracer
from opik import semantic_version
from ... import llm_constants
from ...testlib import (
    ANY_BUT_NONE,
    ANY_DICT,
    ANY_STRING,
    SpanModel,
    TraceModel,
    assert_equal,
)
from .constants import (
    EXPECTED_SHORT_OPENAI_USAGE_LOGGED_FORMAT,
    EXPECTED_FULL_OPENAI_USAGE_LOGGED_FORMAT,
)

LANGCHAIN_OPENAI_VERSION_NEWER_THAN_0_3_35 = (
    semantic_version.SemanticVersion.parse(
        importlib.metadata.version("langchain_openai")
    )
    >= "0.3.35"
)


@pytest.mark.parametrize(
    "llm_model, expected_input_prompt, expected_usage, stream_usage",
    [
        # Legacy langchain_openai.OpenAI is intentionally dropped — it hits the
        # v1/completions endpoint which doesn't serve chat-only models like
        # gpt-5-nano.
        (
            langchain_openai.ChatOpenAI,
            "Given the title of play, write a synopsys for that. Title: Documentary about Bigfoot in Paris.",
            EXPECTED_FULL_OPENAI_USAGE_LOGGED_FORMAT,
            False,
        ),
        (
            langchain_openai.ChatOpenAI,
            "Given the title of play, write a synopsys for that. Title: Documentary about Bigfoot in Paris.",
            EXPECTED_FULL_OPENAI_USAGE_LOGGED_FORMAT,
            True,
        ),
    ],
)
def test_langchain__openai_llm_is_used__token_usage_is_logged__happyflow(
    fake_backend,
    ensure_openai_configured,
    llm_model,
    expected_input_prompt,
    expected_usage,
    stream_usage,
):
    llm_args = {
        "model": llm_constants.OPENAI_GPT_NANO,
        "max_tokens": 10,
        "reasoning_effort": llm_constants.OPENAI_REASONING_EFFORT,
        "name": "custom-openai-llm-name",
    }
    if stream_usage is True:
        llm_args["stream_usage"] = stream_usage

    llm = llm_model(**llm_args)

    template = "Given the title of play, write a synopsys for that. Title: {title}."

    prompt_template = PromptTemplate(input_variables=["title"], template=template)

    synopsis_chain = prompt_template | llm
    test_prompts = {"title": "Documentary about Bigfoot in Paris"}

    callback = OpikTracer(tags=["tag1", "tag2"], metadata={"a": "b"})
    synopsis_chain.invoke(input=test_prompts, config={"callbacks": [callback]})

    callback.flush()

    expected_llm_span_input = {
        "messages": [
            [
                ANY_DICT.containing(
                    {
                        "content": expected_input_prompt,
                        "type": "human",
                    }
                ),
            ]
        ]
    }

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="RunnableSequence",
        input={"title": "Documentary about Bigfoot in Paris"},
        output=ANY_BUT_NONE,
        tags=["tag1", "tag2"],
        metadata={"a": "b", "created_from": "langchain"},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="tool",
                name="PromptTemplate",
                input={"title": "Documentary about Bigfoot in Paris"},
                output={"output": ANY_BUT_NONE},
                metadata={"created_from": "langchain"},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                spans=[],
                source="sdk",
            ),
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="custom-openai-llm-name",
                input=expected_llm_span_input,
                output=ANY_BUT_NONE,
                metadata=ANY_BUT_NONE,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                usage=expected_usage,
                spans=[],
                provider="openai",
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                source="sdk",
            ),
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert len(callback.created_traces()) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_langchain__openai_llm_is_used__sync_stream__token_usage_is_logged__happyflow(
    fake_backend,
    ensure_openai_configured,
):
    callback = OpikTracer(
        tags=["tag3", "tag4"],
        metadata={"c": "d"},
    )

    model = langchain_openai.ChatOpenAI(
        model=llm_constants.OPENAI_GPT_NANO,
        max_tokens=10,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        name="custom-openai-llm-name",
        callbacks=[callback],
        streaming=True,
        # THIS PARAM IS VERY IMPORTANT!
        # if it is explicitly set to True - token usage data will be available
        stream_usage=True,
    )

    template = "Given the title of play, write a synopsys for that. Title: {title}."
    prompt_template = PromptTemplate(input_variables=["title"], template=template)

    chain = prompt_template | model

    def stream_generator(chain, inputs):
        for chunk in chain.stream(inputs, config={"callbacks": [callback]}):
            yield chunk

    def invoke_generator(chain, inputs):
        for chunk in stream_generator(chain, inputs):
            print(chunk)

    inputs = {"title": "The Hobbit"}

    invoke_generator(chain, inputs)

    callback.flush()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="RunnableSequence",
        input={"title": "The Hobbit"},
        output=ANY_DICT,
        tags=["tag3", "tag4"],
        metadata={
            "c": "d",
            "created_from": "langchain",
        },
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="PromptTemplate",
                input={"title": "The Hobbit"},
                output=ANY_BUT_NONE,
                tags=None,
                metadata={
                    "created_from": "langchain",
                },
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                type="tool",
                model=None,
                provider=None,
                usage=None,
                spans=[],
                source="sdk",
            ),
            SpanModel(
                id=ANY_BUT_NONE,
                name="custom-openai-llm-name",
                input={
                    "messages": [
                        [
                            ANY_DICT.containing(
                                {
                                    "content": "Given the title of play, write a synopsys for that. Title: The Hobbit."
                                }
                            )
                        ]
                    ]
                },
                output=ANY_BUT_NONE,
                tags=None,
                metadata=ANY_DICT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                type="llm",
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                provider="openai",
                usage=ANY_DICT.containing(EXPECTED_SHORT_OPENAI_USAGE_LOGGED_FORMAT),
                source="sdk",
            ),
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert len(callback.created_traces()) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


@pytest.mark.skipif(
    LANGCHAIN_OPENAI_VERSION_NEWER_THAN_0_3_35,
    reason="In newer versions usage is logged anyway",
)
def test_langchain__openai_llm_is_used__async_astream__no_token_usage_is_logged__happyflow(
    fake_backend,
    ensure_openai_configured,
):
    """
    In `astream` mode, the `token_usage` is not provided by langchain.
    For trace `input` always will be = {"input": ""}
    """
    callback = OpikTracer(
        tags=["tag3", "tag4"],
        metadata={"c": "d"},
    )

    model = langchain_openai.ChatOpenAI(
        model=llm_constants.OPENAI_GPT_NANO,
        max_tokens=10,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        name="custom-openai-llm-name",
        callbacks=[callback],
        # `stream_usage` param is VERY IMPORTANT!
        # if it is explicitly set to True - token usage data will be available
        # "stream_usage": True,
    )

    template = "Given the title of play, write a synopsys for that. Title: {title}."
    prompt_template = PromptTemplate(input_variables=["title"], template=template)

    chain = prompt_template | model

    async def stream_generator(chain, inputs):
        async for chunk in chain.astream(inputs, config={"callbacks": [callback]}):
            yield chunk

    async def invoke_generator(chain, inputs):
        async for chunk in stream_generator(chain, inputs):
            print(chunk)

    inputs = {"title": "The Hobbit"}

    asyncio.run(invoke_generator(chain, inputs))

    callback.flush()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="RunnableSequence",
        input={"title": "The Hobbit"},
        output=ANY_DICT,
        tags=["tag3", "tag4"],
        metadata={
            "c": "d",
            "created_from": "langchain",
        },
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="PromptTemplate",
                input={"title": "The Hobbit"},
                output=ANY_BUT_NONE,
                tags=None,
                metadata={
                    "created_from": "langchain",
                },
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                type="tool",
                model=None,
                provider=None,
                usage=None,
                spans=[],
                source="sdk",
            ),
            SpanModel(
                id=ANY_BUT_NONE,
                name="custom-openai-llm-name",
                input={
                    "messages": [
                        [
                            ANY_DICT.containing(
                                {
                                    "content": "Given the title of play, write a synopsys for that. Title: The Hobbit.",
                                    "type": "human",
                                }
                            ),
                        ]
                    ]
                },
                output=ANY_BUT_NONE,
                tags=None,
                metadata=ANY_DICT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                type="llm",
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                provider="openai",
                usage=None,
                spans=[],
                source="sdk",
            ),
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert len(callback.created_traces()) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


@pytest.mark.skipif(
    LANGCHAIN_OPENAI_VERSION_NEWER_THAN_0_3_35,
    reason="In newer versions usage is logged anyway",
)
def test_langchain__openai_llm_is_used__sync_stream__no_token_usage_is_logged__happyflow(
    fake_backend,
    ensure_openai_configured,
):
    callback = OpikTracer(
        tags=["tag3", "tag4"],
        metadata={"c": "d"},
    )

    model = langchain_openai.ChatOpenAI(
        model=llm_constants.OPENAI_GPT_NANO,
        max_tokens=10,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        name="custom-openai-llm-name",
        callbacks=[callback],
        streaming=True,
        # `stream_usage` param is VERY IMPORTANT!
        # if it is explicitly set to True - token usage data will be available
        # "stream_usage": True,
    )

    template = "Given the title of play, write a synopsys for that. Title: {title}."
    prompt_template = PromptTemplate(input_variables=["title"], template=template)

    chain = prompt_template | model

    def stream_generator(chain, inputs):
        for chunk in chain.stream(inputs, config={"callbacks": [callback]}):
            yield chunk

    def invoke_generator(chain, inputs):
        for chunk in stream_generator(chain, inputs):
            print(chunk)

    inputs = {"title": "The Hobbit"}

    invoke_generator(chain, inputs)

    callback.flush()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="RunnableSequence",
        input={"title": "The Hobbit"},
        output=ANY_DICT,
        tags=["tag3", "tag4"],
        metadata={
            "c": "d",
            "created_from": "langchain",
        },
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="PromptTemplate",
                input={"title": "The Hobbit"},
                output=ANY_BUT_NONE,
                tags=None,
                metadata={
                    "created_from": "langchain",
                },
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                type="tool",
                model=None,
                provider=None,
                usage=None,
                spans=[],
                source="sdk",
            ),
            SpanModel(
                id=ANY_BUT_NONE,
                name="custom-openai-llm-name",
                input={
                    "messages": [
                        [
                            {
                                "content": "Given the title of play, write a synopsys for that. Title: The Hobbit.",
                                "additional_kwargs": {},
                                "response_metadata": {},
                                "type": "human",
                                "name": None,
                                "id": None,
                                "example": False,
                            }
                        ]
                    ]
                },
                output=ANY_BUT_NONE,
                tags=None,
                metadata=ANY_DICT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                type="llm",
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                provider="openai",
                usage=None,
                spans=[],
                source="sdk",
            ),
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert len(callback.created_traces()) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_langchain__openai_llm_is_used__error_occurred_during_openai_call__error_info_is_logged(
    fake_backend,
):
    llm = langchain_openai.OpenAI(
        max_tokens=10, name="custom-openai-llm-name", api_key="incorrect-api-key"
    )

    template = "Given the title of play, write a synopsys for that. Title: {title}."

    prompt_template = PromptTemplate(input_variables=["title"], template=template)

    synopsis_chain = prompt_template | llm
    test_prompts = {"title": "Documentary about Bigfoot in Paris"}

    callback = OpikTracer(tags=["tag1", "tag2"], metadata={"a": "b"})
    with pytest.raises(Exception):
        synopsis_chain.invoke(input=test_prompts, config={"callbacks": [callback]})

    callback.flush()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="RunnableSequence",
        input={"title": "Documentary about Bigfoot in Paris"},
        output=None,
        tags=["tag1", "tag2"],
        metadata={
            "a": "b",
            "created_from": "langchain",
        },
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        error_info={
            "exception_type": ANY_STRING,
            "traceback": ANY_STRING,
            "message": None,
        },
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="tool",
                name="PromptTemplate",
                input={"title": "Documentary about Bigfoot in Paris"},
                output={"output": ANY_BUT_NONE},
                metadata={
                    "created_from": "langchain",
                },
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                spans=[],
                source="sdk",
            ),
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="custom-openai-llm-name",
                input={
                    "prompts": [
                        "Given the title of play, write a synopsys for that. Title: Documentary about Bigfoot in Paris."
                    ]
                },
                output=None,
                metadata=ANY_BUT_NONE,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                usage=None,
                error_info={
                    "exception_type": ANY_STRING,
                    "traceback": ANY_STRING,
                    "message": None,
                },
                spans=[],
                source="sdk",
            ),
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert len(callback.created_traces()) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_langchain__find_token_usage_dict__multi_turn_returns_latest():
    """
    Test that find_token_usage_dict returns the most recent usage_metadata.

    This is a regression test for a bug where the first token usage was always returned
    instead of the most recent one in multi-turn conversations.
    """
    from opik.integrations.langchain.provider_usage_extractors.langchain_run_helpers import (
        helpers,
    )

    multi_turn_run_dict = {
        "id": "run-123",
        "name": "ChatOpenAI",
        "inputs": {
            "messages": [{"role": "user", "content": "what is the weather in sf"}]
        },
        "outputs": {
            "generations": [
                [
                    {
                        "message": {
                            "content": "I'll check the weather for you.",
                            "kwargs": {
                                "usage_metadata": {
                                    "input_tokens": 150,
                                    "output_tokens": 25,
                                    "total_tokens": 175,
                                }
                            },
                        }
                    }
                ]
            ]
        },
        "events": [
            {
                "event": "on_chat_model_stream",
                "data": {
                    "chunk": {
                        "kwargs": {
                            "usage_metadata": {
                                "input_tokens": 150,
                                "output_tokens": 25,
                                "total_tokens": 175,
                            }
                        }
                    }
                },
            },
            {
                "event": "on_chat_model_stream",
                "data": {
                    "chunk": {
                        "kwargs": {
                            "usage_metadata": {
                                "input_tokens": 190,
                                "output_tokens": 13,
                                "total_tokens": 203,
                            }
                        }
                    }
                },
            },
        ],
    }

    candidate_keys = {"input_tokens", "output_tokens", "total_tokens"}
    result = helpers.find_token_usage_dict(
        multi_turn_run_dict, candidate_keys, all_keys_should_match=False
    )

    assert result is not None
    assert result["input_tokens"] == 190
    assert result["output_tokens"] == 13
    assert result["total_tokens"] == 203