"""Cohort assist aggregation tests (TASK-10-03, D-062, D-063, REQ-NFR-ASSIST-04).

Tests the _aggregate_assist branch in server/cohort/aggregator.py with a
mocked PgStore (no Postgres required). Verifies:
  - _aggregate_assist() upserts the 5 core assist metrics + p95 latency.
  - k-anon suppression: <10 distinct learners → suppressed.
  - Idempotent upsert: same session_outcome twice → same aggregate.
  - assist_guardrail_block_rate = blocks / turns.
  - The practice branch (_aggregate_practice) is unchanged (backward compat).
  - The dashboard endpoints return assist rows (cohort + failure-patterns).

D-063 (binding): assist does NOT update mastery. The _aggregate_assist branch
computes NO mastery metrics (no rubric scores, no gate_open_rate).
"""

from __future__ import annotations

import datetime as _dt
from unittest.mock import AsyncMock, MagicMock

import pytest

from server.cohort.aggregator import (
    K_ANON_THRESHOLD,
    _aggregate_assist,
    _aggregate_practice,
    aggregate_session,
)
from server.cohort.hook import on_session_end


def _mock_pg_store():
    store = MagicMock()
    store.upsert_cohort_aggregate = AsyncMock()
    return store


def _assist_outcome(
    learner_ref: str,
    path: str = "customer_service",
    turn_count: int = 20,
    blocks: int = 2,
    p95_latency_ms: float | None = 580.0,
    cost_cents: int = 20,
    outcome: str = "completed",
) -> dict:
    return {
        "learner_ref": learner_ref,
        "path": path,
        "scenario_id": f"assist:refund",
        "outcome": outcome,
        "session_type": "assist",
        "rubric_scores": [],  # D-063: no rubric scores for assist
        "failure_mode": None,
        "branch_path": [],
        "assist_turn_count": turn_count,
        "guardrail_blocks": blocks,
        "assist_p95_latency_ms": p95_latency_ms,
        "assist_p50_latency_ms": 500.0,
        "assist_p99_latency_ms": 620.0,
        "assist_within_pilot": True,
        "assist_cost_cents": cost_cents,
        "timestamp": _dt.datetime.now(_dt.timezone.utc).isoformat(),
    }


def _practice_outcome(learner_ref: str, path: str = "customer_service") -> dict:
    return {
        "learner_ref": learner_ref,
        "path": path,
        "scenario_id": f"{path}_v01",
        "outcome": "pass",
        "session_type": "practice",
        "rubric_scores": [{"criterion_id": "empathy", "score": 4.0}],
        "failure_mode": None,
        "branch_path": ["accept"],
        "timestamp": _dt.datetime.now(_dt.timezone.utc).isoformat(),
    }


# ── Assist metrics upserted ─────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_aggregate_assist_upserts_5_core_metrics():
    """_aggregate_assist upserts the 5 core assist metrics (REQ-NFR-ASSIST-04)."""
    store = _mock_pg_store()
    await _aggregate_assist(store, _assist_outcome("learner-1"))
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    assert "assist_shifts_count" in metrics
    assert "assist_active_learners_count" in metrics
    assert "assist_turns_count" in metrics
    assert "assist_avg_turns_per_shift" in metrics
    assert "assist_guardrail_block_rate" in metrics


@pytest.mark.asyncio
async def test_aggregate_assist_upserts_p95_latency():
    """_aggregate_assist upserts assist_p95_latency_ms (D-072, TASK-09-01)."""
    store = _mock_pg_store()
    await _aggregate_assist(store, _assist_outcome("learner-1", p95_latency_ms=580.0))
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    assert "assist_p95_latency_ms" in metrics


@pytest.mark.asyncio
async def test_aggregate_assist_upserts_avg_cost():
    """_aggregate_assist upserts assist_avg_cost_per_shift (TASK-11-01)."""
    store = _mock_pg_store()
    await _aggregate_assist(store, _assist_outcome("learner-1", cost_cents=25))
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    assert "assist_avg_cost_per_shift" in metrics


