#!/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}")
        
    # Build an in-memory index of content_hash -> list of index entries for quick dedupe lookups.
    # If older index lines lack a content_hash but reference a body file, hash that file once at startup.
    in_memory_index = {}
    try:
        if index_path.exists():
            with open(index_path, 'r', encoding='utf-8') as idxf:
                for line in idxf:
                    try:
                        j = json.loads(line)
                    except Exception:
                        continue
                    ch = j.get('content_hash')
                    body = j.get('body_file')
                    if not ch and body and os.path.exists(body):
                        try:
                            with open(body, 'rb') as bf:
                                ch = hashlib.sha256(bf.read()).hexdigest()
                        except Exception:
                            ch = None
                    if ch:
                        in_memory_index.setdefault(ch, []).append(j)
    except Exception:
        in_memory_index = {}

    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 attempt_autologin(page, username: str, password: str, post_login_selector: str, wait_after: float = 1.0):
            """Async form-based autologin similar to save_page_snapshots.attempt_autologin."""
            if not username or not password:
                return False
            # selectors
            user_selectors = ['input[name="username"]', 'input.usernameSelector', 'input[type="email"]', 'input[name*="email" i]', 'input[name*="user" i]', 'input[type="text"]']
            pass_selectors = ['input[type="password"]', 'input[name="password"]', 'input.passwordSelector']

            user_sel = None
            for sel in user_selectors:
                try:
                    el = await page.query_selector(sel)
                    if el:
                        user_sel = sel
                        break
                except Exception:
                    continue

            pass_sel = None
            for sel in pass_selectors:
                try:
                    el = await page.query_selector(sel)
                    if el:
                        pass_sel = sel
                        break
                except Exception:
                    continue

            if not user_sel or not pass_sel:
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tAUTOLOGIN\tselectors_missing\n")
                except Exception:
                    pass
                return False

            try:
                await page.fill(user_sel, username)
                await asyncio.sleep(0.1)
                await page.fill(pass_sel, password)
                await asyncio.sleep(0.1)
                # try submit
                submit_selectors = ['button[type=submit]', 'input[type=submit]', 'button:has-text("Sign in")', 'button:has-text("Sign In")', 'button:has-text("Log in")']
                clicked = False
                for s in submit_selectors:
                    try:
                        btn = await page.query_selector(s)
                        if btn:
                            await btn.click()
                            clicked = True
                            break
                    except Exception:
                        continue
                if not clicked:
                    try:
                        await page.press(pass_sel, 'Enter')
                    except Exception:
                        pass

                # wait for post-login selector
                try:
                    if post_login_selector:
                        await page.wait_for_selector(post_login_selector, timeout=10000)
                except Exception:
                    pass
                await asyncio.sleep(wait_after)
                return True
            except Exception:
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tAUTOLOGIN\tfailed_exception\n")
                except Exception:
                    pass
                return False


        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)
                            try:
                                print(f"{now_ts()}\tSAVED\t{body_path}")
                            except Exception:
                                pass
                            excerpt = pretty[:800]
                        except Exception:
                            fname = f"{safe}_{h}_{ts}.txt"
                            body_path = str(output / fname)
                            save_text(output / fname, txt)
                            try:
                                print(f"{now_ts()}\tSAVED\t{body_path}")
                            except Exception:
                                pass
                            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)
                            try:
                                print(f"{now_ts()}\tSAVED\t{body_path}")
                            except Exception:
                                pass
                        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)
                            try:
                                print(f"{now_ts()}\tSAVED\t{body_path}")
                            except Exception:
                                pass
                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)
                            try:
                                print(f"{now_ts()}\tSAVED\t{body_path}")
                            except Exception:
                                pass
                        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)
                            try:
                                print(f"{now_ts()}\tSAVED\t{body_path}")
                            except Exception:
                                pass
                    except Exception:
                        # skip if still failing
                        return

                # compute content hash for dedupe
                content_hash = None
                try:
                    if 'txt' in locals() and isinstance(txt, str):
                        content_hash = hashlib.sha256(txt.encode('utf-8')).hexdigest()
                    elif 'pretty' in locals() and isinstance(pretty, str):
                        content_hash = hashlib.sha256(pretty.encode('utf-8')).hexdigest()
                    elif 'data' in locals() and isinstance(data, (bytes, bytearray)):
                        content_hash = hashlib.sha256(data).hexdigest()
                except Exception:
                    content_hash = None

                # prepare the meta filename we'll write for this response
                new_meta_name = f"{safe}_{h}_{ts}.meta.json"
                new_meta_abs = os.path.abspath(str(output / new_meta_name))

                # If older responses exist with the same content, delete them (fast in-memory lookup)
                try:
                    matches = in_memory_index.get(content_hash, []) if content_hash else []
                    for j in matches:
                        body = j.get('body_file')
                        metaf = j.get('meta_file')
                        try:
                            if body and os.path.exists(body) and os.path.abspath(body) != os.path.abspath(body_path):
                                os.remove(body)
                                try:
                                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                                        dbg.write(f"{now_ts()}\tDELETED_OLD_BODY\t{body}\n")
                                except Exception:
                                    pass
                        except Exception:
                            pass
                        try:
                            if metaf and os.path.exists(metaf) and os.path.abspath(metaf) != new_meta_abs:
                                os.remove(metaf)
                                try:
                                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                                        dbg.write(f"{now_ts()}\tDELETED_OLD_META\t{metaf}\n")
                                except Exception:
                                    pass
                        except Exception:
                            pass
                except Exception:
                    pass

                meta = {
                    'url': url,
                    'resource_type': rtype,
                    'status': resp.status,
                    'timestamp': ts,
                    'response_headers': resp_headers,
                    'response_body_file': body_path,
                    'response_excerpt': excerpt,
                    'content_hash': content_hash,
                }
                meta_name = f"{safe}_{h}_{ts}.meta.json"
                save_text(output / meta_name, json.dumps(meta, ensure_ascii=False, indent=2))
                try:
                    entry = {'url': url, 'resource_type': rtype, 'timestamp': ts, 'body_file': body_path, 'meta_file': str(output / meta_name), 'excerpt': excerpt, 'content_hash': content_hash}
                    with open(index_path, 'a', encoding='utf-8') as idx:
                        idx.write(json.dumps(entry, ensure_ascii=False) + '\n')
                    # update in-memory index
                    if content_hash:
                        in_memory_index.setdefault(content_hash, []).append(entry)
                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

        # attempt autologin (form-based) if credentials provided
        uname = args.username or os.getenv('STATDX_USERNAME')
        pwd = args.password or os.getenv('STATDX_PASSWORD')
        if uname and pwd:
            try:
                ok = await attempt_autologin(page, uname, pwd, args.post_login_selector)
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tAUTOLOGIN_RESULT\t{ok}\n")
                except Exception:
                    pass
            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('--username', help='STATdx username (or use STATDX_USERNAME env var)')
    parser.add_argument('--password', help='STATdx password (or use STATDX_PASSWORD env var)')
    parser.add_argument('--post-login-selector', default='#ds-app', help='Selector that indicates a successful login')
    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