import time

from tradingagents.agents.utils.agent_states import extract_research_provenance
import tradingagents.graph.setup as graph_setup_module
from tradingagents.graph.setup import GraphSetup


def _setup() -> GraphSetup:
    return GraphSetup(
        quick_thinking_llm=None,
        deep_thinking_llm=None,
        tool_nodes={},
        bull_memory=None,
        bear_memory=None,
        trader_memory=None,
        invest_judge_memory=None,
        portfolio_manager_memory=None,
        conditional_logic=None,
        analyst_node_timeout_secs=0.01,
        research_node_timeout_secs=0.01,
    )


def test_manager_guard_fallback_marks_degraded_synthesis():
    setup = _setup()
    state = {
        "investment_debate_state": {
            "history": "Bull Analyst: case",
            "bull_history": "Bull Analyst: case",
            "bear_history": "",
            "current_response": "Bull Analyst: case",
            "judge_decision": "",
            "count": 1,
            "research_status": "full",
            "research_mode": "debate",
            "timed_out_nodes": [],
            "degraded_reason": None,
            "covered_dimensions": ["bull"],
            "manager_confidence": None,
        }
    }

    result = setup._apply_research_fallback(
        state,
        node_name="Research Manager",
        dimension="manager",
        reason="research_manager_timeout",
        started_at=0.0,
    )

    debate = result["investment_debate_state"]
    assert debate["research_status"] == "degraded"
    assert debate["research_mode"] == "degraded_synthesis"
    assert debate["timed_out_nodes"] == ["Research Manager"]
    assert result["investment_plan"].startswith("Recommendation: HOLD")


def test_bull_guard_success_records_coverage():
    setup = _setup()
    state = {
        "investment_debate_state": {
            "history": "",
            "bull_history": "",
            "bear_history": "",
            "current_response": "",
            "judge_decision": "",
            "count": 0,
            "research_status": "full",
            "research_mode": "debate",
            "timed_out_nodes": [],
            "degraded_reason": None,
            "covered_dimensions": [],
            "manager_confidence": None,
        }
    }
    result = {
        "investment_debate_state": {
            "history": "Bull Analyst: ok",
            "bull_history": "Bull Analyst: ok",
            "bear_history": "",
            "current_response": "Bull Analyst: ok",
            "judge_decision": "",
            "count": 1,
        }
    }

    updated = setup._apply_research_success(state, result, dimension="bull")
    debate = updated["investment_debate_state"]
    assert debate["research_status"] == "full"
    assert debate["research_mode"] == "debate"
    assert debate["covered_dimensions"] == ["bull"]


def test_manager_success_sets_confidence_without_changing_shape():
    setup = _setup()
    state = {
        "investment_debate_state": {
            "history": "Bull Analyst: case\nBear Analyst: counter",
            "bull_history": "Bull Analyst: case",
            "bear_history": "Bear Analyst: counter",
            "current_response": "Bear Analyst: counter",
            "judge_decision": "",
            "count": 2,
            "research_status": "full",
            "research_mode": "debate",
            "timed_out_nodes": [],
            "degraded_reason": None,
            "covered_dimensions": ["bull", "bear"],
            "manager_confidence": None,
        }
    }
    result = {
        "investment_debate_state": {
            "history": "Bull Analyst: case\nBear Analyst: counter",
            "bull_history": "Bull Analyst: case",
            "bear_history": "Bear Analyst: counter",
            "current_response": "Recommendation: BUY",
            "judge_decision": "Recommendation: BUY",
            "count": 2,
        },
        "investment_plan": "Recommendation: BUY",
    }

    updated = setup._apply_research_success(state, result, dimension="manager")
    debate = updated["investment_debate_state"]
    assert updated["investment_plan"] == "Recommendation: BUY"
    assert debate["judge_decision"] == "Recommendation: BUY"
    assert debate["research_status"] == "full"
    assert debate["research_mode"] == "debate"
    assert debate["covered_dimensions"] == ["bull", "bear", "manager"]
    assert debate["manager_confidence"] == 1.0


