"""
Unit tests for the LangChain/LangGraph integration.

Tests the pure helper functions and the LangChainCallback internal state machine.
No LangChain execution or network required — callbacks are invoked directly.

LangChainCallback gracefully falls back when langchain-core is not installed
(BaseCallbackHandler = object), so these tests run in CI without framework deps.
"""
import types
import uuid
from datetime import datetime, timezone

from visibe.integrations.langchain import (
    estimate_tokens,
    _merge_callback_into_config,
    LangChainCallback,
)


# ------------------------------------------------------------------
# estimate_tokens
# ------------------------------------------------------------------

def test_estimate_tokens_empty():
    assert estimate_tokens("") == 0


def test_estimate_tokens_returns_int():
    result = estimate_tokens("Hello world")
    assert isinstance(result, int)
    assert result >= 0


def test_estimate_tokens_proportional():
    short = estimate_tokens("Hi")
    long = estimate_tokens("Hi " * 100)
    assert long > short


def test_estimate_tokens_fallback_math():
    # Without tiktoken: len(text) // 4
    # With tiktoken: real count — still positive and proportional
    result = estimate_tokens("x" * 400)
    assert result >= 1


def test_estimate_tokens_none_guard():
    # Passing None should return 0 (falsy short-circuit)
    assert estimate_tokens(None) == 0


# ------------------------------------------------------------------
# _merge_callback_into_config
# ------------------------------------------------------------------

class _FakeCallback:
    pass


def test_merge_config_none_config():
    cb = _FakeCallback()
    result = _merge_callback_into_config(None, cb)
    assert result == {"callbacks": [cb]}


def test_merge_config_no_callbacks_key():
    cb = _FakeCallback()
    result = _merge_callback_into_config({"run_name": "test"}, cb)
    assert cb in result["callbacks"]
    assert result["run_name"] == "test"


def test_merge_config_existing_list():
    cb1 = _FakeCallback()
    cb2 = _FakeCallback()
    config = {"callbacks": [cb1]}
    result = _merge_callback_into_config(config, cb2)
    assert cb1 in result["callbacks"]
    assert cb2 in result["callbacks"]


def test_merge_config_no_duplicate():
    cb = _FakeCallback()
    config = {"callbacks": [cb]}
    result = _merge_callback_into_config(config, cb)
    assert result["callbacks"].count(cb) == 1


def test_merge_config_callback_manager_style():
    cb_new = _FakeCallback()
    cb_existing = _FakeCallback()

    class FakeManager:
        handlers = [cb_existing]

    config = {"callbacks": FakeManager()}
    result = _merge_callback_into_config(config, cb_new)
    assert cb_existing in result["callbacks"]
    assert cb_new in result["callbacks"]


def test_merge_config_does_not_mutate_original():
    cb = _FakeCallback()
    original = {"run_name": "test"}
    _merge_callback_into_config(original, cb)
    assert "callbacks" not in original


# ------------------------------------------------------------------
# LangChainCallback — init / defaults
# ------------------------------------------------------------------

def test_callback_default_content_limits():
    cb = LangChainCallback()
    assert cb._llm_content_limit == LangChainCallback.DEFAULT_LLM_CONTENT_LIMIT
    assert cb._tool_content_limit == LangChainCallback.DEFAULT_TOOL_CONTENT_LIMIT


def test_callback_custom_content_limit():
    cb = LangChainCallback(content_limit=200)
    assert cb._llm_content_limit == 200
    assert cb._tool_content_limit == 200


def test_callback_initial_state():
    cb = LangChainCallback()
    assert cb.llm_calls == []
    assert cb.tool_calls == []
    assert cb.errors == []
    assert cb.steps == []
    assert cb.status == "completed"
    assert cb.start_time is None
    assert cb.prompt is None


# ------------------------------------------------------------------
# _truncate_tool
# ------------------------------------------------------------------

def test_truncate_tool_none():
    cb = LangChainCallback()
    assert cb._truncate_tool(None) == ""


def test_truncate_tool_empty():
    cb = LangChainCallback()
    assert cb._truncate_tool("") == ""


def test_truncate_tool_within_limit():
    cb = LangChainCallback(content_limit=500)
    assert cb._truncate_tool("hello") == "hello"


