"""Integration tests for Stage 20: Context Budgeting with ResearchOrchestrator."""

from unittest.mock import MagicMock, patch

import pytest

from nsct.orchestration.context_budget import (
    ContextBudgetConfig,
    ContextBudgetExhaustedError,
    ContextBudgetTracker,
)
from nsct.orchestration.orchestrator import ResearchOrchestrator
from nsct.config import AppSettings, AudioConfig, DatabaseConfig, LLMConfig, VisionConfig


@pytest.fixture
def mock_config():
    """Minimal AppSettings for orchestrator tests."""
    llm = LLMConfig(
        base_url="http://test:8000/v1",
        model="test-model",
        max_concurrency=3,
    )
    vision = VisionConfig(
        base_url="http://test:8001/v1",
        model="test-vision",
    )
    audio = AudioConfig(
        base_url="http://test:8002/v1",
        model="test-audio",
    )
    postgres = DatabaseConfig(url="sqlite+aiosqlite:///")
    return AppSettings(llm=llm, vision=vision, audio=audio, postgres=postgres)


def test_orchestrator_has_context_budget_tracker(mock_config):
    """Orchestrator muss context_budget_tracker als Property haben."""
    orchestrator = ResearchOrchestrator(
        config=mock_config,
        research_id="test-research-id",
        query="test query",
    )
    tracker = orchestrator.context_budget_tracker
    assert isinstance(tracker, ContextBudgetTracker)
    assert tracker.total_usage == 0


def test_orchestrator_custom_context_budget(mock_config):
    """Orchestrator kann ein custom ContextBudgetConfig akzeptieren."""
    custom_config = ContextBudgetConfig(
        planner_max_tokens=16000,
        claim_extraction_max_tokens=24000,
        contradiction_max_tokens=32000,
        synthesis_max_tokens=64000,
    )
    orchestrator = ResearchOrchestrator(
        config=mock_config,
        research_id="test-research-id",
        query="test query",
        context_budget_config=custom_config,
    )
    assert orchestrator._context_budget_config.planner_max_tokens == 16000
    assert orchestrator.context_budget_tracker._config.planner_max_tokens == 16000


def test_orchestrator_track_tokens(mock_config):
    """Orchestrator muss track_tokens für Stages unterstützen."""
    orchestrator = ResearchOrchestrator(
        config=mock_config,
        research_id="test-research-id",
        query="test query",
    )
    # Track planner tokens
    orchestrator._track_context_tokens("planner", 5000)
    assert orchestrator.context_budget_tracker._usage.get("planner", 0) == 5000

    # Track again — accumulation works
    orchestrator._track_context_tokens("planner", 3000)
    assert orchestrator.context_budget_tracker._usage.get("planner", 0) == 8000

    # get_usage returns correct data
    usage = orchestrator.context_budget_tracker.get_usage()
    assert usage["planner"]["usage_tokens"] == 8000.0
    assert usage["planner"]["limit_tokens"] == 12000.0


def test_orchestrator_context_budget_exhaustion(mock_config):
    """Orchestrator wirft ContextBudgetExhaustedError bei Budget-Überschreitung."""
    orchestrator = ResearchOrchestrator(
        config=mock_config,
        research_id="test-research-id",
        query="test query",
    )
    # Track 12000 tokens to reach the limit
    orchestrator._track_context_tokens("planner", 12000)
    # 12000 tokens tracked → at limit → is_exhausted returns True
    assert orchestrator.context_budget_tracker.is_exhausted("planner") is True

    # One more token should raise
    with pytest.raises(ContextBudgetExhaustedError) as exc_info:
        orchestrator._track_context_tokens("planner", 1)

    assert exc_info.value.stage_name == "planner"
    assert exc_info.value.used_tokens == 12000
    assert exc_info.value.limit_tokens == 12000


def test_orchestrator_reset_clears_context_budget(mock_config):
    """Orchestrator reset sollte auch context_budget_tracker zurücksetzen."""
    orchestrator = ResearchOrchestrator(
        config=mock_config,
        research_id="test-research-id",
        query="test query",
    )
    orchestrator._track_context_tokens("planner", 5000)
    orchestrator._track_context_tokens("synthesis", 10000)
    assert orchestrator.context_budget_tracker.total_usage == 15000

    # Simulate reset of context budget tracker
    orchestrator.context_budget_tracker.reset_all()
    assert orchestrator.context_budget_tracker.total_usage == 0


def test_orchestrator_budget_usage_in_report(mock_config):
    """Context budget usage sollte im Report enthalten sein."""
    orchestrator = ResearchOrchestrator(
        config=mock_config,
        research_id="test-research-id",
        query="test query",
    )
    orchestrator._track_context_tokens("planner", 5000)
    orchestrator._track_context_tokens("claim_extraction", 8000)

    usage = orchestrator.context_budget_tracker.get_usage()

    assert "planner" in usage
    assert "claim_extraction" in usage
    assert usage["planner"]["usage_tokens"] == 5000.0
    assert usage["claim_extraction"]["usage_tokens"] == 8000.0
    # synthesis and contradiction should be present but zero
    assert usage["synthesis"]["usage_tokens"] == 0.0
    assert usage["contradiction"]["usage_tokens"] == 0.0