"""Unit test suite for Task A52: Local Models, Supervisor, and Dual Coder Pipeline."""
import json
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from antigravity_provider.router.dual_coder_pipeline import (
    DualCoderPipeline,
    DualCoderResult,
    JudgeVerdict,
)
from antigravity_provider.router.local_supervisor import (
    ContextExhaustedError,
    IndivisibleTaskError,
    LocalSupervisor,
    ModelMemoryRecord,
    ServerPropsResult,
    SupervisorOutcome,
    TokenCountResult,
)
from antigravity_provider.router.role_registry import CANONICAL_ROLES, RoleRegistry


# =====================================================================
# 1. Role Registry Tests (P0-4)
# =====================================================================
def test_role_registry_contains_local_supervisor():
    """Verify local-supervisor canonical role is properly defined in RoleRegistry."""
    assert "local-supervisor" in CANONICAL_ROLES
    role = RoleRegistry.get_role("local-supervisor")
    assert role is not None
    assert role.role_id == "local-supervisor"
    assert role.display_name_ru == "Надзиратель локальных моделей"
    assert role.tier == "governance"
    assert "task-splitter" in role.capabilities
    assert "token-counter" in role.capabilities
    assert "memory-tracker" in role.capabilities
    assert role.is_implemented is True


def test_role_registry_local_supervisor_aliases():
    """Verify all aliases resolve to local-supervisor."""
    for alias in ["local-supervisor", "local_supervisor", "надзиратель локальных моделей", "supervisor", "локальный надзиратель"]:
        assert RoleRegistry.resolve_canonical_role(alias) == "local-supervisor"


# =====================================================================
# 2. Local Supervisor Props & Token Counting Tests (P0-5)
# =====================================================================
def test_supervisor_query_props_success():
    supervisor = LocalSupervisor(base_url="http://mock-server:8081")
    fake_props_response = json.dumps({
        "default_generation_settings": {"n_ctx": 65536},
        "total_slots": 1,
        "model_path": "/srv/ai/models/qwen3-coder-30b-a3b/Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf",
    }).encode("utf-8")

    mock_resp = MagicMock()
    mock_resp.read.return_value = fake_props_response
    mock_resp.__enter__.return_value = mock_resp

    with patch("urllib.request.urlopen", return_value=mock_resp):
        res = supervisor.query_server_props()
        assert res.n_ctx == 65536
        assert res.total_slots == 1
        assert "Qwen3-Coder-30B" in res.model_name
        assert res.is_measured is True


def test_supervisor_query_props_offline_fallback():
    supervisor = LocalSupervisor(base_url="http://unreachable-host:8081")
    with patch("urllib.request.urlopen", side_effect=Exception("Connection refused")):
        res = supervisor.query_server_props()
        assert res.n_ctx == 65536
        assert res.is_measured is False


def test_supervisor_count_tokens_exact_via_api():
    supervisor = LocalSupervisor(base_url="http://mock-server:8081")
    fake_tok_response = json.dumps({
        "tokens": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    }).encode("utf-8")

    mock_resp = MagicMock()
    mock_resp.read.return_value = fake_tok_response
    mock_resp.__enter__.return_value = mock_resp

    with patch("urllib.request.urlopen", return_value=mock_resp):
        res = supervisor.count_tokens("def add(a, b): return a + b")
        assert res.tokens_count == 10
        assert res.is_estimated is False
        assert res.method == "tokenize_api"


def test_supervisor_count_tokens_fallback_heuristic():
    supervisor = LocalSupervisor(base_url="http://unreachable-host:8081")
    with patch("urllib.request.urlopen", side_effect=Exception("Timeout")):
        sample_text = "A" * 70
        res = supervisor.count_tokens(sample_text)
        assert res.is_estimated is True
        assert res.tokens_count == int(70 / 3.5)
        assert res.method == "char_heuristic"


# =====================================================================
# 3. Semantic Task Splitting Tests (P0-6)
# =====================================================================
def test_supervisor_split_task_fits_in_single_chunk():
    supervisor = LocalSupervisor()
    with patch.object(supervisor, "count_tokens", return_value=TokenCountResult(tokens_count=500, is_estimated=False, method="mock")):
        chunks = supervisor.split_task_semantically("def small_task(): pass", max_chunk_tokens=2000)
        assert len(chunks) == 1
        assert chunks[0] == "def small_task(): pass"


