#!/usr/bin/env python3
"""Run a read-only, parameterized BigQuery query using ADC."""

from __future__ import annotations

import argparse
import base64
import datetime as dt
import decimal
import hashlib
import json
import os
import re
import sys
from pathlib import Path
from typing import Any

from bq_validate import validate_sql
from runtime_config import load_runtime_config
from semantic_model import load_semantic_model


FORBIDDEN_STATEMENTS = re.compile(
    r"\b(ALTER|CALL|CREATE|DELETE|DROP|EXPORT|GRANT|INSERT|LOAD|MERGE|REVOKE|TRUNCATE|UPDATE)\b",
    re.IGNORECASE,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run a read-only BigQuery query and save JSONL plus job metadata."
    )
    parser.add_argument(
        "--project",
        default=os.environ.get("BQ_JOB_PROJECT") or None,
        help=(
            "Optional project that creates the query job. Leave unset to infer "
            "it from Application Default Credentials."
        ),
    )
    parser.add_argument(
        "--location",
        default=os.environ.get("BQ_LOCATION") or None,
        help=(
            "Optional BigQuery job location. Leave unset or use 'auto' to let "
            "BigQuery infer it from referenced tables."
        ),
    )
    parser.add_argument("--sql-file", type=Path, required=True)
    parser.add_argument("--params-file", type=Path, help="JSON named scalar parameters.")
    parser.add_argument("--semantic-model", type=Path)
    parser.add_argument("--output", type=Path, required=True, help="JSONL result path.")
    parser.add_argument("--metadata-output", type=Path, required=True)
    parser.add_argument(
        "--max-bytes-billed",
        type=int,
        default=int(os.environ.get("BQ_MAX_BYTES_BILLED", str(10 * 1024**3))),
    )
    execution = parser.add_mutually_exclusive_group()
    execution.add_argument("--dry-run", action="store_true")
    execution.add_argument(
        "--auto-execute",
        action="store_true",
        help=(
            "Perform a dry run and automatically execute when estimated bytes "
            "are within --max-bytes-billed."
        ),
    )
    parser.add_argument(
        "--max-result-rows",
        type=int,
        help="Maximum rows written to JSONL; defaults to the semantic policy.",
    )
    parser.add_argument("--run-id", default="manual-invoice")
    args = parser.parse_args()
    if args.location and args.location.lower() == "auto":
        args.location = None
    return args


