"""Tests for Vertex._accumulate_upstream_token_usage() and _extract_token_usage()."""

from types import SimpleNamespace
from unittest.mock import MagicMock

from lfx.graph.schema import ResultData
from lfx.graph.vertex.base import Vertex
from lfx.schema.properties import Usage


def _make_edge(source_id: str, target_id: str):
    """Create a lightweight edge stub with source_id and target_id."""
    return SimpleNamespace(source_id=source_id, target_id=target_id)


def _make_vertex_stub(
    vertex_id: str,
    *,
    is_output: bool = False,
    token_usage: Usage | None = None,
    own_token_usage: Usage | None = None,
):
    """Create a lightweight stub that behaves like a Vertex for accumulation tests."""
    stub = MagicMock(spec=Vertex)
    stub.id = vertex_id
    stub.is_output = is_output

    if token_usage is not None:
        stub.result = ResultData(token_usage=token_usage)
    else:
        stub.result = None

    if own_token_usage is not None:
        stub.custom_component = SimpleNamespace(_token_usage=own_token_usage)
    else:
        stub.custom_component = None

    return stub


def _wire_graph(output: MagicMock, vertices: dict[str, MagicMock], edges: list):
    """Wire a mock graph with edges and vertex_map, and bind real traversal methods."""
    graph = MagicMock()
    graph.edges = edges
    graph.get_vertex = lambda vid: vertices[vid]
    output.graph = graph
    # Bind real methods so they work on the mock
    output._get_all_upstream_vertices = lambda: Vertex._get_all_upstream_vertices(output)
    output._accumulate_upstream_token_usage = lambda: Vertex._accumulate_upstream_token_usage(output)


