#!/usr/bin/env python3
"""Auto-resolve an O11Y cluster, run file-backed PromQL, and persist evidence."""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import subprocess
from pathlib import Path
from typing import Any

import requests

from o11y_common import get_auth0_token, o11y_403_message, resolve_cluster, write_private_json


def parse_time(value: str) -> dt.datetime:
    try:
        parsed = dt.datetime.fromtimestamp(float(value), tz=dt.timezone.utc)
    except ValueError:
        parsed = dt.datetime.fromisoformat(value.replace("Z", "+00:00"))
        if parsed.tzinfo is None:
            parsed = parsed.replace(tzinfo=dt.timezone.utc)
        parsed = parsed.astimezone(dt.timezone.utc)
    return parsed


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Resolve a cluster automatically and query O11Y Prometheus metrics."
    )
    parser.add_argument("--cluster-id", required=True)
    parser.add_argument("--tenant-id", help="Optional expected tenant ID for safety validation.")
    parser.add_argument("--provider", choices=("aws", "gcp", "alicloud"), help="Optional hint.")
    parser.add_argument("--region", help="Optional region hint.")
    parser.add_argument("--deploy-type", help="Optional expected deploy type.")
    parser.add_argument("--target", choices=("app", "infra"), default="app")
    parser.add_argument("--query-file", type=Path, required=True)
    parser.add_argument("--start", required=True, help="Inclusive epoch or ISO-8601; naive values are UTC.")
    parser.add_argument("--end-exclusive", required=True, help="Exclusive epoch or ISO-8601.")
    parser.add_argument("--step", type=int, default=60, help="Prometheus step in seconds.")
    parser.add_argument("--timeout", type=int, default=60)
    parser.add_argument("--cluster-output", type=Path, required=True)
    parser.add_argument("--raw-output", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True, help="Normalized JSONL samples.")
    parser.add_argument("--metadata-output", type=Path, required=True)
    args = parser.parse_args()
    args.start_time = parse_time(args.start)
    args.end_time = parse_time(args.end_exclusive)
    if args.start_time >= args.end_time:
        raise ValueError("Metrics start time must be before end-exclusive.")
    if args.step < 1:
        raise ValueError("Metrics step must be positive.")
    return args


def read_query(path: Path) -> str:
    if path.is_symlink() or not path.is_file():
        raise ValueError(f"PromQL query file is missing or a symlink: {path}")
    query = path.read_text(encoding="utf-8").strip()
    if not query:
        raise ValueError("PromQL query file is empty.")
    return query


def request_metrics(
    endpoint: str,
    token: str,
    query: str,
    start: int,
    end_inclusive: int,
    step: int,
    timeout: int,
) -> tuple[int, dict[str, Any]]:
    parameters = {
        "query": query,
        "start": str(start),
        "end": str(end_inclusive),
        "step": str(step),
    }
    initial_status_code = 0
    try:
        response = requests.get(
            endpoint,
            headers={"Authorization": f"Bearer {token}", "Accept": "application/json"},
            params=parameters,
            timeout=timeout,
        )
        initial_status_code = response.status_code
        try:
            body = response.json()
        except json.JSONDecodeError:
            body = {"status": "error", "error": response.text[:2000]}
        if response.status_code == 200:
            return response.status_code, body
    except requests.RequestException as exc:
        body = {"status": "error", "error": type(exc).__name__}

    command = ["curl", "-sS", "--max-time", str(timeout), endpoint, "-G"]
    for key, value in parameters.items():
        command.extend(["--data-urlencode", f"{key}={value}"])
    command.extend(
        ["-H", f"Authorization: Bearer {token}", "-H", "Accept: application/json", "-w", "\n%{http_code}"]
    )
    try:
        completed = subprocess.run(command, check=False, capture_output=True, text=True)
        raw_body, _, raw_status = completed.stdout.rpartition("\n")
        if not raw_status.isdigit():
            return initial_status_code, body
        status_code = int(raw_status)
        try:
            curl_body = json.loads(raw_body)
        except json.JSONDecodeError:
            curl_body = {"status": "error", "error": raw_body[:2000]}
        # Preserve an initial 403 so a failed/unavailable curl fallback cannot
        # suppress the single forced re-authentication retry.
        if initial_status_code == 403 and status_code != 200:
            return initial_status_code, curl_body
        return status_code, curl_body
    except (OSError, ValueError):
        return initial_status_code, body