def test_bear_guard_exception_returns_degraded_argument(monkeypatch):
    def broken_bear(_llm, _memory):
        def node(_state):
            raise ConnectionError("downstream unavailable")

        return node

    monkeypatch.setattr(graph_setup_module, "create_bear_researcher", broken_bear)
    setup = _setup()
    wrapped = setup._guard_research_node("Bear Researcher", None, None)
    state = {
        "investment_debate_state": {
            "history": "Bull Analyst: case",
            "bull_history": "Bull Analyst: case",
            "bear_history": "",
            "current_response": "Bull Analyst: case",
            "judge_decision": "",
            "count": 1,
            "research_status": "full",
            "research_mode": "debate",
            "timed_out_nodes": [],
            "degraded_reason": None,
            "covered_dimensions": ["bull"],
            "manager_confidence": None,
        }
    }

    result = wrapped(state)

    debate = result["investment_debate_state"]
    assert debate["research_status"] == "degraded"
    assert debate["research_mode"] == "degraded_synthesis"
    assert debate["degraded_reason"] == "bear_researcher_connectionerror"
    assert debate["timed_out_nodes"] == []
    assert debate["count"] == 2
    assert debate["current_response"].startswith(
        "Bear Analyst: [DEGRADED] Bear Researcher unavailable (bear_researcher_connectionerror)."
    )
    assert debate["history"].startswith("Bull Analyst: case\nBear Analyst: [DEGRADED]")
    assert debate["bear_history"].startswith("\nBear Analyst: [DEGRADED]")


def test_guard_timeout_returns_without_waiting_for_node_completion(monkeypatch):
    def slow_bull(_llm, _memory):
        def node(_state):
            time.sleep(0.2)
            return {"investment_debate_state": {"history": "", "bull_history": "", "bear_history": "", "current_response": "", "judge_decision": "", "count": 1}}
        return node

    monkeypatch.setattr(graph_setup_module, "create_bull_researcher", slow_bull)
    setup = _setup()
    wrapped = setup._guard_research_node("Bull Researcher", None, None)
    state = {
        "investment_debate_state": {
            "history": "",
            "bull_history": "",
            "bear_history": "",
            "current_response": "",
            "judge_decision": "",
            "count": 0,
            "research_status": "full",
            "research_mode": "debate",
            "timed_out_nodes": [],
            "degraded_reason": None,
            "covered_dimensions": [],
            "manager_confidence": None,
        }
    }

    started = time.monotonic()
    result = wrapped(state)
    elapsed = time.monotonic() - started

    assert elapsed < 0.1
    debate = result["investment_debate_state"]
    assert debate["research_status"] == "degraded"
    assert debate["research_mode"] == "degraded_synthesis"
    assert debate["timed_out_nodes"] == ["Bull Researcher"]


def test_analyst_guard_timeout_returns_degraded_report_quickly():
    setup = _setup()

    def slow_node(_state):
        time.sleep(0.2)
        return {"messages": [], "market_report": "ok"}

    wrapped = setup._guard_analyst_node(
        "Market Analyst",
        slow_node,
        report_field="market_report",
    )

    started = time.monotonic()
    result = wrapped({"messages": []})
    elapsed = time.monotonic() - started

    assert elapsed < 0.1
    assert result["market_report"].startswith("[DEGRADED] Market Analyst unavailable")
    assert result["messages"][0].content.startswith("[DEGRADED] Market Analyst unavailable")


def test_extract_research_provenance_returns_subset():
    payload = extract_research_provenance(
        {
            "research_status": "degraded",
            "research_mode": "degraded_synthesis",
            "timed_out_nodes": ["Bull Researcher"],
            "degraded_reason": "bull_researcher_timeout",
            "covered_dimensions": ["market", "bull"],
            "manager_confidence": 0.0,
            "history": "ignored",
        }
    )

    assert payload == {
        "research_status": "degraded",
        "research_mode": "degraded_synthesis",
        "timed_out_nodes": ["Bull Researcher"],
        "degraded_reason": "bull_researcher_timeout",
        "covered_dimensions": ["market", "bull"],
        "manager_confidence": 0.0,
    }


def test_extract_research_provenance_ignores_non_mapping():
    assert extract_research_provenance(None) is None
    assert extract_research_provenance("bad") is None