def strip_comments_and_literals(sql: str) -> str:
    without_comments = re.sub(r"--[^\n]*|/\*.*?\*/", " ", sql, flags=re.DOTALL)
    return re.sub(r"'(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"|`[^`]*`", " ", without_comments)


def validate_read_only_sql(sql: str) -> None:
    normalized = strip_comments_and_literals(sql).strip()
    if not normalized:
        raise ValueError("M1 BigQuery SQL must not be empty.")
    if FORBIDDEN_STATEMENTS.search(normalized):
        raise ValueError("M1 BigQuery SQL contains a forbidden write or DDL statement.")


def convert_parameter_value(type_name: str, value: Any) -> Any:
    type_name = type_name.upper()
    if value is None:
        return None
    if type_name == "INT64":
        return int(value)
    if type_name == "FLOAT64":
        return float(value)
    if type_name == "BOOL":
        if isinstance(value, bool):
            return value
        return str(value).lower() in {"1", "true", "yes"}
    if type_name in {"NUMERIC", "BIGNUMERIC"}:
        return decimal.Decimal(str(value))
    if type_name == "DATE":
        return dt.date.fromisoformat(str(value))
    if type_name == "DATETIME":
        return dt.datetime.fromisoformat(str(value))
    if type_name == "TIMESTAMP":
        return dt.datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    return str(value)


def load_query_parameters(bigquery: Any, params_file: Path | None) -> list[Any]:
    if params_file is None:
        return []
    raw = json.loads(params_file.read_text(encoding="utf-8"))
    if not isinstance(raw, dict):
        raise ValueError("Parameter file must be a JSON object keyed by parameter name.")

    parameters = []
    for name, spec in raw.items():
        if not isinstance(spec, dict) or "type" not in spec or "value" not in spec:
            raise ValueError(f"Parameter {name!r} must contain type and value.")
        type_name = str(spec["type"]).upper()
        value = convert_parameter_value(type_name, spec["value"])
        parameters.append(bigquery.ScalarQueryParameter(name, type_name, value))
    return parameters


def json_value(value: Any) -> Any:
    if isinstance(value, decimal.Decimal):
        return str(value)
    if isinstance(value, (dt.date, dt.datetime, dt.time)):
        return value.isoformat()
    if isinstance(value, bytes):
        return base64.b64encode(value).decode("ascii")
    if isinstance(value, tuple):
        return [json_value(item) for item in value]
    if isinstance(value, list):
        return [json_value(item) for item in value]
    if isinstance(value, dict):
        return {key: json_value(item) for key, item in value.items()}
    return value


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def main() -> int:
    config_path = load_runtime_config()
    args = parse_args()
    sql = args.sql_file.read_text(encoding="utf-8")
    semantic_model = load_semantic_model(args.semantic_model)
    sql_validation = validate_sql(sql, args.semantic_model)
    if not sql_validation["valid"]:
        raise ValueError("SQL policy validation failed: " + "; ".join(sql_validation["errors"]))
    validate_read_only_sql(sql)
    max_result_rows = args.max_result_rows
    if max_result_rows is None:
        max_result_rows = int(
            semantic_model.get("policy", {}).get("maximum_result_rows", 10000)
        )
    if max_result_rows < 1:
        raise ValueError("--max-result-rows must be at least 1.")

    try:
        from google.cloud import bigquery
    except ImportError:
        print(
            "google-cloud-bigquery is required. Install dependencies from requirements.txt.",
            file=sys.stderr,
        )
        return 2

    query_parameters = load_query_parameters(bigquery, args.params_file)

    def job_config(dry_run: bool) -> Any:
        return bigquery.QueryJobConfig(
            dry_run=dry_run,
            use_query_cache=not dry_run,
            maximum_bytes_billed=args.max_bytes_billed,
            query_parameters=query_parameters,
            labels={"application": "manual-invoice"},
        )

    client = bigquery.Client(project=args.project)
    dry_run_estimate: dict[str, Any] | None = None
    if args.auto_execute:
        estimate_job = client.query(sql, job_config=job_config(True), location=args.location)
        estimated_bytes = int(getattr(estimate_job, "total_bytes_processed", 0) or 0)
        dry_run_estimate = {
            "location": getattr(estimate_job, "location", None) or args.location,
            "total_bytes_processed": estimated_bytes,
            "maximum_bytes_billed": args.max_bytes_billed,
        }
        if estimated_bytes > args.max_bytes_billed:
            raise ValueError(
                "Dry-run estimate exceeds maximum bytes billed: "
                f"{estimated_bytes} > {args.max_bytes_billed}."
            )
        job = client.query(sql, job_config=job_config(False), location=args.location)
    else:
        job = client.query(sql, job_config=job_config(args.dry_run), location=args.location)

    metadata: dict[str, Any] = {
        "run_id": args.run_id,
        "dry_run": args.dry_run,
        "auto_execute": args.auto_execute,
        "dry_run_estimate": dry_run_estimate,
        "config_path": str(config_path) if config_path else None,
        "requested_job_project": args.project,
        "job_project": client.project,
        "requested_location": args.location,
        "location": getattr(job, "location", None) or args.location,
        "sql_file": str(args.sql_file),
        "sql_sha256": hashlib.sha256(sql.encode("utf-8")).hexdigest(),
        "sql_validation": sql_validation,
        "semantic_model_path": semantic_model["_path"],
        "maximum_bytes_billed": args.max_bytes_billed,
        "maximum_result_rows": max_result_rows,
        "job_id": getattr(job, "job_id", None),
        "total_bytes_processed": getattr(job, "total_bytes_processed", None),
        "total_bytes_billed": getattr(job, "total_bytes_billed", None),
        "cache_hit": getattr(job, "cache_hit", None),
    }

    if args.dry_run:
        write_json(args.metadata_output, metadata)
        print(json.dumps(metadata, ensure_ascii=False))
        return 0

    rows = job.result(max_results=max_result_rows)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    row_count = 0
    with args.output.open("w", encoding="utf-8") as handle:
        for row in rows:
            record = {key: json_value(row[key]) for key in row.keys()}
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")
            row_count += 1

    total_available_rows = getattr(rows, "total_rows", None)
    truncated = total_available_rows is not None and total_available_rows > row_count

    metadata.update(
        {
            "row_count": row_count,
            "total_available_rows": total_available_rows,
            "rows_written": row_count,
            "truncated": truncated,
            "started_at": job.started.isoformat() if job.started else None,
            "ended_at": job.ended.isoformat() if job.ended else None,
            "total_bytes_processed": job.total_bytes_processed,
            "total_bytes_billed": job.total_bytes_billed,
            "cache_hit": job.cache_hit,
            "schema": [
                {"name": field.name, "type": field.field_type, "mode": field.mode}
                for field in (job.schema or rows.schema or [])
            ],
        }
    )
    write_json(args.metadata_output, metadata)
    print(json.dumps(metadata, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())