#!/usr/bin/env python3
"""Search and parse a public sample. Full downloads require an account credential.

Run: python3 first-retrieval.py [--mcp] [--download] [--check]
--check validates without saving a file (used by the scheduled live check).
"""
import argparse
import csv
import datetime
import io
import json
import pathlib
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

DATASET_ID = "8bc44aae-5b51-4b9d-9d4a-9078f3f9f553"
QUERY = "US federal holidays 2026"
MAX_BYTES = 2 * 1024 * 1024


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--mcp", action="store_true", help="Use hosted public MCP")
    parser.add_argument("--check", action="store_true", help="Validate without saving")
    parser.add_argument("--api-url", default="https://api.databazaar.io")
    parser.add_argument("--download", action="store_true", help="Retrieve full data with DATABAZAAR_API_KEY from your account")
    parser.add_argument("--check-access", action="store_true", help="Verify anonymous full-data access is rejected")
    args = parser.parse_args()
    token = os.environ.get("DATABAZAAR_API_KEY", "") if args.download else ""
    if args.download and not token:
        raise RuntimeError("Full downloads require an account. Sign up at https://databazaar.io/signup, create a key at /operator/keys, and set DATABAZAAR_API_KEY. Without --download this script only previews a sample.")
    started = time.monotonic()
    base = args.api_url.rstrip("/")
    stages = {}

    def fetch(url, payload=None):
        remaining = 115 - (time.monotonic() - started)
        if remaining <= 0:
            raise RuntimeError("Retrieval exceeded the two-minute target")
        headers = {"User-Agent": "DataBazaar-Quickstart/1.0" + (" synthetic-check" if args.check else "")}
        if token and urllib.parse.urlsplit(url).netloc == urllib.parse.urlsplit(base).netloc:
            headers["Authorization"] = "Bearer " + token
        if payload is not None:
            headers.update({"Content-Type": "application/json", "Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-03-26"})
        request = urllib.request.Request(url, data=json.dumps(payload).encode() if payload is not None else None, headers=headers)
        with urllib.request.urlopen(request, timeout=min(20, remaining)) as response:
            raw = response.read(MAX_BYTES + 1)
        if len(raw) > MAX_BYTES:
            raise RuntimeError("Example grew beyond 2 MB; choose a smaller first dataset")
        return raw.decode("utf-8-sig")

    def rpc(method, params):
        raw = fetch(base + ("/mcp" if args.download else "/mcp/public"), {"jsonrpc": "2.0", "id": 1, "method": method, "params": params})
        # Streamable HTTP servers may return either JSON or SSE.
        messages = [json.loads(line[6:]) for line in raw.splitlines() if line.startswith("data: ")] if raw.startswith("event:") else [json.loads(raw)]
        response = next(message for message in messages if message.get("id") == 1)
        if "error" in response:
            raise RuntimeError(response["error"]["message"])
        return response["result"]

    def call(tool, arguments, rest_path):
        if not args.mcp:
            return json.loads(fetch(base + rest_path))
        result = rpc("tools/call", {"name": tool, "arguments": arguments})
        text = next(item["text"] for item in result["content"] if item["type"] == "text")
        if result.get("isError"):
            raise RuntimeError(text)
        return json.loads(text)

    if args.mcp:
        rpc("initialize", {"protocolVersion": "2025-03-26", "capabilities": {}, "clientInfo": {"name": "databazaar-first-retrieval", "version": "1.0"}})
        tools = {tool["name"] for tool in rpc("tools/list", {})["tools"]}
        required = {"search_datasets", "get_dataset", "preview_sample"} | ({"get_download_url"} if args.download else set())
        if not required.issubset(tools):
            raise RuntimeError("Public MCP connection is missing retrieval tools")
        stages["connected"] = round(time.monotonic() - started, 3)

    search = call("search_datasets", {"query": QUERY, "pricing_type": "free", "limit": 10}, "/datasets?" + urllib.parse.urlencode({"query": QUERY, "pricing_type": "free", "limit": 10}))
    listing = next((row for row in search["data"] if row["id"] == DATASET_ID), None)
    if not listing:
        raise RuntimeError("Example listing is unavailable. Search /datasets?pricing_type=free for another dataset")
    stages["searched"] = round(time.monotonic() - started, 3)
    detail = call("get_dataset", {"dataset_id": DATASET_ID}, "/datasets/" + DATASET_ID)
    if detail["pricing_type"] != "free" or detail.get("price") not in (None, 0):
        raise RuntimeError("Example is no longer free; no purchase was attempted")
    sample = call("preview_sample", {"dataset_id": DATASET_ID}, "/datasets/" + DATASET_ID + "/sample")
    preview = list(csv.DictReader(io.StringIO(fetch(sample["sample_url"]))))
    if not preview or not {"holiday", "date", "observed_date"}.issubset(preview[0]):
        raise RuntimeError("Preview does not contain the expected holiday columns")
    stages["previewed"] = round(time.monotonic() - started, 3)
    if args.check_access:
        if args.download:
            raise RuntimeError("--check-access must be run without --download")
        for suffix, payload in [("download-url", None), ("manifest", None), ("files/0", None), ("query", {})]:
            try:
                fetch(base + "/datasets/" + DATASET_ID + "/" + suffix, payload)
            except urllib.error.HTTPError as error:
                if error.code == 401:
                    continue
                raise RuntimeError("Expected anonymous full-data access to return 401; got " + str(error.code)) from None
            raise RuntimeError("Anonymous full-data access was not rejected: " + suffix)
    if not args.download:
        print(json.dumps({"success": True, "scope": "sample_only", "transport": "mcp" if args.mcp else "rest", "dataset_id": DATASET_ID, "sample_rows": len(preview), "first_row": preview[0], "seconds": round(time.monotonic() - started, 3), "next_step": "Create an account and authorize your agent to retrieve the complete dataset.", "signup_url": "https://databazaar.io/signup?returnTo=%2Fdatasets%2F" + DATASET_ID, "mcp_account_url": base + "/mcp"}, indent=2))
        return
    delivery = call("get_download_url", {"dataset_id": DATASET_ID}, "/datasets/" + DATASET_ID + "/download-url")
    if not delivery.get("download_url"):
        raise RuntimeError("Example now uses a multi-file manifest; inspect manifest_url before downloading")
    raw = fetch(delivery.get("retrieval_url") or delivery["download_url"])
    rows = list(csv.DictReader(io.StringIO(raw)))
    if len(rows) != 11 or any(not row.get("holiday") for row in rows):
        raise RuntimeError("Expected the complete 11-row holiday dataset")
    for row in rows:
        datetime.date.fromisoformat(row["date"])
        datetime.date.fromisoformat(row["observed_date"])
    elapsed = time.monotonic() - started
    if elapsed >= 120:
        raise RuntimeError("Retrieved data but missed the two-minute target")
    if not args.check:
        pathlib.Path("databazaar-holidays.csv").write_text(raw, encoding="utf-8")
    print(json.dumps({"success": True, "scope": "full_dataset", "transport": "mcp" if args.mcp else "rest", "dataset_id": DATASET_ID, "rows": len(rows), "columns": list(rows[0]), "first_row": rows[0], "seconds": round(elapsed, 3), "stages_seconds": stages, "saved_file": None if args.check else "databazaar-holidays.csv"}, indent=2))


if __name__ == "__main__":
    try:
        main()
    except (urllib.error.URLError, RuntimeError, ValueError, KeyError, StopIteration, OSError) as error:
        print("Retrieval failed: " + str(error), file=sys.stderr)
        sys.exit(1)