def test_truncate_tool_exceeds_limit():
    cb = LangChainCallback(content_limit=10)
    result = cb._truncate_tool("x" * 20)
    assert len(result) == 13  # 10 chars + "..."
    assert result.endswith("...")


def test_truncate_tool_zero_limit():
    cb = LangChainCallback(content_limit=0)
    assert cb._truncate_tool("anything") == ""


# ------------------------------------------------------------------
# _create_step
# ------------------------------------------------------------------

def test_create_step_increments_counter():
    cb = LangChainCallback()
    cb._create_step("llm_call", datetime.now(timezone.utc))
    cb._create_step("tool_call", datetime.now(timezone.utc))
    assert cb.step_counter == 2


def test_create_step_appends_to_steps():
    cb = LangChainCallback()
    cb._create_step("llm_call", datetime.now(timezone.utc), model="gpt-4o-mini")
    assert len(cb.steps) == 1
    assert cb.steps[0]["type"] == "llm_call"
    assert cb.steps[0]["model"] == "gpt-4o-mini"


def test_create_step_returns_step_id():
    cb = LangChainCallback()
    step_id = cb._create_step("agent_start", datetime.now(timezone.utc))
    assert step_id == "step_1"


def test_create_step_registers_run_id():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    step_id = cb._create_step("agent_start", datetime.now(timezone.utc), run_id=run_id)
    assert cb._run_to_step[str(run_id)] == step_id


def test_create_step_resolves_parent():
    cb = LangChainCallback()
    parent_run_id = uuid.uuid4()
    child_run_id = uuid.uuid4()

    parent_step_id = cb._create_step("agent_start", datetime.now(timezone.utc), run_id=parent_run_id)
    cb._create_step("llm_call", datetime.now(timezone.utc), run_id=child_run_id, parent_run_id=parent_run_id)

    child_step = cb.steps[1]
    assert child_step.get("parent_span_id") == parent_step_id


# ------------------------------------------------------------------
# on_chain_start
# ------------------------------------------------------------------

def test_on_chain_start_sets_start_time():
    cb = LangChainCallback()
    cb.on_chain_start(
        serialized={"name": "MyAgent"},
        inputs={"input": "hello"},
        run_id=uuid.uuid4(),
    )
    assert cb.start_time is not None


def test_on_chain_start_captures_agent_name_from_name():
    cb = LangChainCallback()
    cb.on_chain_start(
        serialized={"name": "WeatherAgent"},
        inputs={},
        run_id=uuid.uuid4(),
    )
    assert "WeatherAgent" in cb.agents_executed


def test_on_chain_start_captures_agent_name_from_id():
    cb = LangChainCallback()
    cb.on_chain_start(
        serialized={"id": ["langchain", "agents", "WeatherBot"]},
        inputs={},
        run_id=uuid.uuid4(),
    )
    assert "WeatherBot" in cb.agents_executed


def test_on_chain_start_captures_prompt_from_input():
    cb = LangChainCallback()
    cb.on_chain_start(
        serialized={"name": "MyAgent"},
        inputs={"input": "What is the weather?"},
        run_id=uuid.uuid4(),
    )
    assert cb.prompt == "What is the weather?"


def test_on_chain_start_captures_prompt_from_question():
    cb = LangChainCallback()
    cb.on_chain_start(
        serialized={"name": "MyAgent"},
        inputs={"question": "What is 2+2?"},
        run_id=uuid.uuid4(),
    )
    assert cb.prompt == "What is 2+2?"


def test_on_chain_start_no_duplicate_agents():
    cb = LangChainCallback()
    run_id_1 = uuid.uuid4()
    run_id_2 = uuid.uuid4()
    cb.on_chain_start(serialized={"name": "MyAgent"}, inputs={}, run_id=run_id_1)
    cb.on_chain_start(serialized={"name": "MyAgent"}, inputs={}, run_id=run_id_2)
    assert cb.agents_executed.count("MyAgent") == 1


# ------------------------------------------------------------------
# on_llm_start / on_llm_end
# ------------------------------------------------------------------

