#!/usr/bin/env python3
"""Async passive XHR/fetch capture using Playwright's async API.

This script mirrors the sync version but uses async/await to avoid blocking
the event loop. It attaches listeners before navigating and prints a small
heartbeat that shows how many responses we've captured so you can see
activity live.

Usage:
  python scrapers/capture_passive_playwright_async.py --output-dir xhr_captured_async --continuous
"""
import argparse
import asyncio
import json
import mimetypes
import hashlib
import os
from datetime import datetime, timezone
from pathlib import Path

from dotenv import load_dotenv
load_dotenv()

try:
    from playwright.async_api import async_playwright
except Exception:
    raise


def sanitize(s: str) -> str:
    out = []
    for ch in s:
        if ch.isalnum() or ch in ('-', '_', '.'):
            out.append(ch)
        else:
            out.append('_')
    name = ''.join(out)
    if len(name) > 180:
        name = name[:120] + '_' + hashlib.sha1(s.encode('utf-8')).hexdigest()[:8]
    return name


def guess_ext(ctype: str) -> str:
    if not ctype:
        return 'bin'
    mt = ctype.split(';')[0].strip()
    ext = mimetypes.guess_extension(mt)
    if ext:
        return ext.lstrip('.')
    if mt.startswith('image/'):
        return mt.split('/')[1]
    if mt == 'application/json':
        return 'json'
    if mt == 'text/html':
        return 'html'
    if mt.startswith('text/'):
        return 'txt'
    return 'bin'


def now_ts():
    return datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')


def save_text(path: Path, text: str):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        f.write(text)


def save_binary(path: Path, data: bytes):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'wb') as f:
        f.write(data)


