tests/test_agent_flow.py

"""Tests for agents/agent_flow.py — models, routing, readiness, and collection helpers."""

from __future__ import annotations

from unittest.mock import patch, MagicMock

import pytest

from agents.agent_flow import (
    ADVISOR_IDS,
    GENERATOR_IDS,
    AgentDispatchFlow,
    AgentFlowState,
    AgentResponse,
    ChatTurnResponse,
    PreviewData,
    check_readiness,
    collect_responses,
    extract_code,
    route_agents,
)
from agents.design_state import DesignState
from agents.gap_analyzer import GeneratedQuestionCard


# ── Helpers ───────────────────────────────────────────────────────────────────


def _make_response(
    agent_id: str = "design",
    message: str = "Looks good.",
    code: str | None = None,
) -> AgentResponse:
    return AgentResponse.from_agent(agent_id, message, code)


# ── TestAgentResponse ─────────────────────────────────────────────────────────


class TestAgentResponse:
    def test_from_agent_populates_fields(self):
        resp = AgentResponse.from_agent("design", "Hello design world")
        assert resp.agent_id == "design"
        assert resp.agent_name == "Design Agent"
        assert resp.message == "Hello design world"
        assert resp.color == "#7c3aed"
        assert resp.avatar == "DA"
        assert resp.code is None

    def test_from_agent_with_code(self):
        code = "import cadquery as cq\nresult = cq.Workplane('XY').box(10, 10, 10)"
        resp = AgentResponse.from_agent("cad", "Here is the code", code=code)
        assert resp.agent_id == "cad"
        assert resp.code == code

    def test_engineering_agent_fields(self):
        resp = AgentResponse.from_agent("engineering", "Wall needs to be 3 mm thick")
        assert resp.agent_name == "Engineering Agent"
        assert resp.color == "#00b4d8"
        assert resp.avatar == "EA"

    def test_cnc_agent_fields(self):
        resp = AgentResponse.from_agent("cnc", "3-axis is fine here")
        assert resp.agent_name == "CNC Agent"
        assert resp.color == "#00e676"
        assert resp.avatar == "CA"

    def test_cam_agent_fields(self):
        resp = AgentResponse.from_agent("cam", "Roughing then finishing")
        assert resp.agent_name == "CAM Agent"
        assert resp.color == "#ff6b35"
        assert resp.avatar == "CM"

    def test_cad_agent_fields(self):
        resp = AgentResponse.from_agent("cad", "Generated model")
        assert resp.agent_name == "CAD Coder"
        assert resp.color == "#ffab40"
        assert resp.avatar == "CC"

    def test_model_dump_keys(self):
        resp = AgentResponse.from_agent("design", "A message")
        data = resp.model_dump()
        assert set(data.keys()) == {
            "agent_id", "agent_name", "message", "color", "avatar", "code"
        }


# ── TestAgentFlowState ────────────────────────────────────────────────────────


class TestAgentFlowState:
    def test_defaults(self):
        state = AgentFlowState()
        assert state.message == ""
        assert state.context == ""
        assert state.model_str == ""
        assert state.mentions == []
        assert state.is_approved_phase is False
        assert state.active_agent_ids == []
        assert state.knowledge_sources_data == []
        assert state.advisor_responses == []
        assert state.cad_response is None
        assert state.cam_response is None
        assert state.cad_code is None
        assert state.cam_plan is None

    def test_with_inputs(self):
        resp = _make_response("engineering", "Solid choice.")
        state = AgentFlowState(
            message="Make a bracket",
            mentions=["design", "engineering"],
            is_approved_phase=True,
            active_agent_ids=["design", "engineering", "cad"],
            advisor_responses=[resp],
        )
        assert state.message == "Make a bracket"
        assert state.mentions == ["design", "engineering"]
        assert state.is_approved_phase is True
        assert "cad" in state.active_agent_ids
        assert len(state.advisor_responses) == 1

    def test_independent_default_lists(self):
        """Mutable defaults must not be shared between instances."""
        s1 = AgentFlowState()
        s2 = AgentFlowState()
        s1.mentions.append("design")
        assert s2.mentions == []


# ── TestExtractCode ───────────────────────────────────────────────────────────


