#!/usr/bin/env python3
"""Query O11Y Loki and persist the query, raw pages, normalized logs, and metadata."""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

from o11y_common import get_auth0_token, o11y_403_message
from runtime_config import load_runtime_config


AUTH_URL = "https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth"


def parse_time(value: str | None, default: dt.datetime) -> dt.datetime:
    if value is None:
        return default
    try:
        return 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)
        return parsed.astimezone(dt.timezone.utc)


def parse_args() -> argparse.Namespace:
    now = dt.datetime.now(dt.timezone.utc)
    parser = argparse.ArgumentParser(description="Query O11Y Loki with persisted evidence.")
    parser.add_argument("--provider", choices=("aws", "gcp", "alicloud"), required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--query-file", type=Path, required=True)
    parser.add_argument("--start", help="Inclusive epoch seconds or ISO-8601; defaults to one hour ago.")
    parser.add_argument("--end", help="Exclusive epoch seconds or ISO-8601; defaults to now.")
    parser.add_argument("--page-limit", type=int, default=1000)
    parser.add_argument("--max-pages", type=int, default=100)
    parser.add_argument("--timeout", type=int, default=30)
    parser.add_argument("--raw-output", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True, help="Normalized JSONL output.")
    parser.add_argument("--metadata-output", type=Path, required=True)
    parser.add_argument("--run-id", default="manual-invoice-loki")
    args = parser.parse_args()
    args.end_time = parse_time(args.end, now)
    args.start_time = parse_time(args.start, args.end_time - dt.timedelta(hours=1))
    if args.start_time >= args.end_time:
        raise ValueError("Loki start time must be before end time.")
    if args.page_limit < 1 or args.page_limit > 1000 or args.max_pages < 1:
        if args.page_limit > 1000:
            raise ValueError("O11Y Loki page limit cannot exceed 1000.")
        raise ValueError("Page limit and max pages must be positive.")
    return args


def request_json(request: urllib.request.Request, timeout: int) -> dict[str, Any]:
    with urllib.request.urlopen(request, timeout=timeout) as response:
        body = json.loads(response.read().decode("utf-8"))
    if not isinstance(body, dict):
        raise ValueError("O11Y response must be a JSON object.")
    return body


def authenticate(github_pat: str, provider: str, timeout: int) -> str:
    payload = json.dumps({"github_token": github_pat, "provider": provider}).encode("utf-8")
    request = urllib.request.Request(
        AUTH_URL,
        data=payload,
        headers={"Content-Type": "application/json", "User-Agent": "manual-invoice"},
        method="POST",
    )
    last_error = "unknown error"
    for attempt in range(3):
        try:
            body = request_json(request, timeout)
            token = str(body.get("auth0_token") or "")
            if not token:
                raise ValueError("O11Y authentication returned no auth0_token.")
            return token
        except urllib.error.HTTPError as exc:
            last_error = f"HTTP {exc.code}"
        except (urllib.error.URLError, TimeoutError) as exc:
            last_error = type(exc).__name__
        if attempt < 2:
            time.sleep(2**attempt)
    raise RuntimeError(f"O11Y authentication failed: {last_error}")


class LokiHTTPError(RuntimeError):
    def __init__(self, status_code: int, detail: str):
        super().__init__(f"Loki query failed with HTTP {status_code}: {detail}")
        self.status_code = status_code


def loki_page(
    endpoint: str,
    auth0_token: str,
    query: str,
    start_ns: int,
    end_ns: int,
    limit: int,
    timeout: int,
) -> dict[str, Any]:
    params = urllib.parse.urlencode(
        {
            "query": query,
            "start": str(start_ns),
            "end": str(end_ns),
            "limit": str(limit),
            "direction": "forward",
        }
    )
    request = urllib.request.Request(
        f"{endpoint}/query_range?{params}",
        headers={
            "Authorization": f"Bearer {auth0_token}",
            "Accept": "application/json",
            "User-Agent": "manual-invoice",
        },
    )
    try:
        body = request_json(request, timeout)
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:1000]
        raise LokiHTTPError(exc.code, detail) from exc
    if body.get("status") != "success":
        raise RuntimeError(f"Loki query failed with status {body.get('status')!r}.")
    return body


def extract_records(page: dict[str, Any]) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for stream in page.get("data", {}).get("result", []):
        labels = stream.get("stream", {})
        for timestamp_ns, line in stream.get("values", []):
            timestamp = dt.datetime.fromtimestamp(
                int(timestamp_ns) / 1_000_000_000, tz=dt.timezone.utc
            )
            records.append(
                {
                    "timestamp_ns": str(timestamp_ns),
                    "timestamp_utc": timestamp.isoformat(),
                    "labels": labels,
                    "line": line,
                }
            )
    return records


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")
    os.chmod(path, 0o600)


