temp/features/steps/context_analysis_agent_coverage_steps.py

"""Step definitions for ContextAnalysisAgent coverage scenarios."""

from __future__ import annotations

import asyncio
import json
import shutil
import tempfile
from collections.abc import Callable
from pathlib import Path
from typing import Any

from behave import given, then, when
from langchain_community.llms import FakeListLLM
from langchain_core.documents import Document

from cleveragents.agents.context_analysis import (
    ContextAnalysisAgent,
    ContextAnalysisState,
)

DEFAULT_LLM_RESPONSES = [
    "Dependencies: ['os', 'sys', 'pathlib']",
    "Relevance: 0.8",
    "Summary: Generated context overview",
] * 20


class RaisingFakeLLM(FakeListLLM):
    """Fake LLM that raises a runtime error on every invocation."""

    def __init__(self, message: str):
        super().__init__(responses=[])
        self._message = message

    def _call(self, prompt: str, **kwargs: Any) -> str:
        raise RuntimeError(self._message)


def _make_state(**overrides: Any) -> ContextAnalysisState:
    state: ContextAnalysisState = {
        "file_paths": [],
        "documents": [],
        "dependencies": {},
        "summary": "",
        "relevance_scores": {},
        "chunks": [],
        "error": None,
    }
    state.update(overrides)
    return state


def _ensure_agent(context: Any, **kwargs: Any) -> ContextAnalysisAgent:
    if hasattr(context, "agent") and context.agent is not None:
        return context.agent

    llm_factory: Callable[[], FakeListLLM] | None = getattr(context, "make_llm", None)
    if llm_factory is not None and "llm" not in kwargs:
        kwargs["llm"] = llm_factory()
    context.agent = ContextAnalysisAgent(**kwargs)
    return context.agent


def _ensure_temp_dir(context: Any) -> Path:
    if not hasattr(context, "temp_dir"):
        context.temp_dir = Path(tempfile.mkdtemp(prefix="context-analysis-"))
    return context.temp_dir


@given("the context analysis agent module is importable")
def step_module_importable(context: Any) -> None:
    ContextAnalysisAgent  # noqa: B018 (import verification)
    ContextAnalysisState  # noqa: B018 (import verification)


@given("I have a mock LLM provider configured for context analysis")
def step_configure_mock_llm(context: Any) -> None:
    def factory() -> FakeListLLM:
        return FakeListLLM(responses=list(DEFAULT_LLM_RESPONSES))

    context.make_llm = factory


@when("I create a ContextAnalysisAgent with default parameters")
def step_create_agent_default(context: Any) -> None:
    context.agent = _ensure_agent(context)


@when(
    "I create a ContextAnalysisAgent with chunk_size {chunk_size:d} and chunk_overlap {chunk_overlap:d}"
)
def step_create_agent_custom_chunks(
    context: Any, chunk_size: int, chunk_overlap: int
) -> None:
    context.agent = _ensure_agent(
        context,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )


@then("the context analysis agent should be initialized successfully")
def step_agent_initialized(context: Any) -> None:
    assert isinstance(context.agent, ContextAnalysisAgent)


@then("the agent should have a chunk_size attribute set to {value:d}")
def step_agent_chunk_size(context: Any, value: int) -> None:
    assert context.agent.chunk_size == value


@then("the agent should have a chunk_overlap attribute set to {value:d}")
def step_agent_chunk_overlap(context: Any, value: int) -> None:
    assert context.agent.chunk_overlap == value


@then("the context analysis agent should have an llm provider configured")
def step_agent_has_llm(context: Any) -> None:
    assert getattr(context.agent, "llm", None) is not None


@given("I have a ContextAnalysisAgent instance")
def step_have_agent_instance(context: Any) -> None:
    context.agent = _ensure_agent(context)


@given('I have a ContextAnalysisAgent instance with an LLM that raises "{message}"')
def step_have_agent_with_raising_llm(context: Any, message: str) -> None:
    context.agent = None
    context.agent = _ensure_agent(context, llm=RaisingFakeLLM(message))


@given(
    "I have a ContextAnalysisAgent instance with chunk_size {chunk_size:d} and chunk_overlap {overlap:d}"
)
def step_have_agent_with_chunks(context: Any, chunk_size: int, overlap: int) -> None:
    context.agent = _ensure_agent(
        context,
        chunk_size=chunk_size,
        chunk_overlap=overlap,
    )


@when("I inspect the workflow graph")
def step_inspect_graph(context: Any) -> None:
    context.graph = context.agent.graph


@then('the graph should contain node "{node_name}"')
def step_graph_contains_node(context: Any, node_name: str) -> None:
    assert node_name in context.graph.nodes