class TestExtractCode:
    def test_fenced_python_block(self):
        text = "Here is the code:\n```python\nimport cadquery as cq\nresult = cq.Workplane('XY').box(1,1,1)\n```\nDone."
        code = extract_code(text)
        assert code is not None
        assert "import cadquery" in code
        assert "```" not in code

    def test_generic_fence_no_language(self):
        text = "```\ncq.Workplane('XY').box(5,5,5)\n```"
        code = extract_code(text)
        assert code is not None
        assert "cq.Workplane" in code

    def test_unfenced_cq_code_import_marker(self):
        text = "import cadquery as cq\nresult = cq.Workplane('XY').box(2,2,2)"
        code = extract_code(text)
        assert code == text.strip()

    def test_unfenced_cq_dot_marker(self):
        text = "cq.Workplane('XY').box(1,1,1)"
        code = extract_code(text)
        assert code == text.strip()

    def test_unfenced_result_marker(self):
        text = "result = cq.Workplane('XY').box(3,3,3)"
        code = extract_code(text)
        assert code == text.strip()

    def test_plain_text_returns_none(self):
        text = "I need more information about the dimensions before I can proceed."
        code = extract_code(text)
        assert code is None

    def test_empty_string_returns_none(self):
        assert extract_code("") is None

    def test_strips_whitespace_from_fenced_block(self):
        text = "```python\n\n  result = 1\n\n```"
        code = extract_code(text)
        assert code == "result = 1"


# ── TestRouteAgents ───────────────────────────────────────────────────────────


class TestRouteAgents:
    def test_approved_phase_returns_config_agents(self):
        active = route_agents("anything", [], is_approved_phase=True)
        # config.yaml: approved_agents: ["cad", "cnc"]
        assert set(active) == {"cad", "cnc"}

    def test_mentions_override_keywords(self):
        active = route_agents(
            "design shape and engineering material",
            mentions=["cnc"],
            is_approved_phase=False,
        )
        assert active == ["cnc"]

    def test_design_keyword(self):
        active = route_agents("I want a nice shape and design", [], False)
        assert "design" in active

    def test_engineering_keyword(self):
        active = route_agents("what material and wall thickness should I use?", [], False)
        assert "engineering" in active

    def test_cnc_keyword(self):
        active = route_agents("can we machine this on a 5-axis cnc mill?", [], False)
        assert "cnc" in active

    def test_cam_keyword(self):
        active = route_agents("I need a toolpath and gcode output", [], False)
        assert "cam" in active

    def test_default_no_keyword_match(self):
        active = route_agents("hello there", [], False)
        assert "design" in active
        assert "engineering" in active

    def test_max_agents_capped_at_3(self):
        # Message with keywords for design, engineering, cnc, cam — max is 3
        msg = "design the shape, check the material thickness, machine it on a cnc, toolpath and gcode"
        active = route_agents(msg, [], False)
        # CAD-trigger words absent → at most 3 from keyword scoring
        # (cad trigger could add a 4th only if triggered — this message has none)
        non_cad = [a for a in active if a != "cad"]
        assert len(non_cad) <= 3

    def test_cad_trigger_appended(self):
        active = route_agents("generate the model for me", [], False)
        assert "cad" in active

    def test_no_cad_trigger_absent(self):
        active = route_agents("what is the wall thickness?", [], False)
        # "generate" etc. not present — cad might still appear via keyword score
        # but should NOT be added by the trigger path if already at max capacity
        # Just verify the function does not crash and returns a non-empty list
        assert len(active) > 0

    def test_cad_not_duplicated_when_already_in_keywords(self):
        # Even if "cad" is in keyword scores AND a trigger word is present,
        # it should appear only once.
        active = route_agents("generate a model using cad", [], False)
        assert active.count("cad") <= 1

    def test_approved_phase_ignores_mentions(self):
        # approved phase always wins
        active = route_agents("", ["design"], is_approved_phase=True)
        assert set(active) == {"cad", "cnc"}


# ── TestCheckReadiness ────────────────────────────────────────────────────────


class TestCheckReadiness:
    def test_ready_when_clean(self):
        responses = [
            _make_response("design", "Looks good, let's proceed"),
            _make_response("engineering", "Dimensions are confirmed"),
        ]
        result = check_readiness(responses, active_agent_ids=["design", "engineering", "cad"])
        assert result == "READY"

    def test_not_ready_when_flagged(self):
        responses = [
            _make_response("design", "Nice shape"),
            _make_response("cad", "NOT READY: missing wall thickness and material"),
        ]
        result = check_readiness(responses, active_agent_ids=["design", "cad"])
        assert result == "NOT_READY"

    def test_skip_generation_when_no_generators(self):
        responses = [_make_response("design", "All good")]
        result = check_readiness(responses, active_agent_ids=["design", "engineering"])
        assert result == "SKIP_GENERATION"

    def test_not_ready_case_insensitive(self):
        # The check uses .upper() so mixed-case prefix must still trigger NOT_READY
        responses = [_make_response("cad", "not ready: dimensions missing")]
        result = check_readiness(responses, active_agent_ids=["cad"])
        assert result == "NOT_READY"

    def test_not_ready_leading_whitespace_stripped(self):
        responses = [_make_response("cad", "  NOT READY: need more info")]
        result = check_readiness(responses, active_agent_ids=["cad"])
        assert result == "NOT_READY"

    def test_cam_counts_as_generator(self):
        responses = [_make_response("cnc", "Setup looks fine")]
        result = check_readiness(responses, active_agent_ids=["cnc", "cam"])
        assert result == "READY"

    def test_empty_responses_with_generators_is_ready(self):
        result = check_readiness([], active_agent_ids=["cad"])
        assert result == "READY"

    def test_empty_active_ids_skips(self):
        result = check_readiness([], active_agent_ids=[])
        assert result == "SKIP_GENERATION"