async def run(args):
    output = Path(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    index_path = output / 'capture_index.jsonl'

    capture_count = 0
    capture_count_lock = asyncio.Lock()

    async def heartbeat():
        while True:
            await asyncio.sleep(5)
            async with capture_count_lock:
                print(f"[heartbeat] captured={capture_count}")

    async with async_playwright() as p:
        browser_type = p.chromium
        # Launch persistent context
        if args.channel:
            try:
                context = await browser_type.launch_persistent_context(user_data_dir=args.profile, headless=args.headless, channel=args.channel)
            except Exception:
                print(f'Launch with channel {args.channel} failed; falling back to bundled Chromium')
                context = await browser_type.launch_persistent_context(user_data_dir=args.profile, headless=args.headless)
        else:
            context = await browser_type.launch_persistent_context(user_data_dir=args.profile, headless=args.headless)

        page = await context.new_page()

        async def on_response(resp):
            nonlocal capture_count
            try:
                req = resp.request
                rtype = (req.resource_type or '').lower()

                # write quick debug line
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tRESP\t{resp.status}\t{rtype}\t{req.url}\n")
                except Exception:
                    pass

                # determine response content-type early so we can allow images
                resp_headers = dict(resp.headers)
                ctype = (resp_headers.get('content-type') or '').lower()

                allow = False
                if not args.capture_types:
                    allow = True
                elif rtype in args.capture_types:
                    allow = True
                # Heuristic: always capture known document endpoints
                if '/document/content/' in req.url or '/document/summary/' in req.url:
                    allow = True
                # Always capture images regardless of capture_types
                if ctype.startswith('image/'):
                    allow = True
                if not allow:
                    try:
                        with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                            dbg.write(f"{now_ts()}\tSKIP\t{rtype}\t{req.url}\t{ctype}\n")
                    except Exception:
                        pass
                    return

                url = req.url
                ts = now_ts()
                safe = sanitize(url.replace('https://', '').replace('http://', ''))
                h = hashlib.sha1((url + ts).encode('utf-8')).hexdigest()[:8]

                body_path = None
                excerpt = None

                try:
                    if 'application/json' in ctype or ctype.startswith('text/') or 'javascript' in ctype:
                        try:
                            txt = await resp.text()
                        except asyncio.CancelledError:
                            return
                        except Exception as e:
                            if 'Target page, context or browser has been closed' in str(e):
                                return
                            raise
                        # try pretty JSON
                        try:
                            parsed = json.loads(txt)
                            pretty = json.dumps(parsed, ensure_ascii=False, indent=2)
                            fname = f"{safe}_{h}_{ts}.json"
                            body_path = str(output / fname)
                            save_text(output / fname, pretty)
                            excerpt = pretty[:800]
                        except Exception:
                            fname = f"{safe}_{h}_{ts}.txt"
                            body_path = str(output / fname)
                            save_text(output / fname, txt)
                            excerpt = txt[:800]
                    else:
                        try:
                            data = await resp.body()
                        except asyncio.CancelledError:
                            return
                        except Exception as e:
                            if 'Target page, context or browser has been closed' in str(e):
                                return
                            raise
                        ext = guess_ext(ctype)
                        # Put images into an images/ subdirectory, other binaries into assets/
                        if ctype.startswith('image/'):
                            subdir = output / 'images'
                            subdir.mkdir(parents=True, exist_ok=True)
                            fname = f"{safe}_{h}_{ts}.{ext}"
                            body_path = str(subdir / fname)
                            save_binary(subdir / fname, data)
                        else:
                            subdir = output / 'assets'
                            subdir.mkdir(parents=True, exist_ok=True)
                            fname = f"{safe}_{h}_{ts}.{ext}"
                            body_path = str(subdir / fname)
                            save_binary(subdir / fname, data)
                except Exception:
                    # final fallback: try binary
                    try:
                        data = await resp.body()
                        ext = guess_ext(ctype)
                        if ctype.startswith('image/'):
                            subdir = output / 'images'
                            subdir.mkdir(parents=True, exist_ok=True)
                            fname = f"{safe}_{h}_{ts}.{ext}"
                            body_path = str(subdir / fname)
                            save_binary(subdir / fname, data)
                        else:
                            subdir = output / 'assets'
                            subdir.mkdir(parents=True, exist_ok=True)
                            fname = f"{safe}_{h}_{ts}.{ext}"
                            body_path = str(subdir / fname)
                            save_binary(subdir / fname, data)
                    except Exception:
                        # skip if still failing
                        return

                meta = {
                    'url': url,
                    'resource_type': rtype,
                    'status': resp.status,
                    'timestamp': ts,
                    'response_headers': resp_headers,
                    'response_body_file': body_path,
                    'response_excerpt': excerpt,
                }
                meta_name = f"{safe}_{h}_{ts}.meta.json"
                save_text(output / meta_name, json.dumps(meta, ensure_ascii=False, indent=2))
                try:
                    with open(index_path, 'a', encoding='utf-8') as idx:
                        idx.write(json.dumps({'url': url, 'resource_type': rtype, 'timestamp': ts, 'body_file': body_path, 'meta_file': str(output / meta_name), 'excerpt': excerpt}, ensure_ascii=False) + '\n')
                except Exception:
                    pass

                async with capture_count_lock:
                    capture_count += 1
            except Exception:
                # swallow unexpected exceptions from handler so event loop stays healthy
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tRESP_HANDLER_ERROR\n")
                except Exception:
                    pass

        # helpers to attach to page
        def attach_page_listeners(pg):
            pg.on('response', lambda r: asyncio.create_task(on_response(r)))
            pg.on('request', lambda req: open(output / '_debug_events.log', 'a', encoding='utf-8').write(f"{now_ts()}\tREQUEST\t{req.method}\t{req.url}\n"))
            pg.on('requestfinished', lambda req: open(output / '_debug_events.log', 'a', encoding='utf-8').write(f"{now_ts()}\tREQUEST_FINISHED\t{req.method}\t{req.url}\n"))

        # attach listeners to existing pages
        for pg in context.pages:
            attach_page_listeners(pg)

        # attach to future pages
        context.on('page', lambda new_page: attach_page_listeners(new_page))

        # context-level response and request listeners
        context.on('response', lambda r: asyncio.create_task(on_response(r)))
        context.on('request', lambda req: open(output / '_debug_events.log', 'a', encoding='utf-8').write(f"{now_ts()}\tCTX_REQUEST\t{req.method}\t{req.url}\n"))

        # Start heartbeat
        hb_task = asyncio.create_task(heartbeat())

        # Navigate after listeners attached
        try:
            await page.goto(args.url)
        except Exception:
            pass

        # If interactive, let user log in; otherwise start capture immediately
        if not args.no_prompt:
            print('When you have logged in in the opened browser, press Enter here to continue and capture...')
            try:
                await asyncio.get_event_loop().run_in_executor(None, input)
            except Exception:
                pass

        print(f'Starting async passive capture (types: {args.capture_types})')

        try:
            if args.continuous:
                # keep running until KeyboardInterrupt
                while True:
                    await asyncio.sleep(1)
            else:
                await asyncio.sleep(5)
        except KeyboardInterrupt:
            print('Interrupted; closing...')
        finally:
            hb_task.cancel()
            try:
                await context.close()
            except Exception:
                pass


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--url', default='https://app.statdx.com/')
    parser.add_argument('--profile', default='playwright_profile')
    parser.add_argument('--output-dir', default='xhr_captured_async')
    parser.add_argument('--headless', action='store_true')
    parser.add_argument('--continuous', action='store_true')
    parser.add_argument('--no-prompt', action='store_true')
    parser.add_argument('--channel', default=os.getenv('PLAYWRIGHT_CHROME_CHANNEL', 'chrome'))
    parser.add_argument('--capture-types', default='xhr,fetch,document,other')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    # normalize capture types to a set
    args.capture_types = {t.strip().lower() for t in args.capture_types.split(',') if t.strip()}
    try:
        asyncio.run(run(args))
    except KeyboardInterrupt:
        pass