class TestAccumulateUpstreamTokenUsage:
    """Tests for Vertex._accumulate_upstream_token_usage()."""

    def test_single_llm_upstream(self):
        """Single LLM predecessor: returns that LLM's token usage."""
        llm = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=100, output_tokens=50, total_tokens=150),
        )
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        _wire_graph(output, {"LLM-1": llm}, [_make_edge("LLM-1", "ChatOutput-1")])

        result = output._accumulate_upstream_token_usage()

        assert isinstance(result, Usage)
        assert result.input_tokens == 100
        assert result.output_tokens == 50
        assert result.total_tokens == 150

    def test_two_serial_llms(self):
        """LLM1 -> LLM2 -> ChatOutput: returns accumulated total from both."""
        llm1 = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=100, output_tokens=50, total_tokens=150),
        )
        llm2 = _make_vertex_stub(
            "LLM-2",
            token_usage=Usage(input_tokens=200, output_tokens=80, total_tokens=280),
        )
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        edges = [_make_edge("LLM-1", "LLM-2"), _make_edge("LLM-2", "ChatOutput-1")]
        _wire_graph(output, {"LLM-1": llm1, "LLM-2": llm2}, edges)

        result = output._accumulate_upstream_token_usage()

        assert isinstance(result, Usage)
        assert result.input_tokens == 300
        assert result.output_tokens == 130
        assert result.total_tokens == 430

    def test_no_llms_upstream(self):
        """No predecessors with token usage: returns None."""
        prompt = _make_vertex_stub("Prompt-1")
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        _wire_graph(output, {"Prompt-1": prompt}, [_make_edge("Prompt-1", "ChatOutput-1")])

        result = output._accumulate_upstream_token_usage()

        assert result is None

    def test_no_predecessors(self):
        """No predecessors at all: returns None."""
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        _wire_graph(output, {}, [])

        result = output._accumulate_upstream_token_usage()

        assert result is None

    def test_diamond_pattern_deduplication(self):
        """Diamond: LLM feeds into two branches that merge at ChatOutput, counted once."""
        llm = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=100, output_tokens=50, total_tokens=150),
        )
        branch_a = _make_vertex_stub("BranchA-1")
        branch_b = _make_vertex_stub("BranchB-1")
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        edges = [
            _make_edge("LLM-1", "BranchA-1"),
            _make_edge("LLM-1", "BranchB-1"),
            _make_edge("BranchA-1", "ChatOutput-1"),
            _make_edge("BranchB-1", "ChatOutput-1"),
        ]
        _wire_graph(output, {"LLM-1": llm, "BranchA-1": branch_a, "BranchB-1": branch_b}, edges)

        result = output._accumulate_upstream_token_usage()

        assert isinstance(result, Usage)
        assert result.input_tokens == 100
        assert result.output_tokens == 50
        assert result.total_tokens == 150

    def test_includes_own_token_usage(self):
        """Output vertex with its own _token_usage includes it in the total."""
        llm = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=100, output_tokens=50, total_tokens=150),
        )
        output = _make_vertex_stub(
            "ChatOutput-1",
            is_output=True,
            own_token_usage=Usage(input_tokens=10, output_tokens=5, total_tokens=15),
        )
        _wire_graph(output, {"LLM-1": llm}, [_make_edge("LLM-1", "ChatOutput-1")])

        result = output._accumulate_upstream_token_usage()

        assert isinstance(result, Usage)
        assert result.input_tokens == 110
        assert result.output_tokens == 55
        assert result.total_tokens == 165

    def test_mixed_predecessors_with_and_without_tokens(self):
        """Only predecessors with token_usage contribute; others are skipped."""
        llm = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=100, output_tokens=50, total_tokens=150),
        )
        prompt = _make_vertex_stub("Prompt-1")
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        edges = [_make_edge("Prompt-1", "LLM-1"), _make_edge("LLM-1", "ChatOutput-1")]
        _wire_graph(output, {"Prompt-1": prompt, "LLM-1": llm}, edges)

        result = output._accumulate_upstream_token_usage()

        assert isinstance(result, Usage)
        assert result.input_tokens == 100
        assert result.output_tokens == 50
        assert result.total_tokens == 150

    def test_predecessor_with_result_but_no_token_usage(self):
        """Predecessor has a result but token_usage is None: skipped."""
        predecessor = _make_vertex_stub("SomeComponent-1")
        predecessor.result = ResultData(token_usage=None)
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        _wire_graph(output, {"SomeComponent-1": predecessor}, [_make_edge("SomeComponent-1", "ChatOutput-1")])

        result = output._accumulate_upstream_token_usage()

        assert result is None

    def test_handles_none_values_in_token_usage(self):
        """Token usage with None values treated as 0."""
        llm = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=None, output_tokens=50, total_tokens=None),
        )
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        _wire_graph(output, {"LLM-1": llm}, [_make_edge("LLM-1", "ChatOutput-1")])

        result = output._accumulate_upstream_token_usage()

        assert isinstance(result, Usage)
        assert result.input_tokens == 0
        assert result.output_tokens == 50
        assert result.total_tokens == 50


class TestExtractTokenUsage:
    """Tests for Vertex._extract_token_usage() dispatch logic."""

    def test_output_vertex_returns_none(self):
        """Output vertex returns None (tokens shown on chat message instead)."""
        llm = _make_vertex_stub(
            "LLM-1",
            token_usage=Usage(input_tokens=100, output_tokens=50, total_tokens=150),
        )
        output = _make_vertex_stub("ChatOutput-1", is_output=True)
        _wire_graph(output, {"LLM-1": llm}, [_make_edge("LLM-1", "ChatOutput-1")])

        result = Vertex._extract_token_usage(output)

        assert result is None

    def test_non_output_vertex_returns_own_usage(self):
        """Non-output vertex returns its own component's _token_usage."""
        expected_usage = Usage(input_tokens=100, output_tokens=50, total_tokens=150)
        llm = _make_vertex_stub(
            "LLM-1",
            is_output=False,
            own_token_usage=expected_usage,
        )

        result = Vertex._extract_token_usage(llm)

        assert isinstance(result, Usage)
        assert result.input_tokens == 100
        assert result.output_tokens == 50
        assert result.total_tokens == 150

    def test_non_output_vertex_without_usage_returns_none(self):
        """Non-output vertex with no _token_usage returns None."""
        vertex = _make_vertex_stub("Prompt-1", is_output=False)

        result = Vertex._extract_token_usage(vertex)

        assert result is None
