#!/usr/bin/env python3
"""Capture scoped MCP retrieval results for a frozen, blind question packet.

This records evidence for independent grading; it does not infer relevance from
keyword overlap or assign semantic scores. Use an authenticated MCP endpoint.
"""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import time
from datetime import UTC, datetime
from pathlib import Path

from mcp import Client
from mcp.client.streamable_http import streamable_http_client


async def run(args: argparse.Namespace) -> None:
    raw = args.questions.read_bytes()
    questions = json.loads(raw)
    if not isinstance(questions, list) or not questions:
        raise ValueError("Questions must be a nonempty JSON list")
    ids: set[str] = set()
    for question in questions:
        if not isinstance(question, dict) or not all(
            isinstance(question.get(field), str) and question[field].strip()
            for field in ("id", "question", "partition")
        ):
            raise ValueError("Each question requires nonempty id, question and partition strings")
        if question["id"] in ids:
            raise ValueError("Question IDs must be unique")
        ids.add(question["id"])

    # Exclusive creation prevents accidentally overwriting an earlier scored run.
    with args.output.open("x", encoding="utf-8") as output:
        output.write(
            json.dumps(
                {
                    "kind": "manifest",
                    "created_utc": datetime.now(UTC).isoformat(),
                    "questions_sha256": hashlib.sha256(raw).hexdigest(),
                    "project_id": args.project,
                    "limit": 5,
                    "budget_tokens": args.budget_tokens,
                    "question_count": len(questions),
                    "collect_top_five": args.collect_top_five,
                    "scoring": "independent relevance judgments required",
                }
            )
            + "\n"
        )
        output.flush()
        async with Client(streamable_http_client(args.mcp_url)) as client:
            for question in questions:
                started = time.monotonic()
                responses = []
                node_ids: list[str] = []
                cursor: str | None = None
                seen_cursors: set[str] = set()
                is_error = False
                pagination_blocked = False
                for _page in range(5 if args.collect_top_five else 1):
                    arguments = {
                        "query": question["question"],
                        "projects": [args.project],
                        "include_shared": True,
                        "all_projects": False,
                        "limit": 5,
                        "budget_tokens": args.budget_tokens,
                        "include_evidence": False,
                    }
                    if cursor is not None:
                        arguments["cursor"] = cursor
                    response = await client.call_tool("search", arguments)
                    is_error = is_error or response.is_error
                    pages = [
                        json.loads(item.text) for item in response.content if item.type == "text"
                    ]
                    responses.extend(pages)
                    if response.is_error or not pages:
                        break
                    envelope = pages[0]
                    rows = envelope.get("data", [])
                    node_ids.extend(row["node"]["node_id"] for row in rows)
                    next_cursor = envelope.get("next_cursor")
                    if len(node_ids) >= 5 or next_cursor is None:
                        break
                    if not rows or next_cursor in seen_cursors:
                        pagination_blocked = True
                        break
                    seen_cursors.add(next_cursor)
                    cursor = next_cursor
                result = {
                    "kind": "retrieval",
                    "question": question,
                    "seconds": round(time.monotonic() - started, 6),
                    "is_error": is_error,
                    "response": responses,
                    "returned_node_ids": node_ids[:5],
                    "pagination_blocked": pagination_blocked,
                }
                output.write(json.dumps(result, ensure_ascii=False) + "\n")
                output.flush()
                print(f"{question['id']}: {'error' if is_error else 'captured'}", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--questions", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--project", required=True)
    parser.add_argument("--mcp-url", required=True)
    parser.add_argument("--budget-tokens", type=int, default=8000)
    parser.add_argument("--collect-top-five", action="store_true")
    args = parser.parse_args()
    if not 256 <= args.budget_tokens <= 50_000:
        parser.error("--budget-tokens must be between 256 and 50000")
    asyncio.run(run(args))


if __name__ == "__main__":
    main()