def normalize_samples(body: dict[str, Any]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for series_index, series in enumerate(body.get("data", {}).get("result", [])):
        labels = series.get("metric", {})
        for timestamp, value in series.get("values", []):
            timestamp_value = float(timestamp)
            rows.append(
                {
                    "series_index": series_index,
                    "timestamp_unix": str(timestamp),
                    "timestamp_utc": dt.datetime.fromtimestamp(
                        timestamp_value, tz=dt.timezone.utc
                    ).isoformat(),
                    "value": str(value),
                    "labels": labels,
                }
            )
    rows.sort(key=lambda row: (float(row["timestamp_unix"]), row["series_index"]))
    return rows


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path = path.expanduser().resolve()
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    os.chmod(path, 0o600)


def main() -> int:
    os.umask(0o077)
    args = parse_args()
    query = read_query(args.query_file)
    context, token = resolve_cluster(
        cluster_id=args.cluster_id,
        expected_tenant_id=args.tenant_id,
        provider_hint=args.provider,
        region_hint=args.region,
        deploy_type_hint=args.deploy_type,
        analysis_timestamp=int(args.end_time.timestamp()) - 1,
        timeout=args.timeout,
    )
    write_private_json(args.cluster_output, context)

    id_key = "app_o11y_id" if args.target == "app" else "infra_o11y_id"
    base_key = "app_metrics_base" if args.target == "app" else "infra_metrics_base"
    o11y_id = str(context.get(id_key) or "")
    metrics_base = str(context.get(base_key) or "").rstrip("/")
    if not o11y_id:
        raise ValueError(f"No {args.target} O11Y metrics target was resolved.")
    if not metrics_base:
        metrics_base = (
            f"https://www.ds.{context['region']}.{context['provider']}.observability.tidbcloud.com"
            f"/internal/metrics/{o11y_id}/api/v1"
        )
    endpoint = f"{metrics_base}/query_range"
    start_epoch = int(args.start_time.timestamp())
    end_inclusive = int(args.end_time.timestamp()) - 1
    status_code, body = request_metrics(
        endpoint,
        token,
        query,
        start_epoch,
        end_inclusive,
        args.step,
        args.timeout,
    )
    reauth_attempted = False
    reauth_succeeded = False
    reauth_error: Exception | None = None
    if status_code == 403:
        reauth_attempted = True
        try:
            token, _, _ = get_auth0_token(
                context["provider"], args.timeout, force_refresh=True
            )
            reauth_succeeded = True
            context["authentication"] = "refreshed_after_403"
            status_code, body = request_metrics(
                endpoint,
                token,
                query,
                start_epoch,
                end_inclusive,
                args.step,
                args.timeout,
            )
        except Exception as exc:
            reauth_error = exc
    raw = {
        "query": query,
        "cluster": context,
        "target": args.target,
        "start_utc_inclusive": args.start_time.isoformat(),
        "end_utc_exclusive": args.end_time.isoformat(),
        "query_end_epoch_inclusive": end_inclusive,
        "step_seconds": args.step,
        "http_status": status_code,
        "reauth_attempted_after_403": reauth_attempted,
        "reauth_succeeded_after_403": reauth_succeeded,
        "response": body,
    }
    write_private_json(args.raw_output, raw)

    success = status_code == 200 and body.get("status") == "success"
    partial = bool(body.get("isPartial") or body.get("data", {}).get("isPartial"))
    rows = normalize_samples(body) if success else []
    write_jsonl(args.output, rows)
    metadata = {
        "source": "o11y_prometheus",
        "cluster_id": args.cluster_id,
        "tenant_id": context["tenant_id"],
        "provider": context["provider"],
        "region": context["region"],
        "deploy_type": context["deploy_type"],
        "biz_type": context["biz_type"],
        "target": args.target,
        "o11y_id": o11y_id,
        "endpoint": endpoint,
        "query_file": str(args.query_file.resolve()),
        "query_sha256": hashlib.sha256(query.encode("utf-8")).hexdigest(),
        "start_utc_inclusive": args.start_time.isoformat(),
        "end_utc_exclusive": args.end_time.isoformat(),
        "step_seconds": args.step,
        "http_status": status_code,
        "success": success,
        "reauth_attempted_after_403": reauth_attempted,
        "reauth_succeeded_after_403": reauth_succeeded,
        "partial": partial,
        "series_count": len(body.get("data", {}).get("result", [])) if success else 0,
        "sample_count": len(rows),
        "raw_output": str(args.raw_output.resolve()),
        "normalized_output": str(args.output.resolve()),
        "cluster_output": str(args.cluster_output.resolve()),
        "retrieved_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
    }
    write_private_json(args.metadata_output, metadata)
    if not success:
        if reauth_attempted and status_code == 403:
            raise RuntimeError(o11y_403_message("metrics", reauth_error))
        if reauth_error is not None:
            raise RuntimeError(o11y_403_message("metrics", reauth_error))
        raise RuntimeError(f"O11Y metrics query failed with HTTP {status_code}.")
    if partial:
        raise RuntimeError("O11Y metrics query returned a partial result; billing use is blocked.")
    print(json.dumps(metadata, ensure_ascii=False))
    return 0


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