#!/usr/bin/env python3
"""E2E test for the _prompt_lock keepalive behaviour.

Usage:
    python3 tests/e2e_prompt_lock.py --base-url https://investibook.tidbcloud.com --token YOUR_TOKEN --workspace-id ws-xxx

What it does:
  1. Sends a long-running prompt (request A) that triggers tool calls.
  2. Immediately sends a second prompt (request B) to the same workspace.
  3. Monitors request B's SSE stream for keepalive comments.
  4. Reports whether keepalives were received and whether B eventually completed.

Expected outcome (after the fix):
  - Request B receives `: keepalive` comments every ~15s while waiting.
  - Once A finishes, B proceeds and gets a normal response.
  - If A takes >10 min, B gets an "Agent is busy" error (not a silent hang).

Before the fix:
  - Request B would receive NO keepalives while waiting for the lock.
  - ALB would kill the connection after ~60s idle, causing a silent hang.
"""
from __future__ import annotations

import argparse
import asyncio
import sys
import time


async def stream_sse(url: str, label: str, results: dict):
    """Connect to an SSE endpoint and collect events."""
    import httpx

    events = []
    keepalives = 0
    first_event_at = None
    start = time.monotonic()

    print(f"[{label}] Connecting...")

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(700, connect=30)) as client:
            async with client.stream("GET", url) as resp:
                print(f"[{label}] Connected (status={resp.status_code})")
                async for line in resp.aiter_lines():
                    elapsed = time.monotonic() - start
                    if line.strip() == ": keepalive":
                        keepalives += 1
                        print(f"[{label}] +{elapsed:.1f}s keepalive #{keepalives}")
                    elif line.startswith("data: "):
                        import json
                        try:
                            event = json.loads(line[6:])
                            etype = event.get("type", "?")
                            if first_event_at is None:
                                first_event_at = elapsed
                            events.append(event)
                            if etype == "text":
                                text = event.get("text", "")[:60]
                                print(f"[{label}] +{elapsed:.1f}s text: {text}")
                            elif etype == "tool":
                                title = event.get("title", "")
                                status = event.get("status", "")
                                print(f"[{label}] +{elapsed:.1f}s tool: {title} ({status})")
                            elif etype == "complete":
                                print(f"[{label}] +{elapsed:.1f}s COMPLETE")
                                break
                            elif etype == "error":
                                print(f"[{label}] +{elapsed:.1f}s ERROR: {event.get('error','')[:100]}")
                            else:
                                print(f"[{label}] +{elapsed:.1f}s {etype}")
                        except Exception:
                            pass
    except Exception as e:
        elapsed = time.monotonic() - start
        print(f"[{label}] +{elapsed:.1f}s EXCEPTION: {e}")

    results[label] = {
        "keepalives": keepalives,
        "events": len(events),
        "first_event_at": first_event_at,
        "duration": time.monotonic() - start,
        "types": [e.get("type") for e in events],
    }


async def main():
    parser = argparse.ArgumentParser(description="E2E test for prompt lock keepalive")
    parser.add_argument("--base-url", required=True, help="e.g. https://investibook.tidbcloud.com")
    parser.add_argument("--token", required=True, help="Auth token")
    parser.add_argument("--workspace-id", required=True, help="Workspace ID to test against")
    parser.add_argument(
        "--message-a",
        default="帮我分析一下当前集群的整体状态",
        help="Long-running prompt for request A",
    )
    parser.add_argument(
        "--message-b",
        default="hello",
        help="Short prompt for request B (sent while A is running)",
    )
    parser.add_argument(
        "--delay", type=float, default=3.0,
        help="Seconds to wait after A starts before sending B (default: 3)",
    )
    args = parser.parse_args()

    from urllib.parse import urlencode

    base = args.base_url.rstrip("/")

    params_a = urlencode({
        "message": args.message_a,
        "workspace_id": args.workspace_id,
        "token": args.token,
    })
    params_b = urlencode({
        "message": args.message_b,
        "workspace_id": args.workspace_id,
        "token": args.token,
    })

    url_a = f"{base}/api/chat/stream?{params_a}"
    url_b = f"{base}/api/chat/stream?{params_b}"

    results: dict = {}

    print("=" * 60)
    print("E2E Prompt Lock Test")
    print("=" * 60)
    print(f"Workspace: {args.workspace_id}")
    print(f"Request A: {args.message_a[:50]}...")
    print(f"Request B: {args.message_b}")
    print(f"Delay before B: {args.delay}s")
    print("=" * 60)

    # Start A
    task_a = asyncio.create_task(stream_sse(url_a, "A", results))

    # Wait a bit, then start B
    await asyncio.sleep(args.delay)
    task_b = asyncio.create_task(stream_sse(url_b, "B", results))

    # Wait for both
    await asyncio.gather(task_a, task_b)

    # Report
    print()
    print("=" * 60)
    print("RESULTS")
    print("=" * 60)

    for label in ["A", "B"]:
        r = results.get(label, {})
        print(f"\n[{label}]")
        print(f"  Duration:    {r.get('duration', 0):.1f}s")
        print(f"  Keepalives:  {r.get('keepalives', 0)}")
        print(f"  Events:      {r.get('events', 0)}")
        print(f"  Event types: {r.get('types', [])[:10]}")

    b = results.get("B", {})
    print()
    if b.get("keepalives", 0) > 0:
        print("PASS: Request B received keepalives while waiting for the lock.")
    else:
        print("WARN: Request B received NO keepalives (lock may not have been contended).")

    if "error" in b.get("types", []):
        errors = [e for e in b.get("types", []) if e == "error"]
        print(f"INFO: Request B received {len(errors)} error event(s).")
    if "complete" in b.get("types", []):
        print("PASS: Request B completed (not a silent hang).")
    else:
        print("FAIL: Request B did NOT complete (possible hang).")


if __name__ == "__main__":
    asyncio.run(main())