@given('I have a temporary test file named "{filename}" with content:')
def step_create_temp_file_with_content(context: Any, filename: str) -> None:
    temp_dir = _ensure_temp_dir(context)
    file_path = temp_dir / filename
    file_path.parent.mkdir(parents=True, exist_ok=True)
    file_path.write_text(context.text.strip())
    context.last_file_path = file_path


@given('I have a temporary directory named "{dirname}"')
def step_create_temp_directory(context: Any, dirname: str) -> None:
    temp_dir = _ensure_temp_dir(context)
    directory_path = temp_dir / dirname
    directory_path.mkdir(parents=True, exist_ok=True)
    context.last_directory = directory_path


@given("I have temporary test files:")
def step_create_multiple_temp_files(context: Any) -> None:
    temp_dir = _ensure_temp_dir(context)
    file_paths: list[Path] = []
    for row in context.table:
        file_path = temp_dir / row["filename"]
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_text(row["content"])
        file_paths.append(file_path)
    context.temp_files = file_paths


def _resolve_paths(context: Any, relative_paths: list[str]) -> list[str]:
    base = getattr(context, "temp_dir", None)
    if base is None:
        return relative_paths
    return [str(Path(base) / rel) for rel in relative_paths]


@when("I execute the load_files node with file paths:")
def step_execute_load_files(context: Any) -> None:
    file_paths = json.loads(context.text.strip())
    absolute_paths = _resolve_paths(context, file_paths)
    initial_state = _make_state(file_paths=absolute_paths)
    result = context.agent._load_files(initial_state)
    initial_state.update(result)
    context.state = initial_state


@then("the state should contain documents")
def step_state_has_documents(context: Any) -> None:
    assert "documents" in context.state


@then("the documents list should have {count:d} documents")
def step_documents_len(context: Any, count: int) -> None:
    assert len(context.state["documents"]) == count


@then('the first document should contain "{snippet}"')
def step_first_document_contains(context: Any, snippet: str) -> None:
    first_doc = context.state["documents"][0]
    assert snippet in first_doc.page_content


@then("there should be no error")
def step_no_error_present(context: Any) -> None:
    assert context.state.get("error") in (None, "")


@then('the state error should contain "{text}"')
def step_state_error_contains(context: Any, text: str) -> None:
    error = context.state.get("error")
    assert error is not None and text in error


@given("I have a state with loaded documents containing:")
def step_state_with_loaded_documents(context: Any) -> None:
    doc = Document(
        page_content=context.text.strip(),
        metadata={"source": "test.py"},
    )
    context.state = _make_state(documents=[doc])


@given('the state error is "{message}"')
def step_set_state_error(context: Any, message: str) -> None:
    if not hasattr(context, "state") or context.state is None:
        context.state = _make_state()
    context.state["error"] = message


@when("I execute the analyze_dependencies node")
def step_execute_analyze_dependencies(context: Any) -> None:
    state = context.state
    result = context.agent._analyze_dependencies(state)
    state.update(result)
    context.state = state


@then("the state should contain dependencies")
def step_state_contains_dependencies(context: Any) -> None:
    assert "dependencies" in context.state


@then("the dependencies should be a dictionary")
def step_dependencies_is_dict(context: Any) -> None:
    assert isinstance(context.state["dependencies"], dict)


@then('the dependencies for "{source}" should be empty')
def step_dependencies_empty(context: Any, source: str) -> None:
    dependencies = context.state["dependencies"]
    assert dependencies.get(source) == []


@given("I have a state with a document of {size:d} characters")
def step_state_with_document_size(context: Any, size: int) -> None:
    content = "x" * size
    doc = Document(page_content=content, metadata={"source": "large.py"})
    context.state = _make_state(documents=[doc])


@when("I execute the chunk_documents node")
def step_execute_chunk_documents(context: Any) -> None:
    state = context.state
    result = context.agent._chunk_documents(state)
    state.update(result)
    context.state = state


@then("the chunks list should have at least {count:d} chunks")
def step_chunks_at_least(context: Any, count: int) -> None:
    assert len(context.state["chunks"]) >= count


@given("I have a state with chunks from {count:d} different files")
def step_state_with_chunks_from_files(context: Any, count: int) -> None:
    chunks = []
    for idx in range(count):
        chunks.append(
            Document(
                page_content=f"chunk {idx}",
                metadata={"source": f"file{idx}.py"},
            )
        )
    context.state = _make_state(chunks=chunks)


