#!/usr/bin/env python3
"""Run daily_prod_backend_log_report.py on the prod host via SSM and download results.

Usage:
    python3 scripts/run_report_on_prod.py --date 2026-06-10
    python3 scripts/run_report_on_prod.py --date 2026-06-11
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import subprocess
import sys
import time
from pathlib import Path


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"
OUTPUT_ROOT = Path(".tmp/latency-reports")


def run_aws(args: list[str]) -> str:
    result = subprocess.run(
        ["aws", *args],
        text=True,
        capture_output=True,
        check=False,
        env=os.environ.copy(),
    )
    if result.returncode != 0:
        msg = result.stderr.strip() or result.stdout.strip()
        print(f"ERROR: aws {' '.join(args)}\n{msg}", file=sys.stderr)
        sys.exit(1)
    return result.stdout


def resolve_instance_id(region: str, name: str) -> str:
    raw = run_aws([
        "ec2", "describe-instances",
        "--region", region,
        "--filters", "Name=instance-state-name,Values=running",
        f"Name=tag:Name,Values={name}",
        "--query", "Reservations[].Instances[].InstanceId",
        "--output", "text",
    ]).strip()
    ids = [x for x in raw.split() if x]
    if len(ids) != 1:
        print(f"ERROR: expected 1 instance named {name!r}, got {ids}", file=sys.stderr)
        sys.exit(1)
    return ids[0]


def send_command(region: str, instance_id: str, script: str) -> str:
    resp = run_aws([
        "ssm", "send-command",
        "--region", region,
        "--instance-ids", instance_id,
        "--document-name", "AWS-RunShellScript",
        "--parameters", json.dumps({"commands": [f"python3 -c {repr(script)}"]}),
        "--output", "json",
    ])
    data = json.loads(resp)
    return data["Command"]["CommandId"]


def wait_command(region: str, instance_id: str, command_id: str, poll: float = 2.0, timeout: float = 300.0) -> str:
    deadline = time.time() + timeout
    while time.time() < deadline:
        resp = run_aws([
            "ssm", "get-command-invocation",
            "--region", region,
            "--instance-id", instance_id,
            "--command-id", command_id,
            "--output", "json",
        ])
        data = json.loads(resp)
        status = data.get("Status", "")
        if status in ("Success",):
            return data.get("StandardOutputContent", "")
        if status in ("Failed", "TimedOut", "Cancelled"):
            err = data.get("StandardErrorContent", "")
            print(f"ERROR: SSM command {status}\n{err}", file=sys.stderr)
            sys.exit(1)
        time.sleep(poll)
    print("ERROR: timed out waiting for SSM command", file=sys.stderr)
    sys.exit(1)


def build_remote_script(log_path: str, report_date: str) -> str:
    return f"""\
import json, math, re, sys
from collections import Counter, defaultdict
from datetime import date, datetime, time, timedelta, timezone
from pathlib import Path
from urllib.parse import parse_qs, unquote, urlsplit
from zoneinfo import ZoneInfo
import base64, gzip

log_path = Path({log_path!r})
report_date_str = {report_date!r}
slow_threshold_ms = 10000.0

tz = ZoneInfo("Asia/Shanghai")
rd = date.fromisoformat(report_date_str)
local_start = datetime.combine(rd, time.min, tzinfo=tz)
local_end = local_start + timedelta(days=1)
utc_start = local_start.astimezone(timezone.utc)
utc_end = local_end.astimezone(timezone.utc)

