#!/usr/bin/env python3
from __future__ import annotations

import argparse
import base64
import gzip
import json
import math
import os
import re
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


DEFAULT_REGION = "us-west-2"
DEFAULT_INSTANCE_NAME = "nexuscrm-prod-runtime-1"
DEFAULT_LOG_PATH = "/home/ubuntu/nexuscrm-deployment/online/envs/multi_prod/current/logs/backend.stderr.log"
DEFAULT_OUTPUT_ROOT = ".tmp/prod-slow-logs"


class DownloadError(RuntimeError):
    pass


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Download slow HTTP request entries from the production multi_prod backend business log via SSM."
    )
    parser.add_argument("--region", default=DEFAULT_REGION)
    parser.add_argument("--instance-id", default="", help="Prod EC2 instance id. Overrides --instance-name.")
    parser.add_argument("--instance-name", default=DEFAULT_INSTANCE_NAME)
    parser.add_argument("--log-path", default=DEFAULT_LOG_PATH)
    parser.add_argument(
        "--date",
        default=datetime.now(timezone.utc).strftime("%Y-%m-%d"),
        help="UTC log date, YYYY-MM-DD.",
    )
    parser.add_argument("--threshold-ms", type=float, default=3000.0)
    parser.add_argument("--output-dir", default="", help="Default: .tmp/prod-slow-logs/YYYYMMDD")
    parser.add_argument("--chunk-size", type=int, default=6000)
    parser.add_argument("--poll-seconds", type=float, default=1.0)
    parser.add_argument("--top", type=int, default=10)
    return parser.parse_args()


def run_aws(args: list[str], *, capture: bool = True) -> str:
    result = subprocess.run(
        ["aws", *args],
        text=True,
        capture_output=True,
        check=False,
        env=os.environ.copy(),
    )
    if result.returncode != 0:
        message = result.stderr.strip() or result.stdout.strip()
        raise DownloadError(f"aws command failed: aws {' '.join(args)}\n{message}")
    return result.stdout if capture else ""


def resolve_instance_id(region: str, instance_id: str, instance_name: str) -> str:
    if instance_id:
        return instance_id
    query = "Reservations[].Instances[].InstanceId"
    raw = run_aws(
        [
            "ec2",
            "describe-instances",
            "--region",
            region,
            "--filters",
            "Name=instance-state-name,Values=running",
            f"Name=tag:Name,Values={instance_name}",
            "--query",
            query,
            "--output",
            "text",
        ]
    ).strip()
    ids = [item for item in raw.split() if item]
    if len(ids) != 1:
        raise DownloadError(f"expected exactly one running instance named {instance_name!r}, got {ids}")
    return ids[0]


def remote_extract_script(log_path: str, date_prefix: str, threshold_ms: float, chunk_size: int, top_count: int) -> str:
    # Keep this script self-contained because SSM runs it on the prod host.
    return f"""\
import base64
import gzip
import json
import math
import re
from pathlib import Path

log = Path({log_path!r})
date_prefix = {date_prefix!r}
threshold_ms = {threshold_ms!r}
chunk_size = {chunk_size!r}
top_count = {top_count!r}
stem = date_prefix.replace('-', '')
jsonl_path = Path(f'/tmp/prod-slow-http-{{stem}}.jsonl')
gz_path = Path(f'/tmp/prod-slow-http-{{stem}}.jsonl.gz')
b64_path = Path(f'/tmp/prod-slow-http-{{stem}}.jsonl.gz.b64')
parts_path = Path(f'/tmp/prod-slow-http-{{stem}}.jsonl.gz.b64parts')
meta_path = Path(f'/tmp/prod-slow-http-{{stem}}.metadata.json')
ansi = re.compile(r'\\x1b\\[[0-9;]*m')
kv = re.compile(r'([a-zA-Z_]+)=\\s*([^\\s]+)')
rows = []
with log.open('r', errors='replace') as f:
    for line in f:
        if 'HTTP Request' not in line or f'{{date_prefix}}T' not in line:
            continue
        clean = ansi.sub('', line).strip()
        data = dict(kv.findall(clean))
        try:
            latency = float(data.get('latency', '0'))
        except ValueError:
            latency = 0.0
        if latency < threshold_ms:
            continue
        rows.append({{
            'ts': clean[:20],
            'latency_ms': round(latency, 3),
            'method': data.get('method', ''),
            'path': data.get('path', ''),
            'status': data.get('status', ''),
            'request_id': data.get('request_id', ''),
            'trace_id': data.get('trace_id', ''),
            'ip': data.get('ip', ''),
        }})
rows.sort(key=lambda r: r['latency_ms'], reverse=True)
with jsonl_path.open('w') as f:
    for row in rows:
        f.write(json.dumps(row, ensure_ascii=False) + '\\n')
with gzip.open(gz_path, 'wb') as f:
    f.write(jsonl_path.read_bytes())
encoded = base64.b64encode(gz_path.read_bytes()).decode('ascii')
b64_path.write_text(encoded)
part_count = int(math.ceil(len(encoded) / chunk_size)) if encoded else 0
with parts_path.open('w') as f:
    for index in range(part_count):
        f.write(encoded[index * chunk_size:(index + 1) * chunk_size] + '\\n')
metadata = {{
    'date': date_prefix,
    'threshold_ms': threshold_ms,
    'log_path': str(log),
    'jsonl_path': str(jsonl_path),
    'gzip_path': str(gz_path),
    'parts_path': str(parts_path),
    'row_count': len(rows),
    'part_count': part_count,
    'jsonl_bytes': jsonl_path.stat().st_size,
    'gzip_bytes': gz_path.stat().st_size,
    'base64_bytes': len(encoded),
    'top': rows[:top_count],
}}
meta_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + '\\n')
print(json.dumps(metadata, ensure_ascii=False))
"""


