#!/usr/bin/env python3
"""Spotify API rate limit test — compare rate budgets across Spotify apps.

Reads refresh tokens from a JSON file (flat array of strings or newline-delimited),
refreshes each token, calls Spotify API endpoints, and logs per-request metrics.

Usage:
    export SPOTIFY_CLIENT_ID="..."
    export SPOTIFY_CLIENT_SECRET="..."
    python3 scripts/rate_test.py --file ~/Downloads/spotify-smf/spotify.json --limit 100

    # With concurrency:
    python3 scripts/rate_test.py --file tokens.json --limit 200 --workers 5

    # Read from S3:
    python3 scripts/rate_test.py --s3 s3://dev-mymac80/resonance-engine/smf-tokens/spotify.json --limit 100
"""

import argparse
import base64
import csv
import json
import logging
import os
import random
import statistics
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

import requests

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)-7s %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Metrics collection
# ---------------------------------------------------------------------------

class RequestMetric:
    __slots__ = (
        "seq", "fan_index", "endpoint", "status", "latency_ms",
        "retry_after", "timestamp", "attempt", "error",
    )

    def __init__(self, seq, fan_index, endpoint, status, latency_ms,
                 retry_after=None, timestamp=None, attempt=0, error=None):
        self.seq = seq
        self.fan_index = fan_index
        self.endpoint = endpoint
        self.status = status
        self.latency_ms = latency_ms
        self.retry_after = retry_after
        self.timestamp = timestamp or datetime.now(timezone.utc).isoformat()
        self.attempt = attempt
        self.error = error

    def as_dict(self):
        return {s: getattr(self, s) for s in self.__slots__}


# Global metrics list (thread-safe via GIL for list.append)
_metrics: list[RequestMetric] = []
_seq_counter = 0


def _next_seq():
    global _seq_counter
    _seq_counter += 1
    return _seq_counter


# ---------------------------------------------------------------------------
# Token loading
# ---------------------------------------------------------------------------

def load_tokens_local(path, limit):
    """Load tokens from a local JSON file (flat array of strings)."""
    logger.info("Loading tokens from %s (limit=%s)", path, limit)
    tokens = []
    with open(path) as f:
        for line in f:
            line = line.strip().strip(",")
            if line.startswith('"AQ'):
                token = line.strip('"')
                tokens.append(token)
                if limit and len(tokens) >= limit:
                    break
    logger.info("Loaded %d tokens", len(tokens))
    return tokens


def load_tokens_s3(s3_uri, limit):
    """Load tokens from an S3 JSON file."""
    import boto3
    parts = s3_uri.replace("s3://", "").split("/", 1)
    bucket, key = parts[0], parts[1]
    logger.info("Loading tokens from s3://%s/%s (limit=%s)", bucket, key, limit)

    s3 = boto3.client("s3")
    # Stream the file to avoid loading 1.5GB into memory
    resp = s3.get_object(Bucket=bucket, Key=key)
    tokens = []
    for raw_line in resp["Body"].iter_lines():
        line = raw_line.decode("utf-8").strip().strip(",")
        if line.startswith('"AQ'):
            token = line.strip('"')
            tokens.append(token)
            if limit and len(tokens) >= limit:
                break
    logger.info("Loaded %d tokens", len(tokens))
    return tokens


# ---------------------------------------------------------------------------
# Spotify API (ported from collector-worker/handler.py)
# ---------------------------------------------------------------------------

def refresh_token(token, client_id, client_secret):
    """Refresh a Spotify OAuth token. Returns (access_token, new_refresh_token, metric)."""
    auth_b64 = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
    start = time.monotonic()
    metric = None
    try:
        resp = requests.post(
            "https://accounts.spotify.com/api/token",
            data={"grant_type": "refresh_token", "refresh_token": token},
            headers={
                "Authorization": f"Basic {auth_b64}",
                "Content-Type": "application/x-www-form-urlencoded",
            },
            timeout=10,
        )
        latency = (time.monotonic() - start) * 1000
        metric = RequestMetric(
            seq=_next_seq(), fan_index=-1, endpoint="token_refresh",
            status=resp.status_code, latency_ms=round(latency, 1),
        )
        _metrics.append(metric)

        if resp.status_code != 200:
            error_body = ""
            try:
                error_body = resp.json().get("error_description", resp.text[:100])
            except Exception:
                error_body = resp.text[:100]
            metric.error = error_body
            return None, None, metric

        data = resp.json()
        return data["access_token"], data.get("refresh_token"), metric

    except requests.RequestException as exc:
        latency = (time.monotonic() - start) * 1000
        metric = RequestMetric(
            seq=_next_seq(), fan_index=-1, endpoint="token_refresh",
            status=0, latency_ms=round(latency, 1), error=str(exc),
        )
        _metrics.append(metric)
        return None, None, metric