# ── TestCollectResponses ──────────────────────────────────────────────────────


class TestCollectResponses:
    def test_merges_all(self):
        advisors = [
            _make_response("design", "Shape confirmed"),
            _make_response("engineering", "Material confirmed"),
        ]
        cad = _make_response("cad", "Generated model", code="result = ...")
        cam = _make_response("cam", "Toolpath ready")
        result = collect_responses(advisors, cad, cam)
        assert len(result) == 4
        assert result[0].agent_id == "design"
        assert result[1].agent_id == "engineering"
        assert result[2].agent_id == "cad"
        assert result[3].agent_id == "cam"

    def test_handles_none_cad(self):
        advisors = [_make_response("design", "OK")]
        cam = _make_response("cam", "Plan ready")
        result = collect_responses(advisors, cad_response=None, cam_response=cam)
        assert len(result) == 2
        assert result[1].agent_id == "cam"

    def test_handles_none_cam(self):
        advisors = [_make_response("engineering", "OK")]
        cad = _make_response("cad", "Model done")
        result = collect_responses(advisors, cad_response=cad, cam_response=None)
        assert len(result) == 2
        assert result[1].agent_id == "cad"

    def test_handles_both_none(self):
        advisors = [_make_response("cnc", "Ready")]
        result = collect_responses(advisors, cad_response=None, cam_response=None)
        assert len(result) == 1
        assert result[0].agent_id == "cnc"

    def test_handles_empty_advisors(self):
        cad = _make_response("cad", "Model done")
        cam = _make_response("cam", "Plan done")
        result = collect_responses([], cad, cam)
        assert len(result) == 2

    def test_does_not_mutate_input_list(self):
        advisors = [_make_response("design", "OK")]
        original_len = len(advisors)
        collect_responses(advisors, _make_response("cad", "x"), None)
        assert len(advisors) == original_len

    def test_all_none_returns_empty(self):
        result = collect_responses([], None, None)
        assert result == []


# ── TestAgentDispatchFlow ────────────────────────────────────────────────────


class TestAgentDispatchFlow:
    def test_no_agents_path(self):
        """Flow with no matching agents routes through HAS_ADVISORS (default fallback)."""
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="xyzzy",
            context="",
            model_str="gemini/gemini-2.5-flash",
        ))
        with patch.object(AgentDispatchFlow, '_run_advisor_crew', return_value=[
            AgentResponse.from_agent("design", "I can help."),
        ]):
            flow.kickoff()
        assert flow.state.active_agent_ids == ["design", "engineering"]
        assert len(flow.state.advisor_responses) == 1

    def test_approved_phase_routes_generators_only(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="build it",
            context="## APPROVED PLAN",
            model_str="gemini/gemini-2.5-flash",
            is_approved_phase=True,
        ))
        with patch.object(AgentDispatchFlow, '_run_advisor_crew', return_value=[
            AgentResponse.from_agent("cnc", "Looks machinable."),
        ]), patch.object(AgentDispatchFlow, '_run_single_agent_crew', return_value="NOT READY: need dims"):
            flow.kickoff()
        assert flow.state.active_agent_ids == ["cad", "cnc"]
        assert flow.state.cad_response is not None
        assert flow.state.cad_response.message.startswith("NOT READY")

    def test_mentions_override(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="check",
            context="",
            model_str="gemini/gemini-2.5-flash",
            mentions=["design"],
        ))
        with patch.object(AgentDispatchFlow, '_run_advisor_crew', return_value=[
            AgentResponse.from_agent("design", "Looks good."),
        ]):
            flow.kickoff()
        assert flow.state.active_agent_ids == ["design"]

    def test_generators_only_path(self):
        """Mentions with only generators routes through GENERATORS_ONLY."""
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="generate it",
            context="",
            model_str="gemini/gemini-2.5-flash",
            mentions=["cad"],
        ))
        with patch.object(AgentDispatchFlow, '_run_single_agent_crew', return_value="NOT READY: need specs"):
            flow.kickoff()
        assert flow.state.active_agent_ids == ["cad"]
        assert flow.state.advisor_responses == []
        assert flow.state.cad_response is not None

    def test_collect_results_from_advisor_and_cad(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="generate a bracket preview",
            context="",
            model_str="gemini/gemini-2.5-flash",
        ))
        with patch.object(AgentDispatchFlow, '_run_advisor_crew', return_value=[
            AgentResponse.from_agent("design", "L-bracket idea."),
        ]), patch.object(AgentDispatchFlow, '_run_single_agent_crew', return_value="```python\nimport cadquery as cq\nresult = cq.Workplane('XY').box(10,10,10)\n```"):
            flow.kickoff()
        results = collect_responses(
            flow.state.advisor_responses,
            flow.state.cad_response,
            flow.state.cam_response,
        )
        assert len(results) >= 2  # advisor + cad
        assert flow.state.cad_code is not None