def _make_llm_result(prompt_tokens=50, completion_tokens=30, output_text="Answer"):
    """Build a minimal mock LLMResult (doesn't need langchain-core)."""
    gen = types.SimpleNamespace(
        text=output_text,
        generation_info=None,
        message=None,
    )
    return types.SimpleNamespace(
        generations=[[gen]],
        llm_output={"token_usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens}},
    )


def test_on_llm_start_stores_model():
    cb = LangChainCallback()
    cb.on_llm_start(
        serialized={"kwargs": {"model_name": "gpt-4o-mini"}},
        prompts=["What is 2+2?"],
        run_id=uuid.uuid4(),
    )
    assert cb.current_llm_call is not None
    assert cb.current_llm_call["model"] == "gpt-4o-mini"


def test_on_llm_start_captures_input_text():
    cb = LangChainCallback()
    cb.on_llm_start(
        serialized={"kwargs": {"model_name": "gpt-4o-mini"}},
        prompts=["Calculate 5 * 6"],
        run_id=uuid.uuid4(),
    )
    assert cb.current_llm_call["input_text"] == "Calculate 5 * 6"


def test_on_llm_end_records_call():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_llm_start(
        serialized={"kwargs": {"model_name": "gpt-4o-mini"}},
        prompts=["What is 2+2?"],
        run_id=run_id,
    )
    cb.on_llm_end(response=_make_llm_result(50, 30), run_id=run_id)

    assert len(cb.llm_calls) == 1
    assert cb.llm_calls[0]["input_tokens"] == 50
    assert cb.llm_calls[0]["output_tokens"] == 30
    assert cb.current_llm_call is None


def test_on_llm_end_calculates_cost():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_llm_start(
        serialized={"kwargs": {"model_name": "gpt-4o-mini"}},
        prompts=["Hi"],
        run_id=run_id,
    )
    cb.on_llm_end(response=_make_llm_result(1000, 500), run_id=run_id)
    assert cb.llm_calls[0]["cost"] > 0


def test_on_llm_end_creates_step():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_llm_start(serialized={"kwargs": {"model_name": "gpt-4o-mini"}}, prompts=["Hi"], run_id=run_id)
    cb.on_llm_end(response=_make_llm_result(), run_id=run_id)

    llm_steps = [s for s in cb.steps if s["type"] == "llm_call"]
    assert len(llm_steps) == 1
    assert llm_steps[0]["model"] == "gpt-4o-mini"


# ------------------------------------------------------------------
# on_llm_error
# ------------------------------------------------------------------

def test_on_llm_error_sets_failed_status():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_llm_start(serialized={"kwargs": {"model_name": "gpt-4o-mini"}}, prompts=["Hi"], run_id=run_id)
    cb.on_llm_error(error=RuntimeError("Rate limit"), run_id=run_id)

    assert cb.status == "failed"
    assert len(cb.errors) == 1
    assert cb.errors[0]["error_type"] == "LLM_ERROR"
    assert "Rate limit" in cb.errors[0]["message"]


# ------------------------------------------------------------------
# on_tool_start / on_tool_end
# ------------------------------------------------------------------

def test_on_tool_start_stores_tool():
    cb = LangChainCallback()
    cb.on_tool_start(
        serialized={"name": "calculator"},
        input_str="2 + 2",
        run_id=uuid.uuid4(),
    )
    assert cb.current_tool_call is not None
    assert cb.current_tool_call["tool_name"] == "calculator"
    assert cb.current_tool_call["input"] == "2 + 2"


def test_on_tool_end_records_result():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_tool_start(serialized={"name": "calculator"}, input_str="2+2", run_id=run_id)
    cb.on_tool_end(output="4", run_id=run_id)

    assert len(cb.tool_calls) == 1
    assert cb.tool_calls[0]["tool_name"] == "calculator"
    assert cb.tool_calls[0]["output"] == "4"
    assert cb.tool_calls[0]["status"] == "success"
    assert cb.current_tool_call is None


def test_on_tool_end_creates_step():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_tool_start(serialized={"name": "search"}, input_str="python", run_id=run_id)
    cb.on_tool_end(output="Python is a language", run_id=run_id)

    tool_steps = [s for s in cb.steps if s["type"] == "tool_call"]
    assert len(tool_steps) == 1
    assert tool_steps[0]["tool_name"] == "search"


def test_on_tool_end_truncates_output():
    cb = LangChainCallback(content_limit=10)
    run_id = uuid.uuid4()
    cb.on_tool_start(serialized={"name": "tool"}, input_str="x", run_id=run_id)
    cb.on_tool_end(output="x" * 100, run_id=run_id)

    assert len(cb.tool_calls[0]["output"]) == 13  # 10 + "..."


# ------------------------------------------------------------------
# on_tool_error
# ------------------------------------------------------------------

def test_on_tool_error_sets_failed():
    cb = LangChainCallback()
    run_id = uuid.uuid4()
    cb.on_tool_start(serialized={"name": "web_search"}, input_str="news", run_id=run_id)
    cb.on_tool_error(error=ConnectionError("Timeout"), run_id=run_id)

    assert len(cb.tool_calls) == 1
    assert cb.tool_calls[0]["status"] == "failed"
    assert len(cb.errors) == 1
    assert cb.errors[0]["error_type"] == "TOOL_ERROR"
    assert "Timeout" in cb.errors[0]["message"]


# ------------------------------------------------------------------
# on_chain_error
# ------------------------------------------------------------------

def test_on_chain_error_sets_failed():
    cb = LangChainCallback()
    cb.on_chain_start(serialized={"name": "MyAgent"}, inputs={"input": "hi"}, run_id=uuid.uuid4())
    cb.on_chain_error(error=ValueError("Something went wrong"), run_id=uuid.uuid4())

    assert cb.status == "failed"
    assert len(cb.errors) == 1
    assert cb.errors[0]["error_type"] == "CHAIN_ERROR"
    assert "Something went wrong" in cb.errors[0]["message"]


# ------------------------------------------------------------------
# Full scenario: chain → LLM → tool → LLM
# ------------------------------------------------------------------

def test_full_agent_scenario():
    """Simulate a complete agent trace: chain → LLM call → tool → LLM call."""
    cb = LangChainCallback()

    chain_run = uuid.uuid4()
    llm_run_1 = uuid.uuid4()
    tool_run = uuid.uuid4()
    llm_run_2 = uuid.uuid4()

    # Chain starts
    cb.on_chain_start(
        serialized={"name": "Calculator"},
        inputs={"input": "What is 25 * 4?"},
        run_id=chain_run,
    )

    # First LLM call — decides to use calculator tool
    cb.on_llm_start(
        serialized={"kwargs": {"model_name": "gpt-4o-mini"}},
        prompts=["What is 25 * 4?"],
        run_id=llm_run_1,
        parent_run_id=chain_run,
    )
    cb.on_llm_end(response=_make_llm_result(40, 15), run_id=llm_run_1)

    # Tool call
    cb.on_tool_start(
        serialized={"name": "calculator"},
        input_str="25 * 4",
        run_id=tool_run,
        parent_run_id=chain_run,
    )
    cb.on_tool_end(output="100", run_id=tool_run)

    # Second LLM call — summarizes result
    cb.on_llm_start(
        serialized={"kwargs": {"model_name": "gpt-4o-mini"}},
        prompts=["The calculator returned 100. Summarize."],
        run_id=llm_run_2,
        parent_run_id=chain_run,
    )
    cb.on_llm_end(response=_make_llm_result(60, 20, "25 * 4 = 100"), run_id=llm_run_2)

    # Chain ends
    cb.on_chain_end(outputs={"output": "25 * 4 = 100"}, run_id=chain_run)

    # Assertions
    assert cb.status == "completed"
    assert cb.prompt == "What is 25 * 4?"
    assert "Calculator" in cb.agents_executed
    assert len(cb.llm_calls) == 2
    assert len(cb.tool_calls) == 1
    assert cb.tool_calls[0]["tool_name"] == "calculator"
    assert cb.tool_calls[0]["output"] == "100"
    assert len(cb.errors) == 0

    total_input = sum(c["input_tokens"] for c in cb.llm_calls)
    total_output = sum(c["output_tokens"] for c in cb.llm_calls)
    assert total_input == 100   # 40 + 60
    assert total_output == 35   # 15 + 20

    # Steps timeline: agent_start, 2x llm_call, 1x tool_call
    step_types = [s["type"] for s in cb.steps]
    assert step_types.count("agent_start") == 1
    assert step_types.count("llm_call") == 2
    assert step_types.count("tool_call") == 1
