#!/usr/bin/env python3
"""Minimal BigQuery client: stdlib + `cryptography`. Dry-runs every query and refuses to run past a byte cap.

Credentials: path in env GOOGLE_APPLICATION_CREDENTIALS (service-account JSON). Never logged.
The NEAR public dataset is billed to the caller; the free sandbox gives 1 TB of scanned data a month.
"""
import base64, json, os, pathlib, sys, time, urllib.error, urllib.request

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
from common import MAX_BYTES, OPENER  # noqa: E402

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

TOKEN_URL = "https://oauth2.googleapis.com/token"
SCOPE = "https://www.googleapis.com/auth/bigquery"
API = "https://bigquery.googleapis.com/bigquery/v2"


def _b64(raw):
    return base64.urlsafe_b64encode(raw).rstrip(b"=")


class BQ:
    def __init__(self, cred_path=None, max_bytes=120 * 10**9):
        path = cred_path or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
        if not path:
            raise SystemExit("GOOGLE_APPLICATION_CREDENTIALS is not set")
        with open(path, encoding="utf8") as fh:
            self._cred = json.load(fh)
        self.project = self._cred["project_id"]
        self.max_bytes = int(max_bytes)
        self._token, self._exp = None, 0

    def token(self):
        if self._token and time.time() < self._exp - 60:
            return self._token
        now = int(time.time())
        head = _b64(json.dumps({"alg": "RS256", "typ": "JWT"}).encode())
        claims = _b64(json.dumps({"iss": self._cred["client_email"], "scope": SCOPE, "aud": TOKEN_URL,
                                  "iat": now, "exp": now + 3600}).encode())
        key = serialization.load_pem_private_key(self._cred["private_key"].encode(), password=None)
        sig = key.sign(head + b"." + claims, padding.PKCS1v15(), hashes.SHA256())
        body = ("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion="
                + (head + b"." + claims + b"." + _b64(sig)).decode()).encode()
        req = urllib.request.Request(TOKEN_URL, data=body, headers={"content-type": "application/x-www-form-urlencoded"})
        try:
            with OPENER.open(req, timeout=30) as resp:
                out = json.loads(resp.read(1000000))
        except urllib.error.HTTPError as exc:
            # Google's error code and description name the problem and carry no credentials
            try:
                info = json.loads(exc.read(4000))
                detail = f"{info.get('error')}: {str(info.get('error_description'))[:120]}"
            except (ValueError, AttributeError):
                detail = "no detail"
            raise SystemExit(f"Google token request failed, HTTP {exc.code}, {detail}") from None
        self._token, self._exp = out["access_token"], now + int(out.get("expires_in", 3600))
        return self._token

    def _call(self, url, payload=None, timeout=300):
        data = None if payload is None else json.dumps(payload).encode()
        req = urllib.request.Request(url, data=data, method="POST" if data else "GET",
                                     headers={"authorization": f"Bearer {self.token()}", "content-type": "application/json"})
        try:
            with OPENER.open(req, timeout=timeout) as resp:
                return json.loads(resp.read(MAX_BYTES))
        except urllib.error.HTTPError as exc:
            # error bodies can name the project and account, and this log may be public: keep only the reason code
            try:
                reason = json.loads(exc.read(20000))["error"]["errors"][0]["reason"]
            except (ValueError, KeyError, IndexError, TypeError):
                reason = "unknown"
            raise RuntimeError(f"BigQuery HTTP {exc.code}: {reason}") from None

    def dry_run(self, sql):
        out = self._call(f"{API}/projects/{self.project}/jobs", {
            "configuration": {"dryRun": True, "query": {"query": sql, "useLegacySql": False}}})
        return int(out["statistics"]["totalBytesProcessed"])

    def query(self, sql, label="query", cap=None, poll_s=4, max_wait_s=900):
        """Returns (rows as dicts, stats). Raises if the dry run exceeds the cap."""
        cap = int(cap or self.max_bytes)
        scanned = self.dry_run(sql)
        if scanned > cap:
            raise RuntimeError(f"{label}: dry run {scanned/1e9:.1f} GB exceeds cap {cap/1e9:.1f} GB")
        job = self._call(f"{API}/projects/{self.project}/jobs", {
            "configuration": {"query": {"query": sql, "useLegacySql": False, "maximumBytesBilled": str(cap)}}})
        ref = job["jobReference"]
        loc = ref.get("location", "US")
        deadline = time.time() + max_wait_s
        while True:
            res = self._call(f"{API}/projects/{self.project}/queries/{ref['jobId']}?location={loc}&timeoutMs=20000&maxResults=20000")
            if res.get("jobComplete"):
                break
            if time.time() > deadline:
                raise RuntimeError(f"{label}: timed out")
            time.sleep(poll_s)
        fields = [f["name"] for f in res["schema"]["fields"]]
        rows, page = [], res
        while True:
            rows.extend({n: c.get("v") for n, c in zip(fields, r["f"])} for r in page.get("rows", []))
            tok = page.get("pageToken")
            if not tok:
                break
            page = self._call(f"{API}/projects/{self.project}/queries/{ref['jobId']}?location={loc}&pageToken={tok}&maxResults=20000")
        return rows, {"bytes_processed": int(res.get("totalBytesProcessed") or scanned), "dry_run_bytes": scanned,
                      "rows": len(rows), "job_id": ref["jobId"]}