def test_supervisor_split_task_across_files():
    supervisor = LocalSupervisor()
    task = """--- file: a.py
def func_a():
    pass

--- file: b.py
def func_b():
    pass
"""
    def mock_count(text):
        if "a.py" in text and "b.py" in text:
            return TokenCountResult(tokens_count=250, is_estimated=False, method="mock")
        return TokenCountResult(tokens_count=80, is_estimated=False, method="mock")

    with patch.object(supervisor, "count_tokens", side_effect=mock_count):
        chunks = supervisor.split_task_semantically(task, max_chunk_tokens=150)
        assert len(chunks) == 2
        assert "a.py" in chunks[0]
        assert "b.py" in chunks[1]


def test_supervisor_split_task_across_functions():
    supervisor = LocalSupervisor()
    task = """class MyService:
    def method_one(self):
        print("1")

    def method_two(self):
        print("2")
"""
    def mock_count(text):
        if "method_one" in text and "method_two" in text:
            return TokenCountResult(tokens_count=200, is_estimated=False, method="mock")
        return TokenCountResult(tokens_count=40, is_estimated=False, method="mock")

    with patch.object(supervisor, "count_tokens", side_effect=mock_count):
        chunks = supervisor.split_task_semantically(task, max_chunk_tokens=100)
        assert len(chunks) >= 2


def test_supervisor_indivisible_task_raises_error():
    supervisor = LocalSupervisor()
    indivisible_giant_line = "x = " + ("1+" * 5000)
    with patch.object(supervisor, "count_tokens", return_value=TokenCountResult(tokens_count=10000, is_estimated=False, method="mock")):
        with pytest.raises(IndivisibleTaskError):
            supervisor.split_task_semantically(indivisible_giant_line, max_chunk_tokens=500)


# =====================================================================
# 4. Outcome & Reasoning Exhaustion (A39) Tests (P0-7)
# =====================================================================
def test_supervisor_detect_outcome_success():
    supervisor = LocalSupervisor()
    resp = {
        "choices": [{"message": {"content": "def add(a, b): return a + b"}}],
        "timings": {"predicted_n": 25},
    }
    outcome, desc = supervisor.detect_outcome(resp, elapsed_sec=2.5)
    assert outcome == SupervisorOutcome.SUCCESS
    assert "Успешно" in desc


def test_supervisor_detect_outcome_timeout():
    supervisor = LocalSupervisor()
    outcome, desc = supervisor.detect_outcome(None, error=TimeoutError("Request timed out"), elapsed_sec=185.0)
    assert outcome == SupervisorOutcome.TIMEOUT
    assert "Превышен таймаут" in desc


def test_supervisor_detect_outcome_a39_reasoning_exhaustion():
    supervisor = LocalSupervisor()
    # Case A39: content is empty or whitespace, but predicted_n > 50 or reasoning_content exists
    resp = {
        "choices": [{
            "message": {
                "content": "",
                "reasoning_content": "Let me think about how to write the adder function... wait..."
            },
            "finish_reason": "length"
        }],
        "timings": {"predicted_n": 1500},
    }
    outcome, desc = supervisor.detect_outcome(resp, elapsed_sec=45.0)
    assert outcome == SupervisorOutcome.REASONING_EXHAUSTED
    assert "A39" in desc
    assert "enable_thinking: false" in desc


# =====================================================================
# 5. AI-Memory Integration Tests (P0-8)
# =====================================================================
def test_supervisor_memory_tracking_roundtrip(tmp_path: Path):
    mem_file = tmp_path / "local_models_memory.json"
    supervisor = LocalSupervisor(memory_path=mem_file)

    rec1 = supervisor.record_working_volume(
        gguf_name="Qwen3-Coder-30B-A3B",
        prompt_tokens=16000,
        output_tokens=250,
        outcome=SupervisorOutcome.SUCCESS,
        speed_tps=110.5,
        task_id="T01",
    )
    assert rec1.successful_dispatches == 1
    assert rec1.safe_chunk_tokens == 16000
    assert rec1.avg_generation_tps > 100.0

    # Fetch from memory
    fetched = supervisor.get_model_memory("Qwen3-Coder-30B-A3B")
    assert fetched is not None
    assert fetched.last_working_context == 16000

    # Record failure -> safe chunk reduced
    rec2 = supervisor.record_working_volume(
        gguf_name="Qwen3-Coder-30B-A3B",
        prompt_tokens=48000,
        output_tokens=0,
        outcome=SupervisorOutcome.TIMEOUT,
        speed_tps=0.0,
        task_id="T02",
    )
    assert rec2.failed_dispatches == 1
    assert rec2.safe_chunk_tokens < 48000


