#!/usr/bin/env python3
"""
Passive XHR/fetch capture using Playwright.

This script launches a persistent browser profile (so you can log in),
optionally attempts auto-login, then attaches a response handler and saves
every XHR/fetch response body and metadata into `--output-dir` while you
browse. It is intentionally passive: it doesn't reload the page unless you
ask it to; it simply listens and saves responses as they happen.

Usage:
  python scrapers/capture_passive_playwright.py --output-dir xhr_captured --continuous

Options:
  --username/--password or STATDX_USERNAME/STATDX_PASSWORD via .env for auto-login
  --no-prompt to start listening immediately
  --capture-types to adjust which resource types to record (default: xhr,fetch)
"""
from pathlib import Path
import argparse
import json
import sys
import time
import mimetypes
import hashlib
from datetime import datetime, timezone
import os

from loguru import logger

from playwright.sync_api import sync_playwright
import asyncio
try:
    from playwright._impl._errors import TargetClosedError as _PlaywrightTargetClosedError
    TARGET_CLOSED_EXCEPTIONS = (_PlaywrightTargetClosedError,)
except Exception:
    # If we cannot import Playwright internal error types, fall back to Exception
    TARGET_CLOSED_EXCEPTIONS = (Exception,)
from dotenv import load_dotenv

logger.remove()
logger.add("capture_passive_playwright.log", rotation="10 MB")
logger.add(sys.stderr, level="TRACE")
# Keep detailed logs in the file only to avoid noisy console output while you browse.
# We will print a short summary to stdout when the browser/context is closed.

load_dotenv()


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_binary(path: Path, data: bytes):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'wb') as f:
        f.write(data)


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 try_pretty_json(s: str):
    try:
        parsed = json.loads(s)
        return json.dumps(parsed, ensure_ascii=False, indent=2)
    except Exception:
        return None