ANSI_RE = re.compile(r"\\x1b\\[[0-9;]*m")
TIMESTAMP_RE = re.compile(r"^(?P<ts>\\d{{4}}-\\d{{2}}-\\d{{2}}T\\d{{2}}:\\d{{2}}:\\d{{2}}(?:\\.\\d+)?Z)\\b")
FIELD_RE = re.compile(
    r"\\b(?P<key>ip|latency|method|path|request_id|status|trace_id)="
    r"(?P<value>.*?)(?=\\s+(?:ip|latency|method|path|request_id|status|trace_id)=|$)"
)
STATIC_EXT_RE = re.compile(r"\\.(?:css|js|map|png|jpg|jpeg|gif|svg|ico|webp|woff|woff2|ttf|eot)(?:\\?|$)", re.IGNORECASE)
SALESFORCE_ID_RE = re.compile(r"\\b[0-9A-Za-z]{{15,18}}\\b")

def route_path(p):
    return urlsplit(p).path or p

def request_category(p):
    cp = route_path(p)
    if cp in ("/health", "/healthz", "/ready", "/version.json"):
        return "ignored"
    if STATIC_EXT_RE.search(cp):
        return "ignored"
    if cp.startswith("/services/data/") and "/connect/notifications" in cp:
        return "background"
    if cp.startswith("/webhooks/"):
        return "integration"
    return "interactive"

def normalize_route(p):
    cp = route_path(p)
    parts = [x for x in cp.split("/") if x]
    normalized = []
    for i, part in enumerate(parts):
        prev = parts[i-1] if i else ""
        if SALESFORCE_ID_RE.fullmatch(part):
            normalized.append(":id")
        elif prev == "sobjects" and i+1 < len(parts):
            normalized.append(part)
        elif prev not in ("services", "data") and i > 0 and SALESFORCE_ID_RE.search(part):
            normalized.append(":id")
        else:
            normalized.append(part)
    route = "/" + "/".join(normalized)
    route = re.sub(r"(/services/data/v[0-9.]+/sobjects/[^/]+)/[^/]+$", r"\\1/:id", route)
    route = re.sub(r"(/services/data/v[0-9.]+/ui-api/(?:object-info|layouts|record-ui)/[^/]+)/[^/]+", r"\\1/:id", route)
    return route

def compute_percentile(values, ratio):
    if not values:
        return 0.0
    ordered = sorted(values)
    index = max(0, math.ceil(ratio * len(ordered)) - 1)
    return float(ordered[index])

all_route_latencies = defaultdict(list)
all_latencies = []
slow_routes = defaultdict(list)
total_by_category = Counter()
slow_by_category = Counter()
parsed_requests = 0

with log_path.open("r", errors="ignore") as f:
    for line in f:
        if "HTTP Request" not in line:
            continue
        clean = ANSI_RE.sub("", line).strip()
        m = TIMESTAMP_RE.search(clean)
        if not m:
            continue
        ts = datetime.fromisoformat(m.group("ts").replace("Z", "+00:00"))
        if not (utc_start <= ts < utc_end):
            continue
        fields = {{mm.group("key"): mm.group("value").strip() for mm in FIELD_RE.finditer(clean)}}
        if not all(fields.get(k) for k in ("latency", "method", "path", "status")):
            continue
        try:
            latency_ms = float(fields["latency"])
        except ValueError:
            continue
        parsed_requests += 1
        path = fields["path"]
        category = request_category(path)
        total_by_category[category] += 1
        if category != "ignored":
            route = normalize_route(path)
            all_route_latencies[route].append(latency_ms)
            all_latencies.append(latency_ms)
        if latency_ms >= slow_threshold_ms:
            slow_by_category[category] += 1
            if category != "ignored":
                slow_routes[normalize_route(path)].append(latency_ms)

route_latency_profile = []
for route, latencies in all_route_latencies.items():
    slow_count = len(slow_routes.get(route, []))
    profile = {{
        "route": route,
        "request_count": len(latencies),
        "slow_count": slow_count,
        "slow_rate": round((slow_count / len(latencies)) * 100, 2) if latencies else 0.0,
        "avg_latency_ms": round(sum(latencies) / len(latencies), 2),
        "p50_latency_ms": round(compute_percentile(latencies, 0.50), 2),
        "p95_latency_ms": round(compute_percentile(latencies, 0.95), 2),
        "p99_latency_ms": round(compute_percentile(latencies, 0.99), 2),
        "max_latency_ms": round(max(latencies), 2),
    }}
    route_latency_profile.append(profile)