# =====================================================================
# 6. Dual Coder Pipeline Tests (P0-9)
# =====================================================================
def test_dual_coder_pipeline_success_winner_a():
    coder_a = MagicMock(return_value={"content": "def solution_a(): return True", "tokens_generated": 40})
    coder_b = MagicMock(return_value={"content": "def solution_b(): return False", "tokens_generated": 45})
    judge = MagicMock(return_value={
        "verdict": "ACCEPTED_A",
        "commentary": "Solution A is correct and handles edge cases.",
        "tokens_consumed": 120,
    })

    pipeline = DualCoderPipeline(coder_a_fn=coder_a, coder_b_fn=coder_b, judge_fn=judge, max_rounds=3, enabled=True)
    result = pipeline.run_pipeline("Write boolean check", judge_model_name="claude-3-5-sonnet")

    assert result.success is True
    assert result.final_verdict == JudgeVerdict.ACCEPTED_A
    assert result.winning_coder == "coder-a"
    assert result.total_rounds == 1
    assert result.total_judge_calls == 1
    assert result.total_judge_tokens == 120
    assert "solution_a" in result.final_code


def test_dual_coder_pipeline_rework_then_winner_b():
    call_count = {"judge": 0}

    def mock_coder_a(prompt, feedback):
        return {"content": f"code_a_round_{call_count['judge']}", "tokens_generated": 30}

    def mock_coder_b(prompt, feedback):
        return {"content": f"code_b_round_{call_count['judge']}", "tokens_generated": 35}

    def mock_judge(prompt, code_a, code_b, model):
        call_count["judge"] += 1
        if call_count["judge"] == 1:
            return {
                "verdict": "REWORK_BOTH",
                "commentary": "Both missing docstrings.",
                "feedback_for_a": "Add docstring to A",
                "feedback_for_b": "Add docstring to B",
                "tokens_consumed": 100,
            }
        return {
            "verdict": "ACCEPTED_B",
            "commentary": "B added perfect docstrings.",
            "tokens_consumed": 110,
        }

    pipeline = DualCoderPipeline(coder_a_fn=mock_coder_a, coder_b_fn=mock_coder_b, judge_fn=mock_judge, max_rounds=3, enabled=True)
    result = pipeline.run_pipeline("Write function with docs")

    assert result.success is True
    assert result.final_verdict == JudgeVerdict.ACCEPTED_B
    assert result.winning_coder == "coder-b"
    assert result.total_rounds == 2
    assert result.total_judge_calls == 2
    assert result.total_judge_tokens == 210


def test_dual_coder_pipeline_stagnation_detection():
    # Both coders return identical unchanged code on round 2
    coder_a = MagicMock(return_value={"content": "unchanged_code_a", "tokens_generated": 20})
    coder_b = MagicMock(return_value={"content": "unchanged_code_b", "tokens_generated": 25})
    judge = MagicMock(return_value={"verdict": "REWORK_BOTH", "commentary": "Still broken", "tokens_consumed": 50})

    pipeline = DualCoderPipeline(coder_a_fn=coder_a, coder_b_fn=coder_b, judge_fn=judge, max_rounds=3, enabled=True)
    result = pipeline.run_pipeline("Fix algorithm")

    assert result.success is False
    assert result.final_verdict == JudgeVerdict.STAGNATION
    assert "Застревание" in result.failure_reason
    assert result.total_rounds == 2


def test_dual_coder_pipeline_rounds_exhausted():
    round_cnt = [0]
    def mock_a(prompt, fb):
        round_cnt[0] += 1
        return {"content": f"varying_code_a_step_{round_cnt[0]}", "tokens_generated": 20}

    def mock_b(prompt, fb):
        return {"content": f"varying_code_b_step_{round_cnt[0]}", "tokens_generated": 25}

    judge = lambda prompt, ca, cb, m: {"verdict": "REWORK_BOTH", "commentary": "Needs more work", "tokens_consumed": 50}

    pipeline = DualCoderPipeline(coder_a_fn=mock_a, coder_b_fn=mock_b, judge_fn=judge, max_rounds=3, enabled=True)
    result = pipeline.run_pipeline("Tough problem")

    assert result.success is False
    assert result.final_verdict == JudgeVerdict.ROUNDS_EXHAUSTED
    assert result.total_rounds == 3
    assert result.total_judge_calls == 3
    assert result.total_judge_tokens == 150