import json
from unittest.mock import MagicMock

from agent_factories import (
    make_agent_orchestrator,
    make_final_report_response,
    make_text_response,
    make_tool_use_response,
)
from spark_advisor_analyzer.agent.tools import AgentToolName
from spark_advisor_models.config import AiSettings
from spark_advisor_models.model import Severity
from spark_advisor_models.testing import make_job

DEFAULT_MAX_ITERATIONS = AiSettings().max_agent_iterations


class TestAgentOrchestrator:
    def test_two_turn_flow(self) -> None:
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [
            make_tool_use_response(AgentToolName.GET_JOB_OVERVIEW, {}),
            make_final_report_response(),
        ]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert result.ai_report.severity == Severity.WARNING
        assert len(result.ai_report.recommendations) == 1
        assert mock_client.create_message.call_count == 2

    def test_three_turn_with_rules_and_stage(self) -> None:
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [
            make_tool_use_response(AgentToolName.RUN_RULES_ENGINE, {}, "call_1"),
            make_tool_use_response(AgentToolName.GET_STAGE_DETAILS, {"stage_id": 0}, "call_2"),
            make_final_report_response("call_3"),
        ]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert mock_client.create_message.call_count == 3
        assert len(result.rule_results) > 0

    def test_text_only_response_gets_nudge(self) -> None:
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [
            make_text_response(),
            make_final_report_response(),
        ]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert mock_client.create_message.call_count == 2

    def test_max_iterations_forces_report(self) -> None:
        overview_resp = make_tool_use_response(AgentToolName.GET_JOB_OVERVIEW, {})
        mock_client = MagicMock()
        mock_client.create_message.side_effect = (
            [overview_resp] * DEFAULT_MAX_ITERATIONS + [make_final_report_response()]
        )

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert mock_client.create_message.call_count == DEFAULT_MAX_ITERATIONS + 1

    def test_tool_error_sent_back_to_llm(self) -> None:
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [
            make_tool_use_response(AgentToolName.GET_STAGE_DETAILS, {"stage_id": 999}),
            make_final_report_response(),
        ]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        second_call_messages = mock_client.create_message.call_args_list[1].kwargs["messages"]
        last_msg = second_call_messages[-1]
        assert "error" in str(last_msg)

    def test_submit_final_report_terminates_immediately(self) -> None:
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [make_final_report_response()]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert mock_client.create_message.call_count == 1

    def test_rules_only_fallback_on_force_failure(self) -> None:
        overview_resp = make_tool_use_response(AgentToolName.GET_JOB_OVERVIEW, {})
        empty_response = make_text_response()

        mock_client = MagicMock()
        mock_client.create_message.side_effect = [overview_resp] * DEFAULT_MAX_ITERATIONS + [empty_response]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is None
        assert len(result.rule_results) > 0

    def test_suggested_config_extracted(self) -> None:
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [make_final_report_response()]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert "spark.sql.adaptive.skewJoin.enabled" in result.ai_report.suggested_config

    def test_multiple_tools_in_single_response(self) -> None:
        tool_call_1 = MagicMock()
        tool_call_1.id = "t1"
        tool_call_1.function.name = AgentToolName.GET_JOB_OVERVIEW
        tool_call_1.function.arguments = json.dumps({})

        tool_call_2 = MagicMock()
        tool_call_2.id = "t2"
        tool_call_2.function.name = AgentToolName.RUN_RULES_ENGINE
        tool_call_2.function.arguments = json.dumps({})

        message = MagicMock()
        message.content = "I'll analyze the job overview and run rules."
        message.tool_calls = [tool_call_1, tool_call_2]

        multi_tool_response = MagicMock()
        multi_tool_response.choices = [MagicMock()]
        multi_tool_response.choices[0].message = message

        mock_client = MagicMock()
        mock_client.create_message.side_effect = [multi_tool_response, make_final_report_response()]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert mock_client.create_message.call_count == 2

    def test_malformed_final_report_triggers_fallback(self) -> None:
        malformed = make_tool_use_response(
            AgentToolName.SUBMIT_FINAL_REPORT,
            {"summary": "ok"},
        )
        mock_client = MagicMock()
        mock_client.create_message.side_effect = [malformed, make_final_report_response()]

        orch = make_agent_orchestrator(mock_client)
        result = orch.run(make_job())

        assert result.ai_report is not None
        assert mock_client.create_message.call_count == 2