route_latency_profile.sort(key=lambda x: (-x["request_count"], -x["p95_latency_ms"], x["route"]))

parsed_total = sum(total_by_category.values())
slow_total = sum(slow_by_category.values())
interactive_total = total_by_category["interactive"]
interactive_slow = slow_by_category["interactive"]

report = {{
    "report_date": report_date_str,
    "timezone": "Asia/Shanghai",
    "local_window_start": local_start.isoformat(),
    "local_window_end": local_end.isoformat(),
    "summary": {{
        "parsed_requests_in_window": parsed_requests,
        "slow_threshold_ms": slow_threshold_ms,
        "slow_request_count": slow_total,
        "slow_request_rate": round((slow_total / parsed_total) * 100, 2) if parsed_total else 0.0,
        "interactive_request_count": interactive_total,
        "interactive_slow_request_count": interactive_slow,
        "interactive_slow_request_rate": round((interactive_slow / interactive_total) * 100, 2) if interactive_total else 0.0,
        "all_request_p50_ms": round(compute_percentile(all_latencies, 0.50), 2),
        "all_request_p95_ms": round(compute_percentile(all_latencies, 0.95), 2),
        "all_request_p99_ms": round(compute_percentile(all_latencies, 0.99), 2),
        "all_request_avg_ms": round(sum(all_latencies) / len(all_latencies), 2) if all_latencies else 0.0,
    }},
    "route_latency_profile": route_latency_profile,
}}

payload = json.dumps(report, ensure_ascii=False)
compressed = gzip.compress(payload.encode())
encoded = base64.b64encode(compressed).decode()
print("===REPORT_START===")
print(encoded)
print("===REPORT_END===")
"""


def main() -> int:
    parser = argparse.ArgumentParser(description="Run latency report on prod via SSM.")
    parser.add_argument("--date", required=True, help="Report date YYYY-MM-DD (Asia/Shanghai)")
    parser.add_argument("--region", default=DEFAULT_REGION)
    parser.add_argument("--instance-name", default=DEFAULT_INSTANCE_NAME)
    parser.add_argument("--instance-id", default="")
    parser.add_argument("--log-path", default=DEFAULT_LOG_PATH)
    args = parser.parse_args()

    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
    output_file = OUTPUT_ROOT / f"latency-report-{args.date}.json"

    if output_file.exists():
        print(f"Report already exists: {output_file}")
        return 0

    print(f"Resolving instance...")
    instance_id = args.instance_id or resolve_instance_id(args.region, args.instance_name)
    print(f"Instance: {instance_id}")

    script = build_remote_script(args.log_path, args.date)
    print(f"Sending report script for {args.date}...")
    command_id = send_command(args.region, instance_id, script)
    print(f"Command: {command_id}, waiting...")

    output = wait_command(args.region, instance_id, command_id, timeout=600)

    start_marker = "===REPORT_START==="
    end_marker = "===REPORT_END==="
    if start_marker not in output or end_marker not in output:
        print(f"ERROR: markers not found in output. Output:\n{output[:2000]}", file=sys.stderr)
        return 1

    encoded = output.split(start_marker)[1].split(end_marker)[0].strip()
    import gzip as gz
    payload = gz.decompress(base64.b64decode(encoded)).decode()
    report = json.loads(payload)

    output_file.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"Report saved: {output_file}")
    print(f"  requests: {report['summary']['parsed_requests_in_window']}")
    print(f"  slow_rate: {report['summary']['slow_request_rate']}%")
    print(f"  all_p95: {report['summary']['all_request_p95_ms']}ms")
    print(f"  routes: {len(report['route_latency_profile'])}")
    return 0


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