@pytest.mark.asyncio
async def test_aggregate_assist_no_mastery_metrics():
    """D-063: _aggregate_assist computes NO mastery metrics."""
    store = _mock_pg_store()
    await _aggregate_assist(store, _assist_outcome("learner-1"))
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    # No mastery metrics should be present.
    assert "gate_open_rate" not in metrics
    assert "median_mastery_score" not in metrics
    assert not any(m.startswith("rubric_criterion_mean:") for m in metrics)
    # No practice metrics either (assist is a separate branch).
    assert "sessions_count" not in metrics


# ── k-anonymity suppression ──────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_assist_9_learners_suppressed():
    """<10 distinct learners → all assist cells suppressed."""
    store = _mock_pg_store()
    for i in range(9):
        await _aggregate_assist(store, _assist_outcome(f"learner-{i}"))
    suppressed = [c for c in store.upsert_cohort_aggregate.call_args_list if c.args[6] is True]
    non_suppressed = [c for c in store.upsert_cohort_aggregate.call_args_list if c.args[6] is False]
    assert suppressed, "assist cells should be suppressed with <10 learners"
    assert not non_suppressed, "no assist cell should be non-suppressed with 9 learners"


@pytest.mark.asyncio
async def test_assist_10_learners_not_suppressed():
    """≥10 distinct learners → assist cells not suppressed."""
    store = _mock_pg_store()
    for i in range(10):
        await _aggregate_assist(store, _assist_outcome(f"learner-{i}"))
    non_suppressed = [c for c in store.upsert_cohort_aggregate.call_args_list if c.args[6] is False]
    assert non_suppressed, "assist cells should NOT be suppressed at 10 learners"
    for c in non_suppressed:
        assert c.args[4] is not None, "non-suppressed cell value must not be None"


# ── Idempotent upsert ───────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_assist_idempotent_same_outcome_twice():
    """Re-running with the same outcome produces consistent upserts (idempotent)."""
    store = _mock_pg_store()
    outcome = _assist_outcome("learner-x")
    await _aggregate_assist(store, outcome)
    first_call_count = store.upsert_cohort_aggregate.call_count
    await _aggregate_assist(store, outcome)
    second_call_count = store.upsert_cohort_aggregate.call_count
    # Both runs produce upsert calls (the DB ON CONFLICT makes them idempotent).
    assert second_call_count >= first_call_count
    assert store.upsert_cohort_aggregate.called


# ── assist_guardrail_block_rate = blocks / turns ─────────────────────────────


@pytest.mark.asyncio
async def test_assist_guardrail_block_rate_computed():
    """assist_guardrail_block_rate = blocks / turns (safety signal)."""
    store = _mock_pg_store()
    # 10 learners so the cell is not suppressed (we can read the value).
    for i in range(10):
        await _aggregate_assist(store, _assist_outcome(f"learner-{i}", turn_count=20, blocks=2))
    block_rate_cells = [
        c for c in store.upsert_cohort_aggregate.call_args_list
        if c.args[1] == "assist_guardrail_block_rate" and c.args[6] is False
    ]
    assert block_rate_cells, "should have a non-suppressed assist_guardrail_block_rate cell"
    # The running mean of per-shift block rates (2/20 = 0.1) → ~0.1.
    rate = block_rate_cells[-1].args[4]
    assert rate is not None
    assert 0.05 <= rate <= 0.15  # ~0.1 with running-mean drift


@pytest.mark.asyncio
async def test_assist_zero_turns_block_rate_is_zero():
    """0 turns → block_rate = 0.0 (no division by zero)."""
    store = _mock_pg_store()
    for i in range(10):
        await _aggregate_assist(store, _assist_outcome(f"learner-{i}", turn_count=0, blocks=0))
    block_rate_cells = [
        c for c in store.upsert_cohort_aggregate.call_args_list
        if c.args[1] == "assist_guardrail_block_rate" and c.args[6] is False
    ]
    assert block_rate_cells
    rate = block_rate_cells[-1].args[4]
    assert rate == 0.0


# ── Practice branch unchanged (backward compat) ─────────────────────────────