def attempt_autologin(page, username, password, post_login_selector, wait_after=3.0):
    if not username or not password:
        return False
    logger.info('Attempting auto-login')
    user_selectors = ['input[name="username"]', 'input.usernameSelector', 'input[type="email"]', 'input[name*="email" i]', 'input[name*="user" i]']
    pass_selectors = ['input[type="password"]', 'input[name="password"]', 'input.passwordSelector']

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

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

    if not user_sel or not pass_sel:
        logger.warning('Could not find login fields for auto-login')
        return False

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

        if post_login_selector:
            try:
                page.wait_for_selector(post_login_selector, timeout=10000)
            except Exception:
                logger.info('Post-login selector not found in timeout')

        time.sleep(wait_after)
        return True
    except Exception:
        logger.exception('Auto-login failed')
        return False


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--url', '-u', default='https://app.statdx.com/')
    parser.add_argument('--profile', '-p', default='playwright_profile')
    parser.add_argument('--output-dir', '-o', default='xhr_captured')
    parser.add_argument('--headless', action='store_true')
    parser.add_argument('--continuous', action='store_true', help='Keep listening and saving as you browse')
    parser.add_argument('--no-prompt', action='store_true', help='Start listening immediately without waiting for Enter')
    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')
    parser.add_argument('--capture-types', default='xhr,fetch,document,other', help='Comma-separated resource types to capture (empty = all)')
    args = parser.parse_args()

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

    username = args.username or os.getenv('STATDX_USERNAME')
    password = args.password or os.getenv('STATDX_PASSWORD')
    capture_types = {t.strip().lower() for t in args.capture_types.split(',') if t.strip()}

    with sync_playwright() as p:
        # launch persistent context
        try:
            context = p.chromium.launch_persistent_context(user_data_dir=args.profile, headless=args.headless, channel='chrome')
        except Exception:
            context = p.chromium.launch_persistent_context(user_data_dir=args.profile, headless=args.headless)

        page = context.new_page()
        # NOTE: don't navigate until we've attached listeners so we don't miss
        # early requests from the initial load.

        # response handler (define before we attach listeners)
        def on_response(resp):
            """Safe response handler: reads response bodies but tolerates browser/context close.

            Playwright can raise TargetClosedError or asyncio.CancelledError when the
            browser or context closes while we're trying to read a body. Those are
            normal during shutdown and should not produce noisy tracebacks. This
            handler therefore catches those and returns quietly.
            """
            logger.debug(f'Response: {getattr(resp, "status", "?")} {getattr(resp, "url", "?")}')
            try:
                req = resp.request
                rtype = (req.resource_type or '').lower()
                # Log every response minimally so user sees activity
                logger.info('Response event: %s %s %s', getattr(resp, 'status', '?'), rtype, req.url)

                # write debug line to disk so we can verify events even if console is quiet
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\t{getattr(resp, 'status', '?')}\t{rtype}\t{req.url}\n")
                except Exception:
                    pass

                # If capture_types is empty, capture everything. Otherwise check membership.
                allow = False
                if not capture_types:
                    allow = True
                elif rtype in capture_types:
                    allow = True
                # Heuristic: always capture document content endpoints we care about
                if '/document/content/' in req.url or '/document/summary/' in req.url:
                    allow = True
                if not allow:
                    return

                url = req.url
                ts = now_ts()
                status = getattr(resp, 'status', None)
                req_headers = dict(req.headers)
                resp_headers = dict(resp.headers)

                meta = {
                    'url': url,
                    'method': req.method,
                    'resource_type': rtype,
                    'status': status,
                    'timestamp': ts,
                    'request_headers': req_headers,
                    'response_headers': resp_headers,
                }

                # request post data if any
                try:
                    post = req.post_data
                    meta['request_post_data'] = post
                except Exception:
                    meta['request_post_data'] = None

                ctype = (resp_headers.get('content-type') or '').lower()
                safe = sanitize(url.replace('https://', '').replace('http://', ''))
                h = hashlib.sha1((url + ts).encode('utf-8')).hexdigest()[:8]

                body_path = None
                excerpt = None
                # try textual
                try:
                    if 'application/json' in ctype or ctype.startswith('text/') or 'javascript' in ctype:
                        try:
                            txt = resp.text()
                        except TARGET_CLOSED_EXCEPTIONS + (asyncio.CancelledError,):
                            logger.debug('Response body unavailable (target closed) for %s', getattr(resp, 'url', '<unknown>'))
                            return
                        # pretty-print JSON if possible
                        pretty = try_pretty_json(txt)
                        if pretty is not None:
                            fname = f"{safe}_{h}_{ts}.json"
                            body_path = str(output / fname)
                            save_text(output / fname, pretty)
                            excerpt = pretty[:800]
                        else:
                            fname = f"{safe}_{h}_{ts}.txt"
                            body_path = str(output / fname)
                            save_text(output / fname, txt)
                            excerpt = txt[:800]
                    else:
                        try:
                            data = resp.body()
                        except TARGET_CLOSED_EXCEPTIONS + (asyncio.CancelledError,):
                            logger.debug('Response body unavailable (target closed) for %s', getattr(resp, 'url', '<unknown>'))
                            return
                        ext = guess_ext(ctype)
                        # For images, avoid per-capture timestamps: use stable name with hash
                        if ctype.startswith('image/'):
                            fname = f"{safe}_{h}.{ext}"
                        else:
                            fname = f"{safe}_{h}_{ts}.{ext}"
                        body_path = str(output / fname)
                        save_binary(output / fname, data)
                except Exception as e:
                    try:
                        # fallback to binary
                        try:
                            data = resp.body()
                        except TARGET_CLOSED_EXCEPTIONS + (asyncio.CancelledError,):
                            logger.debug('Response body unavailable (target closed) for %s', getattr(resp, 'url', '<unknown>'))
                            return
                        ext = guess_ext(ctype)
                        if ctype.startswith('image/'):
                            fname = f"{safe}_{h}.{ext}"
                        else:
                            fname = f"{safe}_{h}_{ts}.{ext}"
                        body_path = str(output / fname)
                        save_binary(output / fname, data)
                    except Exception:
                        logger.debug('Failed to save body for %s: %s', url, e)

                meta['response_body_file'] = body_path
                meta['response_excerpt'] = excerpt

                meta_name = f"{safe}_{h}_{ts}.json"
                meta_path = output / meta_name
                save_text(meta_path, json.dumps(meta, ensure_ascii=False, indent=2))

                # append index
                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(meta_path), 'excerpt': excerpt}, ensure_ascii=False) + '\n')
                except Exception:
                    logger.exception('Failed to write index entry')

                logger.info('Captured %s %s -> %s', rtype, url, meta_path.name)
            except Exception as e:
                # suppress noisy Playwright shutdown traces
                if isinstance(e, TARGET_CLOSED_EXCEPTIONS + (asyncio.CancelledError,)) or \
                   (hasattr(e, 'args') and any('Target page, context or browser has been closed' in str(a) for a in e.args)):
                    logger.debug('Response handler aborted: target closed')
                    return
                logger.exception('Error in response handler')

        # helper to attach listeners to a page
        def attach_listeners_to_page(pg):
            try:
                pg.on('response', on_response)
                # log requests and finished requests for debug
                def on_request(req):
                    try:
                        with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                            dbg.write(f"{now_ts()}\tREQUEST\t{req.method}\t{req.url}\n")
                    except Exception:
                        pass

                def on_request_finished(req):
                    try:
                        with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                            dbg.write(f"{now_ts()}\tREQUEST_FINISHED\t{req.method}\t{req.url}\n")
                    except Exception:
                        pass

                pg.on('request', on_request)
                pg.on('requestfinished', on_request_finished)
                # also log request failures for visibility
                def on_req_failed(req):
                    try:
                        logger.warning('Request failed: %s %s', req.failure, req.url)
                    except Exception:
                        pass
                pg.on('requestfailed', on_req_failed)
                logger.debug('Attached listeners to page: %s', pg.url)
            except Exception:
                logger.exception('Failed to attach listeners to page')

            try:
                # navigation and lifecycle events
                def on_navigate():
                    try:
                        with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                            dbg.write(f"{now_ts()}\tNAVIGATE\t{pg.url}\n")
                    except Exception:
                        pass
                pg.on('load', lambda: on_navigate())
                pg.on('domcontentloaded', lambda: on_navigate())
                pg.on('framenavigated', lambda frame: on_navigate())
            except Exception:
                pass

        # attach to any existing pages in the context
        try:
            for pg in context.pages:
                attach_listeners_to_page(pg)
        except Exception:
            logger.debug('No existing pages to attach')

        # attach to future pages
        try:
            context.on('page', lambda new_page: attach_listeners_to_page(new_page))
        except Exception:
            logger.debug('Could not attach context.page listener')

        # also attach at context level to catch responses not tied to a single page
        try:
            context.on('response', lambda resp: on_response(resp))
        except Exception:
            logger.debug('Context-level response listener not available')

        # add context-level request listeners too (some fetches/service-workers
        # may appear at the context level rather than page)
        try:
            def _ctx_request(req):
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tCTX_REQUEST\t{req.method}\t{req.url}\n")
                except Exception:
                    pass
            def _ctx_req_finished(req):
                try:
                    with open(output / '_debug_events.log', 'a', encoding='utf-8') as dbg:
                        dbg.write(f"{now_ts()}\tCTX_REQUEST_FINISHED\t{req.method}\t{req.url}\n")
                except Exception:
                    pass
            context.on('request', lambda r: _ctx_request(r))
            context.on('requestfinished', lambda r: _ctx_req_finished(r))
        except Exception:
            logger.debug('Context-level request listeners not available')

        # try auto-login
        if username and password:
            attempt_autologin(page, username, password, args.post_login_selector)

        # navigate after listeners attached so we capture initial load
        try:
            logger.info('Opening %s', args.url)
            page.goto(args.url)
        except Exception:
            logger.debug('Initial navigation failed or was skipped')

        # If interactive and not no-prompt, let user finish manual login/navigation
        if not args.no_prompt:
            try:
                input('When you have logged in and are ready, press Enter to start passive capture...\n')
            except Exception:
                pass

        logger.info(f'Starting passive response capture (types: {", ".join(sorted(capture_types))})')

        # attach the response handler to the initial page
        page.on('response', on_response)

        logger.info('Passive listener attached. Continue browsing in the opened browser window.')

        if args.continuous:
            try:
                while True:
                    time.sleep(1)
            except KeyboardInterrupt:
                logger.info('Interrupted by user; finishing and closing browser')
        else:
            # one-shot: wait a short while to gather responses then exit
            time.sleep(5)

        try:
            context.close()
        except Exception:
            pass


if __name__ == '__main__':
    main()