"""Stage 4 (R3-Q) P0-4: atomic seat-task initialization and safe recovery.

The old initializer created tasks one by one and short-circuited with
``if existing: return existing`` — a crash after the first create permanently
dropped the remaining required seats on recovery, and the quality gate could
then pass on an incomplete task set.

These tests prove:

- ``get_or_create_tasks`` creates the FULL expected set in one transaction;
- after a simulated partial initialization (only the first task exists),
  recovery fills the missing tasks and never touches existing rows;
- a frozen contract / lifecycle progress is NEVER overwritten;
- two concurrent recovery calls converge on one complete set without raising
  unique-key violations;
- ``QualitySeatTasksExecutor._load_or_create_tasks`` derives the expected set
  from the dispatch assignments and fills gaps;
- the quality gate blocks when an expected required task is missing.
"""

from __future__ import annotations

import asyncio
from datetime import datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from app.maf_roundtable.workflow_executors import (
    QualityGateExecutor,
    QualitySeatTasksExecutor,
    RoundtableState,
)
from deerflow.persistence.maf_roundtable import make_maf_roundtable_store
from deerflow.persistence.maf_roundtable.model import (
    Base,
    MafRoundtableRunRow,
    MafRoundtableSessionRow,
)

CONTRACT_A = {
    "objective": "原始合同目标 A",
    "expected_deliverable": "结论+证据",
    "acceptance_criteria": ["证据"],
    "required": True,
    "max_rework_attempts": 2,
    "fallback_policy": "reassign_then_gap",
}
CONTRACT_B = {
    "objective": "试图覆盖的合同 B",
    "expected_deliverable": "别的",
    "acceptance_criteria": [],
    "required": True,
    "max_rework_attempts": 0,
    "fallback_policy": "controlled_gap",
}


@pytest_asyncio.fixture
async def task_store(tmp_path):
    engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'tasks.db'}", echo=False)
    try:
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        sf = async_sessionmaker(engine, expire_on_commit=False)
        store = make_maf_roundtable_store(sf)
        now = datetime(2026, 1, 1)
        async with sf() as session:
            session.add(MafRoundtableSessionRow(
                id="t-sess", user_id="u1", title="t", task_text="t",
                mode="structured", status="pending", write_policy="auto",
                manager_agent_id="manager-1", max_rounds=8,
                created_at=now, updated_at=now,
            ))
            session.add(MafRoundtableRunRow(
                id="t-run", session_id="t-sess", status="pending",
                created_at=now, updated_at=now,
            ))
            await session.commit()
        yield store
    finally:
        await engine.dispose()


def _expected_specs() -> list[dict]:
    return [
        {"task_key": "01:seat-a", "agent_id": "seat-a", "contract": dict(CONTRACT_A), "required": True},
        {"task_key": "02:seat-b", "agent_id": "seat-b", "contract": dict(CONTRACT_A), "required": True},
        {"task_key": "03:seat-c", "agent_id": "seat-c", "contract": dict(CONTRACT_A), "required": True},
    ]


# ── repository: atomic full-set creation ────────────────────────────────────


@pytest.mark.asyncio
async def test_get_or_create_tasks_creates_full_set(task_store):
    tasks = await task_store.seat_tasks.get_or_create_tasks(
        session_id="t-sess", run_id="t-run", expected=_expected_specs(),
    )
    assert [t["task_key"] for t in tasks] == ["01:seat-a", "02:seat-b", "03:seat-c"]
    assert all(t["status"] == "assigned" for t in tasks)
    assert all(t["contract"]["objective"] == "原始合同目标 A" for t in tasks)


@pytest.mark.asyncio
async def test_get_or_create_tasks_recovers_partial_initialization(task_store):
    # Simulate the old crash mode: only the FIRST task survived initialization.
    await task_store.seat_tasks.create_task(
        session_id="t-sess", run_id="t-run", task_key="01:seat-a",
        agent_id="seat-a", contract=dict(CONTRACT_A), required=True,
    )

    tasks = await task_store.seat_tasks.get_or_create_tasks(
        session_id="t-sess", run_id="t-run", expected=_expected_specs(),
    )
    # The complete set comes back — the missing seats are filled in.
    assert [t["task_key"] for t in tasks] == ["01:seat-a", "02:seat-b", "03:seat-c"]
    by_key = {t["task_key"]: t for t in tasks}
    # The pre-existing row kept its identity and frozen contract.
    assert by_key["01:seat-a"]["contract"]["objective"] == "原始合同目标 A"
    assert by_key["01:seat-a"]["status"] == "assigned"