def call_spotify_endpoint(access_token, endpoint, params=None, fan_index=0):
    """Call a Spotify API endpoint with retry on 429. Returns (data, metric)."""
    url = f"https://api.spotify.com/v1{endpoint}"
    headers = {"Authorization": f"Bearer {access_token}"}
    max_retries = 3

    for attempt in range(max_retries + 1):
        start = time.monotonic()
        try:
            resp = requests.get(url, headers=headers, params=params, timeout=30)
        except requests.RequestException as exc:
            latency = (time.monotonic() - start) * 1000
            m = RequestMetric(
                seq=_next_seq(), fan_index=fan_index, endpoint=endpoint,
                status=0, latency_ms=round(latency, 1), attempt=attempt,
                error=str(exc),
            )
            _metrics.append(m)
            return None, m

        latency = (time.monotonic() - start) * 1000
        retry_after = None

        if resp.status_code == 429:
            try:
                retry_after = int(resp.headers.get("Retry-After", 5))
            except (ValueError, TypeError):
                retry_after = 5

        m = RequestMetric(
            seq=_next_seq(), fan_index=fan_index, endpoint=endpoint,
            status=resp.status_code, latency_ms=round(latency, 1),
            retry_after=retry_after, attempt=attempt,
        )
        _metrics.append(m)

        if resp.status_code == 429:
            if attempt == max_retries:
                m.error = "rate_limit_exhausted"
                return None, m
            sleep_time = min(retry_after * (2 ** attempt), 30) + random.uniform(0, 1)
            logger.warning(
                "429 on %s attempt %d/%d, Retry-After=%ds, sleeping %.1fs",
                endpoint, attempt + 1, max_retries + 1, retry_after, sleep_time,
            )
            time.sleep(sleep_time)
            continue

        if resp.status_code != 200:
            m.error = f"HTTP {resp.status_code}"
            return None, m

        return resp.json(), m

    return None, None


# ---------------------------------------------------------------------------
# Fan processing
# ---------------------------------------------------------------------------

ENDPOINTS = [
    ("/me/top/artists", {"limit": "50", "time_range": "medium_term"}),
    ("/me/player/recently-played", {"limit": "50"}),
]


def process_fan(fan_index, token, client_id, client_secret):
    """Process a single fan: refresh token, call endpoints."""
    access_token, new_token, refresh_metric = refresh_token(
        token, client_id, client_secret
    )
    if refresh_metric:
        refresh_metric.fan_index = fan_index

    if not access_token:
        logger.warning("Fan %d: token refresh failed — %s", fan_index,
                        refresh_metric.error if refresh_metric else "unknown")
        return {"fan_index": fan_index, "success": False, "error": "token_refresh_failed"}

    results = {"fan_index": fan_index, "success": True, "endpoints": {}}
    for ep, params in ENDPOINTS:
        data, metric = call_spotify_endpoint(access_token, ep, params, fan_index)
        if data:
            item_count = len(data.get("items", []))
            results["endpoints"][ep] = item_count
        else:
            results["endpoints"][ep] = None
            if metric and metric.status == 401:
                logger.warning("Fan %d: 401 on %s — skipping remaining endpoints", fan_index, ep)
                break

    return results


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

