#!/usr/bin/env python3
"""Save XHR/fetch requests and responses (including POST bodies).

Usage:
  1) Start the script with a persistent profile so you can log in manually:
     python scrapers/save_xhr_requests.py --profile .playwright_profile --output xhr_captured --domain app.statdx.com
  2) A browser opens. Log in and navigate to the page that issues XHRs.
  3) Return to the terminal and press Enter to begin capturing. The script will reload the current page to re-trigger XHRs.
  4) When you're done, press Ctrl+C; captured request/response pairs will be saved in the output directory.

Files produced:
  - <sanitized-url>_<hash>_<timestamp>.json  -- metadata (request, response headers, status, any textual body)
  - <same>.<ext> -- binary response body saved when non-text content (images, etc.)

The saved JSON includes (where available): url, method, status, request_headers, request_post_data, response_headers, response_text (if textual), response_body_file (if binary saved).
"""

import os
import argparse
import logging
import time
import hashlib
from datetime import datetime, timezone
from urllib.parse import urlparse
import mimetypes
from playwright.sync_api import sync_playwright

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')


def sanitize_filename(url: str) -> str:
    parsed = urlparse(url)
    base = parsed.netloc + parsed.path
    if parsed.query:
        base += '?' + parsed.query
    base = base.strip('/')
    if not base:
        base = parsed.netloc
    safe = ''.join(c if c.isalnum() or c in ('-', '_', '.') else '_' for c in base)
    h = hashlib.sha1(url.encode('utf-8')).hexdigest()[:8]
    return f"{safe[:200]}_{h}"


def save_binary(body: bytes, out_dir: str, url: str, ts: str, ctype: str):
    ext = ''
    try:
        ext = mimetypes.guess_extension(ctype.split(';')[0]) or ''
    except Exception:
        ext = ''
    fname = f"{sanitize_filename(url)}_{ts}{ext}"
    path = os.path.join(out_dir, fname)
    with open(path, 'wb') as fh:
        fh.write(body)
    return path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--profile', '-p', default='.playwright_profile')
    parser.add_argument('--output', '-o', default='xhr_captured')
    parser.add_argument('--domain', '-d', default='app.statdx.com')
    parser.add_argument('--topics-url', default='https://app.statdx.com/')
    parser.add_argument('--wait', type=float, default=6.0, help='Seconds to wait after reload for XHRs to complete')
    parser.add_argument('--channel', default='chrome', help='Playwright browser channel to use (e.g. chrome). Set to empty to use bundled Chromium')
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)

    with sync_playwright() as p:
        browser = p.chromium
        logging.info('Launching persistent context with profile: %s (channel=%s)', args.profile, args.channel)
        if args.channel:
            context = browser.launch_persistent_context(user_data_dir=args.profile, headless=False, channel=args.channel)
        else:
            context = browser.launch_persistent_context(user_data_dir=args.profile, headless=False)
        page = context.new_page()

        captured = []

        def on_response(resp):
            try:
                req = resp.request
                rtype = req.resource_type
                if args.domain not in req.url:
                    return
                if rtype not in ('xhr', 'fetch'):
                    return

                ts = datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')

                info = {
                    'url': req.url,
                    'method': req.method,
                    'resource_type': rtype,
                    'status': resp.status,
                    'request_headers': dict(req.headers),
                    'response_headers': dict(resp.headers),
                    'timestamp': ts,
                }

                # Request post data (may be None)
                try:
                    post = req.post_data
                    info['request_post_data'] = post
                except Exception:
                    try:
                        info['request_post_data'] = req.post_data_text()
                    except Exception:
                        info['request_post_data'] = None

                # Response body
                ctype = resp.headers.get('content-type', '')
                try:
                    # prefer text for textual content
                    if 'application/json' in ctype or 'text/' in ctype or 'application/javascript' in ctype:
                        txt = resp.text()
                        info['response_text'] = txt
                        info['response_body_file'] = None
                    else:
                        body = resp.body()
                        path = save_binary(body, args.output, req.url, ts, ctype)
                        info['response_text'] = None
                        info['response_body_file'] = path
                except Exception:
                    # fallback: save binary
                    try:
                        body = resp.body()
                        path = save_binary(body, args.output, req.url, ts, ctype)
                        info['response_text'] = None
                        info['response_body_file'] = path
                    except Exception:
                        logging.exception('Failed to read response body for %s', req.url)

                fname = f"{sanitize_filename(req.url)}_{ts}.json"
                out_path = os.path.join(args.output, fname)
                with open(out_path, 'w', encoding='utf-8') as fh:
                    import json

                    json.dump(info, fh, indent=2)

                logging.info('Captured XHR: %s -> %s', req.url, out_path)
                captured.append(out_path)
            except Exception:
                logging.exception('Error in on_response')

        page.on('response', on_response)

        logging.info('Open the browser, log in if needed, and navigate to the page that triggers XHRs. Then return here and press Enter.')
        page.goto(args.topics_url)
        input('Press Enter to start capture (will reload the current page)...')

        logging.info('Reloading page to trigger XHRs...')
        page.reload()

        try:
            time.sleep(args.wait)
            logging.info('Waiting finished. You can interact to trigger more XHRs; press Ctrl+C when done.')
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            logging.info('Finishing capture...')
        finally:
            try:
                context.close()
            except Exception:
                logging.exception('Error closing context')

    logging.info('Done. Captured %d XHRs. Files saved under %s', len(captured), args.output)


if __name__ == '__main__':
    main()