@pytest.mark.asyncio
async def test_aggregate_session_dispatches_to_practice():
    """aggregate_session with session_type='practice' → _aggregate_practice."""
    store = _mock_pg_store()
    await aggregate_session(store, _practice_outcome("learner-1"))
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    # Practice metrics should be present.
    assert "sessions_count" in metrics
    assert "active_learners_count" in metrics
    # Assist metrics should NOT be present (practice branch).
    assert "assist_shifts_count" not in metrics


@pytest.mark.asyncio
async def test_aggregate_session_dispatches_to_assist():
    """aggregate_session with session_type='assist' → _aggregate_assist."""
    store = _mock_pg_store()
    await aggregate_session(store, _assist_outcome("learner-1"))
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    assert "assist_shifts_count" in metrics
    assert "sessions_count" not in metrics


@pytest.mark.asyncio
async def test_aggregate_session_default_is_practice():
    """aggregate_session with no session_type → practice (backward compat)."""
    store = _mock_pg_store()
    outcome = _practice_outcome("learner-1")
    outcome.pop("session_type")  # omit session_type → default practice
    await aggregate_session(store, outcome)
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    assert "sessions_count" in metrics
    assert "assist_shifts_count" not in metrics


# ── No PII in assist upsert calls ───────────────────────────────────────────


@pytest.mark.asyncio
async def test_assist_no_pii_in_upsert_calls():
    """No raw learner_ref leaks into assist aggregate cell args (D-031)."""
    store = _mock_pg_store()
    await _aggregate_assist(store, _assist_outcome("learner-sensitive-id-1234"))
    for c in store.upsert_cohort_aggregate.call_args_list:
        for arg in c.args:
            assert "learner-sensitive-id-1234" not in str(arg), \
                "raw learner_ref must not leak into assist aggregate cell args"
        assert isinstance(c.args[5], int)  # cell_count is an int


# ── Hook dispatches assist correctly ───────────────────────────────────────


@pytest.mark.asyncio
async def test_hook_dispatches_assist_session():
    """on_session_end with session_type='assist' → _aggregate_assist (no error)."""
    store = _mock_pg_store()
    await on_session_end(store, _assist_outcome("learner-1"))
    assert store.upsert_cohort_aggregate.called
    metrics = {c.args[1] for c in store.upsert_cohort_aggregate.call_args_list}
    assert "assist_shifts_count" in metrics


@pytest.mark.asyncio
async def test_hook_assist_no_postgres_is_noop():
    """on_session_end with no Postgres → no-op (assist hook)."""
    await on_session_end(None, _assist_outcome("learner-1"))


# ── Dashboard endpoints return assist rows ──────────────────────────────────


def _make_app_with_assist_rows(rows: list[dict]):
    """Build a minimal FastAPI app with the cohort + failure-patterns routers
    + a mocked pg_store returning `rows`."""
    from fastapi import FastAPI
    from fastapi.testclient import TestClient
    from server.operator.cohort import router as cohort_router
    from server.operator.failure_patterns import router as failure_router

    app = FastAPI()
    pg_store = MagicMock()
    pg_store.pool = MagicMock()
    conn = MagicMock()
    conn.fetch = AsyncMock(return_value=rows)
    cm = MagicMock()
    cm.__aenter__ = AsyncMock(return_value=conn)
    cm.__aexit__ = AsyncMock(return_value=None)
    pg_store.pool.acquire = MagicMock(return_value=cm)
    app.state.pg_store = pg_store
    # Bypass auth for these tests by stubbing current_operator.
    from server.auth.dependencies import current_operator
    from server.auth.models import Operator

    async def _stub_op():
        return Operator(id="op-1", username="tester", display_name="T", role="operator")
    app.dependency_overrides[current_operator] = _stub_op
    app.include_router(cohort_router)
    app.include_router(failure_router)
    return TestClient(app)


def _assist_metric_row(metric: str, value: float, suppressed: bool = False) -> dict:
    return {
        "path": "customer_service",
        "metric": metric,
        "window_start": _dt.date.today() - _dt.timedelta(days=6),
        "window_end": _dt.date.today(),
        "value": value if not suppressed else None,
        "cell_count": 12,
        "cell_suppressed": suppressed,
        "updated_at": _dt.datetime.now(_dt.timezone.utc),
    }


