import json
from typing import Any
from unittest.mock import MagicMock

from spark_advisor_analyzer.agent.context import AgentContext
from spark_advisor_analyzer.agent.orchestrator import AgentOrchestrator
from spark_advisor_analyzer.agent.tools import AgentToolName
from spark_advisor_models.config import AiSettings, Thresholds
from spark_advisor_models.testing import make_job
from spark_advisor_rules import StaticAnalysisService, rules_for_threshold


def make_agent_context(**overrides: object) -> AgentContext:
    return AgentContext(job=make_job(**overrides))


def make_default_static() -> StaticAnalysisService:
    return StaticAnalysisService(rules_for_threshold(Thresholds()))


def make_tool_use_response(tool_name: str, tool_input: dict[str, Any], tool_id: str = "call_1") -> MagicMock:
    tool_call = MagicMock()
    tool_call.id = tool_id
    tool_call.function.name = tool_name
    tool_call.function.arguments = json.dumps(tool_input)

    message = MagicMock()
    message.content = "I'll use this tool to analyze the job."
    message.tool_calls = [tool_call]

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

    return response


def make_final_report_response(tool_id: str = "call_final") -> MagicMock:
    return make_tool_use_response(
        AgentToolName.SUBMIT_FINAL_REPORT,
        {
            "summary": "Job has minor issues.",
            "severity": "warning",
            "recommendations": [
                {
                    "priority": 1,
                    "title": "Enable AQE",
                    "parameter": "spark.sql.adaptive.skewJoin.enabled",
                    "current_value": "false",
                    "recommended_value": "true",
                    "explanation": "Fixes skew automatically.",
                    "estimated_impact": "~40% faster Stage 0",
                    "risk": "Minor overhead.",
                }
            ],
            "causal_chain": "",
        },
        tool_id,
    )


def make_text_response() -> MagicMock:
    message = MagicMock()
    message.content = "Let me analyze this job..."
    message.tool_calls = None

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

    return response


def make_agent_orchestrator(mock_client: MagicMock) -> AgentOrchestrator:
    return AgentOrchestrator(
        client=mock_client,
        static_analysis=StaticAnalysisService(rules_for_threshold(Thresholds())),
        ai_settings=AiSettings(),
    )
