"""Integration test for benchmark Group B: fixture grading -> Q_quality -> reward.

Boots the in-memory store, seeds a completed run whose tasks carry generated code, grades it
against the held-out `add_function` fixture in the sandbox, and asserts Q_quality and reward
become real. Hermetic, no model key. Run from agent_swarm_v6 (install deps first):
    pip install -r orchestrator/requirements.txt -r agent/requirements.txt
    python scripts/test-quality.py
(this test sets REDIS_FAKE / ENABLE_QUALITY_EVAL / HEICODE_SANDBOX_ISOLATED itself.)
"""
import asyncio
import json
import math
import os
import sys
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
os.environ["ENABLE_QUALITY_EVAL"] = "1"  # gate ON for this test (executes code in the sandbox)
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator.redis_client import redis_client
from orchestrator import swarm_runtime as sr_mod
from orchestrator.swarm_runtime import swarm_runtime
from orchestrator.task_queue import task_queue, TaskStatus
from orchestrator.quality import evaluate_run_quality, collect_generated_files
from orchestrator.sandbox import SandboxIsolationError
from benchmark.collectors.run_collector import SwarmRunMetricsCollector

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


async def _noop(self, *a, **k):
    return None


# A correct implementation + the swarm's OWN test (which must NOT be the grader).
IMPL = "def add(a, b):\n    return a + b\n"
RESULT = {
    "success": True,
    "subtasks": [
        {"status": "completed", "files": [{"path": "calc.py", "action": "write", "content": IMPL}]},
        {"status": "completed", "files": [{"path": "test_calc.py", "action": "write",
            "content": "from calc import add\ndef test_self():\n    assert add(1, 1) == 2\n"}]},
    ],
    "usage": {"model_cost_usd": 2.0},
}


async def add_completed_task(run, task_id, result):
    t = await task_queue.create_task(task_id=task_id, description=task_id,
                                     agent_role=task_id.split("-")[-1], depends_on=[], enqueue=False)
    t.status = TaskStatus.COMPLETED
    t.assigned_agent_id = "A"
    t.retry_count = 0
    t.result = json.dumps(result)
    await task_queue._save_task(t)
    await swarm_runtime.attach_task(run, t.task_id)
    return t


async def main():
    await redis_client.connect()
    sr_mod.SwarmRuntime._post_callback = _noop

    body = {
        "mode": "swarm",
        "requirement": {"objective": "Write add(a,b)"},
        "orchestration_plan": {"budget": {"max_cost_usd": 10}},
        "callback": {"url": "http://127.0.0.1:9/cb", "subscribed_events": []},
        "metadata": {"manager_deployment_id": "m-q", "benchmark_fixture_id": "add_function"},
    }
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="cq")
    impl_task = await add_completed_task(run, "t-implementation", RESULT)
    tasks = [impl_task]

    # FAIL-CLOSED: with eval enabled but isolation NOT confirmed, grading must hard-fail (not
    # silently skip and not execute code).
    os.environ.pop("HEICODE_SANDBOX_ISOLATED", None)
    raised = False
    try:
        await evaluate_run_quality(run, tasks)
    except SandboxIsolationError:
        raised = True
    check("fail-closed: eval enabled without isolation raises", raised)
    os.environ["HEICODE_SANDBOX_ISOLATED"] = "1"  # confirm isolation for the rest (CI/test runner)

    # file separation: impl vs the swarm's own tests
    files = collect_generated_files(tasks)
    check("collect splits impl vs agent tests",
          [f.path for f in files["impl"]] == ["calc.py"] and len(files["agent_tests"]) == 1)

    # grade against held-out fixture tests in the sandbox
    quality = await evaluate_run_quality(run, tasks)
    check("quality produced", quality is not None)
    check("fixture TestPassRate = 100 (4/4 held-out)", quality and quality["test_pass_rate"] == 100.0)
    check("Q_quality = 100 (masked: only test_pass present)", quality and quality["q_quality"] == 100.0)
    check("agent self-test kept as separate signal", quality and quality["agent_test_pass_rate"] == 100.0)
    check("code_review/user_acceptance remain uncollected (None)",
          quality and quality["code_review_score"] is None and quality["user_acceptance"] is None)

    # record + simulate a 30s run so V_speed is computable
    await swarm_runtime.record_quality(run, quality)
    run.created_at = run.updated_at - 30.0
    await swarm_runtime.save_run(run)

    collector = SwarmRunMetricsCollector(run.swarm_id)
    m = await collector.collect()
    check("reward is now REAL (not NaN)", not math.isnan(m.reward) and collector.coverage["reward"] is True)
    # gain/p_decision/benchmark must still be NaN — Group B does not close them
    check("gain still NaN (needs baselines)", math.isnan(m.s_gain) and collector.coverage["s_gain"] is False)
    check("benchmark still NaN (aggregate)", math.isnan(m.benchmark) and collector.coverage["benchmark"] is False)

    print()
    if failures:
        print(f"{len(failures)} quality check(s) FAILED: {failures}")
        sys.exit(1)
    print("all Group B quality checks passed")


if __name__ == "__main__":
    asyncio.run(main())