@pytest.mark.asyncio
async def test_get_or_create_tasks_never_overwrites_frozen_contract(task_store):
    existing = await task_store.seat_tasks.create_task(
        session_id="t-sess", run_id="t-run", task_key="01:seat-a",
        agent_id="seat-c", contract=dict(CONTRACT_A), required=True,
    )
    # Lifecycle progressed (e.g. before a restart): reassignment + attempts.
    await task_store.seat_tasks.update_task(existing["id"], {
        "status": "accepted",
        "attempt_count": 3,
        "rework_count": 2,
        "reassign_count": 1,
    })

    specs = _expected_specs()
    specs[0]["contract"] = dict(CONTRACT_B)  # recovery must NOT apply this
    tasks = await task_store.seat_tasks.get_or_create_tasks(
        session_id="t-sess", run_id="t-run", expected=specs,
    )
    by_key = {t["task_key"]: t for t in tasks}
    kept = by_key["01:seat-a"]
    assert kept["id"] == existing["id"]
    assert kept["contract"]["objective"] == "原始合同目标 A"
    assert kept["agent_id"] == "seat-c"
    assert kept["status"] == "accepted"
    assert kept["attempt_count"] == 3
    assert kept["rework_count"] == 2
    assert kept["reassign_count"] == 1


@pytest.mark.asyncio
async def test_get_or_create_tasks_concurrent_recovery(task_store):
    # Two recoveries race on the same empty run: both must return the complete
    # set, neither may raise a unique-key violation, and exactly one row per
    # task_key may exist afterwards.
    results = await asyncio.gather(
        task_store.seat_tasks.get_or_create_tasks(
            session_id="t-sess", run_id="t-run", expected=_expected_specs(),
        ),
        task_store.seat_tasks.get_or_create_tasks(
            session_id="t-sess", run_id="t-run", expected=_expected_specs(),
        ),
    )
    for tasks in results:
        assert [t["task_key"] for t in tasks] == ["01:seat-a", "02:seat-b", "03:seat-c"]

    final = await task_store.seat_tasks.list_by_run("t-run")
    keys = [t["task_key"] for t in final]
    assert sorted(keys) == ["01:seat-a", "02:seat-b", "03:seat-c"]
    assert len(keys) == len(set(keys)), "no duplicate tasks"


# ── executor: expected set from dispatch assignments ────────────────────────


def _make_state(**overrides) -> RoundtableState:
    kwargs = {
        "session_id": "t-sess",
        "run_id": "t-run",
        "task_text": "分析市场",
        "mode": "structured",
        "model_name": None,
        "reasoning_effort": None,
        "max_rounds": 8,
        "write_policy": "auto",
        "manager_agent_id": "manager-1",
        "seat_agent_ids": ["seat-a", "seat-b"],
        "thread_map": {},
        "capability_snapshot": None,
        "seat_timeout_seconds": 5,
        "schema_repair_enabled": True,
        "intent_type": "parallel_consultation",
        "dispatch_plan": {
            "execution_mode": "parallel",
            "assignments": [
                {"agent_id": "seat-a", "task_goal": "目标一", "delivery_requirements": ["证据"]},
                {"agent_id": "seat-b", "task_goal": "目标二", "delivery_requirements": ["建议"]},
            ],
        },
        "dispatch_contract_enabled": True,
        "quality_gate_enabled": True,
    }
    kwargs.update(overrides)
    return RoundtableState(**kwargs)


@pytest.mark.asyncio
async def test_executor_load_or_create_fills_missing_tasks(task_store):
    # Partial init: only the first assignment's task exists.
    await task_store.seat_tasks.create_task(
        session_id="t-sess", run_id="t-run", task_key="01:seat-a",
        agent_id="seat-a", contract=dict(CONTRACT_A), required=True,
    )
    state = _make_state()
    adapter = SimpleNamespace(_service=SimpleNamespace(store=task_store))
    executor = QualitySeatTasksExecutor(adapter, state)

    tasks = await executor._load_or_create_tasks()
    assert [t["task_key"] for t in tasks] == ["01:seat-a", "02:seat-b"]
    by_key = {t["task_key"]: t for t in tasks}
    # Existing contract untouched; missing task created from the assignment.
    assert by_key["01:seat-a"]["contract"]["objective"] == "原始合同目标 A"
    assert by_key["02:seat-b"]["contract"]["objective"] == "目标二"
    assert by_key["02:seat-b"]["status"] == "assigned"

    # Idempotent second recovery.
    again = await executor._load_or_create_tasks()
    assert [t["id"] for t in again] == [t["id"] for t in tasks]


@pytest.mark.asyncio
async def test_quality_gate_blocks_missing_required_task(task_store):
    state = _make_state()
    # Only ONE of the two expected required tasks is represented.
    state.seat_task_statuses = {
        "01:seat-a": {
            "task_key": "01:seat-a", "status": "accepted",
            "agent_id": "seat-a", "required": True,
        },
    }
    adapter = SimpleNamespace(
        _service=SimpleNamespace(store=SimpleNamespace(seat_tasks=SimpleNamespace())),
        _emit_critical=AsyncMock(),
    )
    gate = QualityGateExecutor(adapter, state)
    ctx = SimpleNamespace(send_message=AsyncMock(), yield_output=AsyncMock())
    with pytest.raises(RuntimeError, match="quality gate blocked"):
        await gate.handle({}, ctx)
    assert state.status != "completed"