class TestMemoryHelpers:
    def test_recall_returns_empty_when_no_memory(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="bracket design",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = None
        result = flow._recall_for_agent("design")
        assert result == ""

    def test_recall_formats_matches(self):
        mock_memory = MagicMock()
        mock_match = MagicMock()
        mock_match.record.content = "L-bracket with fillets"
        mock_memory.recall.return_value = [mock_match]

        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="bracket",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = mock_memory
        result = flow._recall_for_agent("design")
        assert "## Relevant context from prior turns" in result
        assert "L-bracket with fillets" in result
        mock_memory.recall.assert_called_once()

    def test_recall_returns_empty_when_no_matches(self):
        mock_memory = MagicMock()
        mock_memory.recall.return_value = []

        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="bracket",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = mock_memory
        result = flow._recall_for_agent("design")
        assert result == ""

    def test_remember_stores_with_scope(self):
        mock_memory = MagicMock()

        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="test",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = mock_memory
        flow._remember_response("engineering", "Use 3mm walls in aluminum.")
        mock_memory.remember.assert_called_once_with(
            "Use 3mm walls in aluminum.",
            scope="/agent/engineering",
        )

    def test_remember_noop_when_no_memory(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="test",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = None
        flow._remember_response("design", "test")  # Should not raise


class TestCollaborationFlag:
    def test_advisors_get_delegation(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="test",
            context="",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = None
        from crewai import LLM
        llm = LLM(model="gemini/gemini-2.5-flash", temperature=0.2)
        agent, task = flow._build_crew_agent("design", llm)
        assert agent.allow_delegation is True

    def test_generators_no_delegation(self):
        flow = AgentDispatchFlow(initial_state=AgentFlowState(
            message="test",
            context="",
            model_str="gemini/gemini-2.5-flash",
        ))
        flow._memory = None
        from crewai import LLM
        llm = LLM(model="gemini/gemini-2.5-flash", temperature=0.2)
        agent, task = flow._build_crew_agent("cad", llm)
        assert agent.allow_delegation is False


# ── TestPreviewData ─────────────────────────────────────────────────────────


class TestPreviewData:
    def test_success_preview(self):
        p = PreviewData(
            success=True,
            part_name="bracket",
            stl_url="/api/models/bracket.stl",
            step_url="/api/models/bracket.step",
            execution={"success": True, "volume_mm3": 1000.0},
            validation={"machinable": True, "axis_recommendation": "3-axis"},
        )
        assert p.success is True
        assert p.part_name == "bracket"

    def test_failure_preview(self):
        p = PreviewData(success=False, error="Execution failed")
        assert p.success is False
        assert p.error == "Execution failed"

    def test_model_dump(self):
        p = PreviewData(success=True, part_name="gear")
        d = p.model_dump()
        assert d["success"] is True
        assert d["cam"] is None
        assert d["gcode_url"] is None


# ── TestChatTurnResponse ────────────────────────────────────────────────────


class TestChatTurnResponse:
    def test_minimal(self):
        r = ChatTurnResponse(design_state=DesignState())
        assert r.responses == []
        assert r.preview is None
        assert r.question_cards == []

    def test_full(self):
        resp = AgentResponse(agent_id="design", agent_name="D", message="hi", color="#fff", avatar="D")
        preview = PreviewData(success=True, part_name="test")
        state = DesignState(material="aluminum")
        card = GeneratedQuestionCard(category="material", question="What material?", responsible_agent="engineering", agent_name="Eng", agent_color="#00e676")
        r = ChatTurnResponse(responses=[resp], preview=preview, design_state=state, question_cards=[card])
        assert len(r.responses) == 1
        assert r.preview.part_name == "test"
        assert r.design_state.material == "aluminum"
        assert len(r.question_cards) == 1

    def test_model_dump_roundtrip(self):
        state = DesignState(part_name="bracket", material="steel")
        r = ChatTurnResponse(design_state=state)
        d = r.model_dump()
        assert d["design_state"]["part_name"] == "bracket"
        assert d["responses"] == []
        assert d["preview"] is None