def test_cohort_endpoint_returns_assist_rows():
    """GET /api/operator/cohort returns assist_shifts_count + assist_turns_count."""
    rows = [
        _assist_metric_row("sessions_count", 15.0),
        _assist_metric_row("active_learners_count", 12.0),
        _assist_metric_row("assist_shifts_count", 8.0),
        _assist_metric_row("assist_turns_count", 160.0),
    ]
    client = _make_app_with_assist_rows(rows)
    r = client.get("/api/operator/cohort")
    assert r.status_code == 200
    data = r.json()
    metrics = {c["metric"] for v in data["views"] for c in v["metrics"]}
    assert "assist_shifts_count" in metrics
    assert "assist_turns_count" in metrics
    assert "sessions_count" in metrics  # practice still present


def test_failure_patterns_endpoint_returns_guardrail_block_rate():
    """GET /api/operator/failure-patterns returns assist_guardrail_block_rate."""
    rows = [
        _assist_metric_row("failure_mode:missed_apology", 3.0),
        _assist_metric_row("branch:escalate", 5.0),
        _assist_metric_row("assist_guardrail_block_rate", 0.08),
    ]
    client = _make_app_with_assist_rows(rows)
    r = client.get("/api/operator/failure-patterns")
    assert r.status_code == 200
    data = r.json()
    metrics = {c["metric"] for v in data["views"] for c in v["metrics"]}
    assert "assist_guardrail_block_rate" in metrics
    assert "failure_mode:missed_apology" in metrics  # practice failure patterns still present


def test_mastery_endpoint_excludes_assist_metrics():
    """D-063: GET /api/operator/mastery does NOT return assist metrics."""
    from fastapi import FastAPI
    from fastapi.testclient import TestClient
    from server.auth.dependencies import current_operator
    from server.auth.models import Operator
    from server.operator.mastery import router as mastery_router

    rows = [
        _assist_metric_row("gate_open_rate", 0.5),
        _assist_metric_row("median_mastery_score", 3.8),
        _assist_metric_row("assist_shifts_count", 8.0),  # should be EXCLUDED
        _assist_metric_row("assist_guardrail_block_rate", 0.08),  # EXCLUDED
    ]
    app = FastAPI()
    pg_store = MagicMock()
    pg_store.pool = MagicMock()
    conn = MagicMock()
    conn.fetch = AsyncMock(return_value=rows)
    cm = MagicMock()
    cm.__aenter__ = AsyncMock(return_value=conn)
    cm.__aexit__ = AsyncMock(return_value=None)
    pg_store.pool.acquire = MagicMock(return_value=cm)
    app.state.pg_store = pg_store

    async def _stub_op():
        return Operator(id="op-1", username="tester", display_name="T", role="operator")
    app.dependency_overrides[current_operator] = _stub_op
    app.include_router(mastery_router)
    client = TestClient(app)
    r = client.get("/api/operator/mastery")
    assert r.status_code == 200
    data = r.json()
    metrics = {c["metric"] for v in data["views"] for c in v["metrics"]}
    assert "gate_open_rate" in metrics
    assert "median_mastery_score" in metrics
    # D-063: assist metrics must NOT appear in the mastery view.
    assert "assist_shifts_count" not in metrics
    assert "assist_guardrail_block_rate" not in metrics


def test_cohort_endpoint_suppressed_assist_cells():
    """Suppressed assist cells have value=null + cell_suppressed=true (k-anon)."""
    rows = [
        _assist_metric_row("assist_shifts_count", 0.0, suppressed=True),
        _assist_metric_row("assist_turns_count", 0.0, suppressed=True),
    ]
    client = _make_app_with_assist_rows(rows)
    r = client.get("/api/operator/cohort")
    assert r.status_code == 200
    data = r.json()
    for v in data["views"]:
        for c in v["metrics"]:
            if c["metric"] in ("assist_shifts_count", "assist_turns_count"):
                assert c["cell_suppressed"] is True
                assert c["value"] is None