def main() -> int:
    os.umask(0o077)
    config_path = load_runtime_config()
    args = parse_args()
    github_pat = os.environ.get("GITHUB_PAT", "")
    if not github_pat:
        raise ValueError("GITHUB_PAT is not configured; run $manual-invoice configure.")

    query = args.query_file.read_text(encoding="utf-8").strip()
    if not query:
        raise ValueError("LogQL query file is empty.")
    cached_token = os.environ.get("O11Y_AUTH0_TOKEN", "")
    cached_provider = os.environ.get("O11Y_PROVIDER", "")
    try:
        cached_expires = int(os.environ.get("O11Y_AUTH0_TOKEN_EXPIRES", "0"))
    except ValueError:
        cached_expires = 0
    if cached_token and cached_provider == args.provider and cached_expires > time.time() + 300:
        auth0_token = cached_token
        authentication = "cached"
    else:
        auth0_token = authenticate(github_pat, args.provider, args.timeout)
        authentication = "fresh"
    endpoint = (
        f"https://www.ds.{args.region}.{args.provider}.observability.tidbcloud.com"
        "/loki/self-monitoring/loki/api/v1"
    )

    start_ns = int(args.start_time.timestamp() * 1_000_000_000)
    end_ns = int(args.end_time.timestamp() * 1_000_000_000)
    cursor_ns = start_ns
    raw_pages: list[dict[str, Any]] = []
    records: list[dict[str, Any]] = []
    possibly_truncated = False
    reauth_attempted = False
    reauth_succeeded = False
    reauth_error: Exception | None = None

    for page_number in range(1, args.max_pages + 1):
        try:
            page = loki_page(
                endpoint,
                auth0_token,
                query,
                cursor_ns,
                end_ns,
                args.page_limit,
                args.timeout,
            )
        except LokiHTTPError as exc:
            if exc.status_code != 403 or reauth_attempted:
                raise
            reauth_attempted = True
            try:
                auth0_token, _, _ = get_auth0_token(
                    args.provider, args.timeout, force_refresh=True
                )
                reauth_succeeded = True
                authentication = "refreshed_after_403"
                page = loki_page(
                    endpoint,
                    auth0_token,
                    query,
                    cursor_ns,
                    end_ns,
                    args.page_limit,
                    args.timeout,
                )
            except Exception as retry_error:
                reauth_error = retry_error
                if isinstance(retry_error, LokiHTTPError) and retry_error.status_code != 403:
                    raise
                raise RuntimeError(o11y_403_message("Loki", reauth_error)) from retry_error
        page_records = extract_records(page)
        raw_pages.append(
            {
                "page": page_number,
                "start_ns": str(cursor_ns),
                "end_ns": str(end_ns),
                "response": page,
            }
        )
        records.extend(page_records)
        if len(page_records) < args.page_limit:
            break
        next_cursor = max(int(record["timestamp_ns"]) for record in page_records) + 1
        if next_cursor >= end_ns:
            break
        cursor_ns = next_cursor
    else:
        possibly_truncated = True

    unique: dict[tuple[str, str, str], dict[str, Any]] = {}
    for record in records:
        key = (
            record["timestamp_ns"],
            json.dumps(record["labels"], sort_keys=True, ensure_ascii=False),
            record["line"],
        )
        unique[key] = record
    normalized = sorted(unique.values(), key=lambda record: int(record["timestamp_ns"]))

    write_json(
        args.raw_output,
        {
            "query": query,
            "provider": args.provider,
            "region": args.region,
            "start_utc": args.start_time.isoformat(),
            "end_utc": args.end_time.isoformat(),
            "pages": raw_pages,
            "reauth_attempted_after_403": reauth_attempted,
            "reauth_succeeded_after_403": reauth_succeeded,
        },
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as handle:
        for record in normalized:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")
    os.chmod(args.output, 0o600)

    metadata = {
        "run_id": args.run_id,
        "source": "o11y_loki",
        "config_path": str(config_path) if config_path else None,
        "provider": args.provider,
        "region": args.region,
        "endpoint": endpoint,
        "authentication": authentication,
        "reauth_attempted_after_403": reauth_attempted,
        "reauth_succeeded_after_403": reauth_succeeded,
        "query_file": str(args.query_file),
        "query_sha256": hashlib.sha256(query.encode("utf-8")).hexdigest(),
        "start_utc": args.start_time.isoformat(),
        "end_utc": args.end_time.isoformat(),
        "duration_seconds": int((args.end_time - args.start_time).total_seconds()),
        "page_limit": args.page_limit,
        "page_count": len(raw_pages),
        "record_count": len(normalized),
        "possibly_truncated": possibly_truncated,
        "raw_output": str(args.raw_output),
        "normalized_output": str(args.output),
        "retrieved_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
    }
    write_json(args.metadata_output, metadata)
    print(json.dumps(metadata, ensure_ascii=False))
    return 0


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