import pytest
from unittest.mock import patch, MagicMock
from langchain_core.messages import AIMessage, HumanMessage


def _run_input_agent_with_llm_response(llm_content, initial_message="I have an idea about bike lanes"):
    """Helper: run input_agent with a mocked LLM returning llm_content."""
    mock_llm = MagicMock()
    mock_llm.invoke.return_value = AIMessage(content=llm_content)

    state = {
        "messages": [HumanMessage(content=initial_message)],
        "location": "New York City, NY",
        "org": "nyc",
        "category": None,
        "collected_details": {},
        "email_cc": [],
        "conversation_complete": False,
        "draft_approved": False,
    }

    with patch("src.agents.input_agent.get_llm", return_value=mock_llm):
        from src.agents.input_agent import input_agent
        return input_agent(state)


def test_extracts_category_idea():
    result = _run_input_agent_with_llm_response(
        "Great idea!\nCATEGORY: idea\nLOCATION: New York City"
    )
    assert result["category"] == "idea"


def test_extracts_category_complaint():
    result = _run_input_agent_with_llm_response(
        "I hear you.\nCATEGORY: complaint\nLOCATION: New York City"
    )
    assert result["category"] == "complaint"


def test_extracts_location_from_tags():
    result = _run_input_agent_with_llm_response(
        "Understood.\nCATEGORY: idea\nLOCATION: Tamil Nadu"
    )
    assert result["location"] == "Tamil Nadu"


def test_category_is_lowercased():
    result = _run_input_agent_with_llm_response(
        "CATEGORY: IDEA\nLOCATION: London"
    )
    assert result["category"] == "idea"


def test_display_text_strips_tag_lines():
    result = _run_input_agent_with_llm_response(
        "Here is my response.\nCATEGORY: idea\nLOCATION: NYC"
    )
    last_ai = result["messages"][-1]
    assert "CATEGORY:" not in last_ai.content
    assert "LOCATION:" not in last_ai.content
    assert "Here is my response." in last_ai.content


def test_category_found_without_location_still_routes():
    result = _run_input_agent_with_llm_response(
        "Your request is clear.\nCATEGORY: complaint"
    )
    assert result["category"] == "complaint"
    assert "location" not in result


def test_sets_conversation_complete_false_on_category():
    result = _run_input_agent_with_llm_response(
        "CATEGORY: idea\nLOCATION: NYC"
    )
    assert result["conversation_complete"] is False


def test_sets_draft_approved_false_on_category():
    result = _run_input_agent_with_llm_response(
        "CATEGORY: idea\nLOCATION: NYC"
    )
    assert result["draft_approved"] is False


def test_no_category_calls_interrupt(mocker):
    mock_llm = MagicMock()
    mock_llm.invoke.return_value = AIMessage(content="Which city is your email about?")
    mock_interrupt = mocker.patch("src.agents.input_agent.interrupt", return_value="London")

    state = {
        "messages": [HumanMessage(content="I have a complaint")],
        "location": "New York City, NY",
        "org": "nyc",
        "category": None,
        "collected_details": {},
        "email_cc": [],
        "conversation_complete": False,
        "draft_approved": False,
    }
    with patch("src.agents.input_agent.get_llm", return_value=mock_llm):
        from src.agents.input_agent import input_agent
        result = input_agent(state)

    mock_interrupt.assert_called_once()
    # User reply should be added as HumanMessage
    human_messages = [m for m in result["messages"] if isinstance(m, HumanMessage)]
    assert any(m.content == "London" for m in human_messages)


def test_location_saved_even_when_no_category(mocker):
    mock_llm = MagicMock()
    mock_llm.invoke.return_value = AIMessage(
        content="Which city?\nLOCATION: Chennai"
    )
    mocker.patch("src.agents.input_agent.interrupt", return_value="complaint about roads")

    state = {
        "messages": [HumanMessage(content="CM Vijay should fix this")],
        "location": "New York City, NY",
        "org": "nyc",
        "category": None,
        "collected_details": {},
        "email_cc": [],
        "conversation_complete": False,
        "draft_approved": False,
    }
    with patch("src.agents.input_agent.get_llm", return_value=mock_llm):
        from src.agents.input_agent import input_agent
        result = input_agent(state)

    assert result.get("location") == "Chennai"