def send_command(region: str, instance_id: str, command: str) -> str:
    payload = json.dumps({"commands": [command]})
    return run_aws(
        [
            "ssm",
            "send-command",
            "--region",
            region,
            "--instance-ids",
            instance_id,
            "--document-name",
            "AWS-RunShellScript",
            "--parameters",
            payload,
            "--query",
            "Command.CommandId",
            "--output",
            "text",
        ]
    ).strip()


def wait_command(region: str, instance_id: str, command_id: str, poll_seconds: float) -> dict[str, Any]:
    while True:
        raw = run_aws(
            [
                "ssm",
                "get-command-invocation",
                "--region",
                region,
                "--command-id",
                command_id,
                "--instance-id",
                instance_id,
                "--output",
                "json",
            ]
        )
        payload = json.loads(raw)
        status = payload.get("Status")
        if status in {"Success", "Failed", "Cancelled", "TimedOut", "Cancelling"}:
            if status != "Success":
                raise DownloadError(
                    f"ssm command {command_id} ended with {status}: {payload.get('StandardErrorContent', '')}"
                )
            return payload
        time.sleep(poll_seconds)


def run_remote_extract(args: argparse.Namespace, instance_id: str) -> dict[str, Any]:
    script = remote_extract_script(args.log_path, args.date, args.threshold_ms, args.chunk_size, args.top)
    encoded = base64.b64encode(script.encode("utf-8")).decode("ascii")
    command = (
        "bash -lc "
        + json.dumps(
            "set -euo pipefail; "
            f"printf %s {encoded} | base64 -d > /tmp/extract_prod_slow_http.py; "
            "python3 /tmp/extract_prod_slow_http.py"
        )
    )
    command_id = send_command(args.region, instance_id, command)
    result = wait_command(args.region, instance_id, command_id, args.poll_seconds)
    stdout = str(result.get("StandardOutputContent") or "").strip()
    if not stdout:
        raise DownloadError("remote extractor produced no output")
    return json.loads(stdout.splitlines()[-1])


def fetch_remote_part(region: str, instance_id: str, remote_parts_path: str, index: int, poll_seconds: float) -> str:
    command_id = send_command(region, instance_id, f"sed -n {index}p {remote_parts_path}")
    result = wait_command(region, instance_id, command_id, poll_seconds)
    return str(result.get("StandardOutputContent") or "").strip()


def write_download(output_dir: Path, metadata: dict[str, Any], parts: list[str]) -> tuple[Path, Path, Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    stem = str(metadata["date"]).replace("-", "")
    b64_path = output_dir / f"prod-slow-http-{stem}.jsonl.gz.b64"
    gz_path = output_dir / f"prod-slow-http-{stem}.jsonl.gz"
    jsonl_path = output_dir / f"prod-slow-http-{stem}.jsonl"
    metadata_path = output_dir / f"prod-slow-http-{stem}.metadata.json"
    b64_path.write_text("".join(parts))
    gz_bytes = base64.b64decode(b64_path.read_text(), validate=True)
    gz_path.write_bytes(gz_bytes)
    jsonl_path.write_bytes(gzip.decompress(gz_bytes))
    metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n")
    return jsonl_path, gz_path, metadata_path


def main() -> int:
    args = parse_args()
    if args.chunk_size < 1000:
        raise DownloadError("--chunk-size must be >= 1000")
    if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", args.date):
        raise DownloadError("--date must be YYYY-MM-DD")

    instance_id = resolve_instance_id(args.region, args.instance_id, args.instance_name)
    print(f"[download] instance_id={instance_id}", flush=True)
    metadata = run_remote_extract(args, instance_id)
    print(
        "[download] remote rows={row_count} gzip_bytes={gzip_bytes} part_count={part_count}".format(**metadata),
        flush=True,
    )
    parts = []
    for index in range(1, int(metadata["part_count"]) + 1):
        parts.append(fetch_remote_part(args.region, instance_id, str(metadata["parts_path"]), index, args.poll_seconds))
        print(f"[download] fetched_part={index}/{metadata['part_count']}", flush=True)

    output_dir = Path(args.output_dir or Path(DEFAULT_OUTPUT_ROOT) / args.date.replace("-", ""))
    jsonl_path, gz_path, metadata_path = write_download(output_dir, metadata, parts)
    row_count = sum(1 for _ in jsonl_path.open())
    if row_count != int(metadata["row_count"]):
        raise DownloadError(f"downloaded row count mismatch: local={row_count} remote={metadata['row_count']}")
    print(f"[download] wrote_jsonl={jsonl_path}")
    print(f"[download] wrote_gzip={gz_path}")
    print(f"[download] wrote_metadata={metadata_path}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except DownloadError as exc:
        print(f"error: {exc}", file=sys.stderr)
        raise SystemExit(1)