def print_summary():
    """Print rate limit analysis summary."""
    if not _metrics:
        print("\nNo metrics collected.")
        return

    total = len(_metrics)
    successes = [m for m in _metrics if m.status == 200]
    rate_limited = [m for m in _metrics if m.status == 429]
    errors = [m for m in _metrics if m.status not in (200, 429) and m.status != 0]
    network_errors = [m for m in _metrics if m.status == 0]
    token_failures = [m for m in _metrics if m.endpoint == "token_refresh" and m.status != 200]

    # Timing
    all_latencies = [m.latency_ms for m in _metrics if m.status in (200, 429)]
    success_latencies = [m.latency_ms for m in successes]

    # Rate limit details
    retry_afters = [m.retry_after for m in rate_limited if m.retry_after is not None]

    # Throughput (wall clock)
    timestamps = [m.timestamp for m in _metrics]
    if len(timestamps) >= 2:
        first = datetime.fromisoformat(timestamps[0])
        last = datetime.fromisoformat(timestamps[-1])
        wall_seconds = max((last - first).total_seconds(), 0.001)
    else:
        wall_seconds = 0.001

    # Per-endpoint breakdown
    api_endpoints = set(m.endpoint for m in _metrics if m.endpoint != "token_refresh")

    print("\n" + "=" * 70)
    print("RATE TEST SUMMARY")
    print("=" * 70)
    print(f"Total requests:     {total}")
    print(f"  Successful (200): {len(successes)}")
    print(f"  Rate limited:     {len(rate_limited)}")
    print(f"  Other errors:     {len(errors)}")
    print(f"  Network errors:   {len(network_errors)}")
    print(f"  Token failures:   {len(token_failures)}")
    print(f"  429 rate:         {len(rate_limited)/max(total,1)*100:.1f}%")
    print()

    if success_latencies:
        print("Latency (successful requests):")
        print(f"  Mean:   {statistics.mean(success_latencies):.0f} ms")
        print(f"  Median: {statistics.median(success_latencies):.0f} ms")
        if len(success_latencies) >= 20:
            sorted_lat = sorted(success_latencies)
            p95_idx = int(len(sorted_lat) * 0.95)
            print(f"  P95:    {sorted_lat[p95_idx]:.0f} ms")
        print()

    if retry_afters:
        print("Retry-After (429 responses):")
        print(f"  Mean:   {statistics.mean(retry_afters):.1f}s")
        print(f"  Median: {statistics.median(retry_afters):.1f}s")
        print(f"  Min:    {min(retry_afters)}s")
        print(f"  Max:    {max(retry_afters)}s")
        print()

    print(f"Wall clock time:    {wall_seconds:.1f}s")
    print(f"Throughput:         {len(successes)/wall_seconds:.1f} successful req/s")
    fan_successes = len(set(m.fan_index for m in successes if m.endpoint != "token_refresh"))
    print(f"Fan throughput:     {fan_successes/wall_seconds:.1f} fans/s")
    print()

    for ep in sorted(api_endpoints):
        ep_metrics = [m for m in _metrics if m.endpoint == ep]
        ep_429 = [m for m in ep_metrics if m.status == 429]
        ep_200 = [m for m in ep_metrics if m.status == 200]
        print(f"  {ep}: {len(ep_200)} ok, {len(ep_429)} rate-limited")

    print("=" * 70)


def export_metrics_csv(path):
    """Export per-request metrics to CSV."""
    with open(path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=RequestMetric.__slots__)
        writer.writeheader()
        for m in _metrics:
            writer.writerow(m.as_dict())
    logger.info("Metrics exported to %s (%d rows)", path, len(_metrics))


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Spotify API rate limit test")
    parser.add_argument("--file", help="Local path to token JSON file")
    parser.add_argument("--s3", help="S3 URI to token JSON file")
    parser.add_argument("--limit", type=int, default=100,
                        help="Max number of tokens to process (default: 100)")
    parser.add_argument("--workers", type=int, default=1,
                        help="Concurrent workers (default: 1)")
    parser.add_argument("--csv", help="Export per-request metrics to CSV file")
    parser.add_argument("--offset", type=int, default=0,
                        help="Skip first N tokens (default: 0)")
    args = parser.parse_args()

    client_id = os.environ.get("SPOTIFY_CLIENT_ID")
    client_secret = os.environ.get("SPOTIFY_CLIENT_SECRET")
    if not client_id or not client_secret:
        print("ERROR: Set SPOTIFY_CLIENT_ID and SPOTIFY_CLIENT_SECRET env vars")
        sys.exit(1)

    # Load tokens
    if args.file:
        tokens = load_tokens_local(args.file, args.limit + args.offset)
    elif args.s3:
        tokens = load_tokens_s3(args.s3, args.limit + args.offset)
    else:
        print("ERROR: Provide --file or --s3")
        sys.exit(1)

    if args.offset:
        tokens = tokens[args.offset:]
    tokens = tokens[:args.limit]

    if not tokens:
        print("No tokens loaded.")
        sys.exit(1)

    logger.info("Processing %d tokens with %d worker(s)", len(tokens), args.workers)
    start_time = time.monotonic()

    if args.workers == 1:
        for i, token in enumerate(tokens):
            result = process_fan(i, token, client_id, client_secret)
            status = "OK" if result["success"] else f"FAIL ({result.get('error', '?')})"
            logger.info("Fan %d/%d: %s", i + 1, len(tokens), status)
    else:
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = {
                executor.submit(process_fan, i, token, client_id, client_secret): i
                for i, token in enumerate(tokens)
            }
            done = 0
            for future in as_completed(futures):
                done += 1
                result = future.result()
                status = "OK" if result["success"] else f"FAIL ({result.get('error', '?')})"
                logger.info("Fan %d/%d: %s", done, len(tokens), status)

    elapsed = time.monotonic() - start_time
    logger.info("Completed in %.1fs", elapsed)

    print_summary()

    if args.csv:
        export_metrics_csv(args.csv)


if __name__ == "__main__":
    main()