@given('I have a state with duplicate chunks from "{source}"')
def step_state_with_duplicate_chunks(context: Any, source: str) -> None:
    chunks = [
        Document(page_content="chunk 0", metadata={"source": source, "chunk_index": 0}),
        Document(page_content="chunk 1", metadata={"source": source, "chunk_index": 1}),
    ]
    context.state = _make_state(chunks=chunks)


@when("I execute the score_relevance node")
def step_execute_score_relevance(context: Any) -> None:
    state = context.state
    result = context.agent._score_relevance(state)
    state.update(result)
    context.state = state


@then("the relevance_scores should be a dictionary")
def step_scores_is_dict(context: Any) -> None:
    assert isinstance(context.state["relevance_scores"], dict)


@then("the relevance_scores should contain {count:d} entries")
def step_scores_has_entries(context: Any, count: int) -> None:
    assert len(context.state["relevance_scores"]) == count


@then("all scores should be between {min_value:f} and {max_value:f}")
def step_scores_in_range(context: Any, min_value: float, max_value: float) -> None:
    for score in context.state["relevance_scores"].values():
        assert min_value <= score <= max_value


@given("I have a complete analysis state with:")
def step_complete_analysis_state(context: Any) -> None:
    values = {row["field"]: int(row["value"]) for row in context.table}
    doc_count = values.get("documents", 0)
    dep_count = values.get("dependencies", 0)
    score_count = values.get("relevance_scores", 0)

    documents = [
        Document(page_content=f"file {i}", metadata={"source": f"file{i}.py"})
        for i in range(doc_count)
    ]
    dependencies = {
        f"file{i}.py": [
            f"dep{i}_{j}" for j in range(dep_count // max(doc_count, 1) or 1)
        ]
        for i in range(doc_count)
    }
    relevance_scores = {f"file{i}.py": 0.6 + (i * 0.05) for i in range(score_count)}
    context.state = _make_state(
        documents=documents,
        dependencies=dependencies,
        relevance_scores=relevance_scores,
        chunks=documents,
    )


@when("I execute the summarize_context node")
def step_execute_summarize_context(context: Any) -> None:
    state = context.state
    result = context.agent._summarize_context(state)
    state.update(result)
    context.state = state


@when("I execute the summarize_context node with a flaky LLM that fails once")
def step_execute_summarize_with_flaky_llm(context: Any) -> None:
    class FlakySummaryLLM(FakeListLLM):
        def __init__(self) -> None:
            super().__init__(responses=["Summary after retry"])
            object.__setattr__(self, "_call_count", 0)

        @property
        def call_count(self) -> int:
            return getattr(self, "_call_count", 0)

        def _call(self, prompt: str, **kwargs: Any) -> str:
            object.__setattr__(self, "_call_count", self.call_count + 1)
            if self.call_count == 1:
                raise RuntimeError("summary transient failure")
            return super()._call(prompt, **kwargs)

    original_llm = context.agent.llm
    flaky_llm = FlakySummaryLLM()
    context.agent.llm = flaky_llm
    try:
        context.summary_result = context.agent._summarize_context(context.state)
    finally:
        context.agent.llm = original_llm
    context.flaky_summary_calls = flaky_llm.call_count


@then("the summarize_context node should succeed after retry")
def step_summarize_retry_success(context: Any) -> None:
    result = getattr(context, "summary_result", {})
    summary = result.get("summary")
    assert isinstance(summary, str) and summary not in {"", "Context analysis failed"}
    assert not result.get("error")
    assert context.flaky_summary_calls >= 2


@then("the state should contain a summary")
def step_state_has_summary(context: Any) -> None:
    summary = context.state.get("summary", "")
    assert isinstance(summary, str) and summary != ""


@then('the summary should equal "{expected}"')
def step_summary_equals(context: Any, expected: str) -> None:
    assert context.state.get("summary") == expected


@when("I run the complete workflow with file paths:")
def step_run_complete_workflow(context: Any) -> None:
    file_paths = json.loads(context.text.strip())
    absolute_paths = _resolve_paths(context, file_paths)
    initial_state = _make_state(file_paths=absolute_paths)
    config = {"configurable": {"thread_id": "sync-workflow"}}
    context.final_state = context.agent.invoke(initial_state, config=config)


@then("the workflow should complete successfully")
def step_workflow_completed(context: Any) -> None:
    assert context.final_state is not None


@then("the final state should contain documents")
def step_final_state_documents(context: Any) -> None:
    assert len(context.final_state.get("documents", [])) >= 0


@then("the final state should contain dependencies")
def step_final_state_dependencies(context: Any) -> None:
    assert "dependencies" in context.final_state


@then("the final state should contain relevance_scores")
def step_final_state_scores(context: Any) -> None:
    assert "relevance_scores" in context.final_state


@then("the final state should contain a summary")
def step_final_state_summary(context: Any) -> None:
    assert "summary" in context.final_state


@given('I have a temporary test file named "{filename}" with content "{content}"')
def step_create_simple_temp_file(context: Any, filename: str, content: str) -> None:
    temp_dir = _ensure_temp_dir(context)
    file_path = temp_dir / filename
    file_path.parent.mkdir(parents=True, exist_ok=True)
    file_path.write_text(content)
    context.last_file_path = file_path


@when("I run the workflow asynchronously with file paths:")
def step_run_workflow_async(context: Any) -> None:
    file_paths = json.loads(context.text.strip())
    absolute_paths = _resolve_paths(context, file_paths)
    initial_state = _make_state(file_paths=absolute_paths)
    config = {"configurable": {"thread_id": "async-workflow"}}

    async def _run() -> ContextAnalysisState:
        return await context.agent.ainvoke(initial_state, config=config)

    context.final_state = asyncio.run(_run())


@then("the async workflow should complete successfully")
def step_async_workflow_completed(context: Any) -> None:
    assert context.final_state is not None


@then("the final state should contain all expected fields")
def step_async_final_state_fields(context: Any) -> None:
    expected_keys = {
        "file_paths",
        "documents",
        "dependencies",
        "summary",
        "relevance_scores",
        "chunks",
        "error",
    }
    assert expected_keys.issubset(context.final_state.keys())


@when("I stream the workflow asynchronously with file paths:")
def step_stream_workflow_async(context: Any) -> None:
    file_paths = json.loads(context.text.strip())
    absolute_paths = _resolve_paths(context, file_paths)
    initial_state = _make_state(file_paths=absolute_paths)
    config = {"configurable": {"thread_id": "async-stream"}}

    async def _run() -> list[dict[str, Any]]:
        events: list[dict[str, Any]] = []
        async for event in context.agent.astream(initial_state, config=config):
            events.append(event)
        return events

    context.stream_events = asyncio.run(_run())


@when("I stream the workflow with file paths:")
def step_stream_workflow(context: Any) -> None:
    file_paths = json.loads(context.text.strip())
    absolute_paths = _resolve_paths(context, file_paths)
    initial_state = _make_state(file_paths=absolute_paths)
    config = {"configurable": {"thread_id": "stream-workflow"}}
    context.stream_events = list(context.agent.stream(initial_state, config=config))


@then("I should receive multiple state updates")
def step_stream_has_updates(context: Any) -> None:
    assert len(context.stream_events) >= 1


@then("each update should correspond to a node execution")
def step_stream_updates_are_nodes(context: Any) -> None:
    expected_nodes = {
        "load_files",
        "analyze_dependencies",
        "chunk_documents",
        "score_relevance",
        "summarize_context",
    }
    for event in context.stream_events:
        assert isinstance(event, dict)
        assert any(key in expected_nodes for key in event)


@when("I parse dependencies from:")
def step_parse_dependencies(context: Any) -> None:
    agent = _ensure_agent(context)
    context.parsed_dependencies = agent._parse_dependencies(context.text.strip())


@then('the parsed dependencies should include "{module}"')
def step_parsed_dependencies_include(context: Any, module: str) -> None:
    assert any(module == dep or module in dep for dep in context.parsed_dependencies)


@when("I parse relevance scores from hints:")
def step_parse_relevance_hints(context: Any) -> None:
    agent = _ensure_agent(context)
    parsed_scores: list[float] = []
    expected_scores: list[float] = []
    for row in context.table:
        parsed_scores.append(agent._parse_relevance_score(row["hint"]))
        expected_scores.append(float(row["expected"]))
    context.parsed_scores = parsed_scores
    context.expected_scores = expected_scores


@then("the parsed scores should match expected values")
def step_parsed_scores_match_expected(context: Any) -> None:
    assert hasattr(context, "parsed_scores")
    assert hasattr(context, "expected_scores")
    for parsed, expected in zip(
        context.parsed_scores, context.expected_scores, strict=False
    ):
        assert abs(parsed - expected) < 1e-9


def after_scenario(context: Any, _scenario: Any) -> None:
    if hasattr(context, "temp_dir") and context.temp_dir.exists():
        shutil.rmtree(context.temp_dir, ignore_errors=True)
    if hasattr(context, "agent"):
        context.agent = None
    if hasattr(context, "state"):
        context.state = None
    if hasattr(context, "final_state"):
        context.final_state = None
    if hasattr(context, "stream_events"):
        context.stream_events = None
    for attr in (
        "parsed_dependencies",
        "parsed_scores",
        "expected_scores",
        "last_directory",
    ):
        if hasattr(context, attr):
            setattr(context, attr, None)