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 pathlib import Path
from typing import Any, Callable

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

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


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


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


@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 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 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, "")


@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])


@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)


@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)


@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


@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 != ""


@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 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.keys())


@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)


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