#!/usr/bin/env python3
"""
Capture all network response bodies using Chrome DevTools Protocol (CDP).
This script launches a persistent Chromium/Chrome profile (so you can log in),
attaches a CDP session to the page, enables Network events and calls
Network.getResponseBody for every responseReceived event — saving the body
(text or binary) and a JSON metadata file per response.

Usage:
  python scrapers/capture_with_cdp.py --url https://app.statdx.com/ --output-dir xhr_captured --capture-wait 10

Notes:
- Requires Playwright and an installed Chrome/Chromium. Use the same Python
  environment you used for other scripts in this repo.
- This method retrieves bodies even for cached or service-worker responses in
  many cases because it asks Chrome directly for the response payload.

"""
import argparse
import base64
import hashlib
import json
import mimetypes
from datetime import timezone
import time
from datetime import datetime
from pathlib import Path

try:
    from loguru import logger
except Exception:
    import logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    logger = logging.getLogger("capture_with_cdp")
from playwright.sync_api import sync_playwright
from dotenv import load_dotenv
import os

# load .env
load_dotenv()


def sanitize_fn(s: str) -> str:
    safe = []
    for ch in s:
        if ch.isalnum() or ch in "._-":
            safe.append(ch)
        else:
            safe.append("-")
    name = "".join(safe)
    # shorten
    if len(name) > 180:
        h = hashlib.sha1(s.encode("utf-8")).hexdigest()[:8]
        name = name[:120] + "-" + h
    return name


def guess_ext_from_mime(mime: str) -> str:
    if not mime:
        return "bin"
    mt = mime.split(";")[0].strip()
    ext = mimetypes.guess_extension(mt)
    if ext:
        return ext.lstrip('.')
    # fallback mapping
    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 write_file(path: Path, data: bytes):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "wb") as f:
        f.write(data)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="https://app.statdx.com/", help="URL to open")
    parser.add_argument("--output-dir", default="xhr_captured", help="Directory to write captures")
    parser.add_argument("--profile", default="playwright_profile", help="Persistent profile directory")
    parser.add_argument("--headless", action="store_true", help="Run headless")
    parser.add_argument("--capture-wait", type=int, default=8, help="Seconds to wait after reload to capture XHRs")
    parser.add_argument("--keep-open", action="store_true", help="Keep browser open after capture")
    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 to wait for after login")
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    logger.info("Starting Playwright and launching Chrome/Chromium (persistent profile: {})", args.profile)

    with sync_playwright() as p:
        # try to use installed Chrome first
        try:
            browser = p.chromium.launch_persistent_context(user_data_dir=args.profile, headless=args.headless, channel="chrome")
            page = browser.pages[0] if browser.pages else browser.new_page()
        except Exception as e:
            logger.warning("Failed to launch Chrome channel: {}. Falling back to chromium bundle.", e)
            browser = p.chromium.launch_persistent_context(user_data_dir=args.profile, headless=args.headless)
            page = browser.pages[0] if browser.pages else browser.new_page()

        # attach CDP session to the page
        try:
            session = browser.new_cdp_session(page)
        except Exception as e:
            logger.error("Failed to create CDP session: {}", e)
            browser.close()
            return

        logger.info("Enabling Network domain on CDP")
        session.send("Network.enable")

        seen = set()

        def on_response_received(params):
            # params contains requestId and response metadata
            try:
                requestId = params.get("requestId")
                response = params.get("response", {})
                url = response.get("url")
                status = response.get("status")
                headers = response.get("headers", {})
                mime = headers.get("content-type") or response.get("mimeType") or ""
                ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")

                key = f"{url}|{status}|{requestId}"
                if key in seen:
                    return
                seen.add(key)

                # attempt to get body
                body = None
                base64_encoded = False
                try:
                    result = session.send("Network.getResponseBody", {"requestId": requestId})
                    if isinstance(result, dict):
                        body = result.get("body")
                        base64_encoded = bool(result.get("base64Encoded"))
                except Exception as e:
                    logger.debug("Network.getResponseBody failed for {}: {}", url, e)

                # fallback: if no body, skip writing but still store metadata
                # build filenames
                safe = sanitize_fn(url.replace('https://', '').replace('http://', ''))
                h = hashlib.sha1((url + str(time.time())).encode("utf-8")).hexdigest()[:8]
                # determine extension
                ext = guess_ext_from_mime(mime)
                body_filename = f"{safe}_{h}_{ts}.{ext}"
                meta_filename = f"{safe}_{h}_{ts}.json"

                meta = {
                    "url": url,
                    "status": status,
                    "timestamp": ts,
                    "requestId": requestId,
                    "mime": mime,
                    "headers": headers,
                }

                if body is not None:
                    try:
                        if base64_encoded:
                            data = base64.b64decode(body)
                        else:
                            data = body.encode("utf-8")
                        write_file(out_dir / body_filename, data)
                        meta["response_body_file"] = str(out_dir / body_filename)
                        # try to create a small excerpt for indexing
                        try:
                            excerpt = data.decode("utf-8", errors="ignore")[:800]
                        except Exception:
                            excerpt = ""
                        meta["response_excerpt"] = excerpt
                        logger.info("Saved body {} ({} bytes) for {}", body_filename, len(data), url)
                    except Exception as e:
                        logger.warning("Failed to save body for {}: {}", url, e)
                else:
                    meta["response_body_file"] = None

                # write metadata
                with open(out_dir / meta_filename, "w", encoding="utf-8") as mf:
                    json.dump(meta, mf, ensure_ascii=False, indent=2)

            except Exception as exc:
                logger.exception("Error handling responseReceived: {}", exc)

        # subscribe to CDP event
        session.on("Network.responseReceived", on_response_received)

        logger.info("Opening page: {}. Please log in if needed, then press ENTER to start capture.", args.url)
        page.goto(args.url)

        # Automatic login if credentials provided
        username = args.username or os.getenv('STATDX_USERNAME')
        password = args.password or os.getenv('STATDX_PASSWORD')
        if username and password:
            logger.info('Attempting automatic login using provided credentials')
            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:
                    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 user_sel and pass_sel:
                try:
                    page.fill(user_sel, username)
                    time.sleep(0.1)
                    page.fill(pass_sel, password)
                    time.sleep(0.1)
                    # try submit methods
                    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 = page.query_selector(s)
                            if btn:
                                btn.click()
                                clicked = True
                                logger.info('Clicked submit button using selector: %s', s)
                                break
                        except Exception:
                            continue
                    if not clicked:
                        page.press(pass_sel, 'Enter')

                    # wait for post-login selector
                    try:
                        if args.post_login_selector:
                            page.wait_for_selector(args.post_login_selector, timeout=10000)
                    except Exception:
                        logger.info('Post-login selector not found within timeout')
                except Exception:
                    logger.exception('Automatic login attempt failed')

        input("After logging in and navigating to the page you want to capture, press Enter to begin capture...\n")

        logger.info("Reloading page to trigger XHRs and network activity")
        page.reload()

        # wait a bit to let XHRs fire and be processed by CDP
        wait = args.capture_wait
        logger.info("Waiting {} seconds to collect responses...", wait)
        time.sleep(wait)

        logger.info("Capture pass complete. Metadata and bodies written to {}", out_dir)

        if args.keep_open:
            logger.info("Keeping browser open. Close manually when done.")
            try:
                while True:
                    time.sleep(1)
            except KeyboardInterrupt:
                logger.info("Received interrupt, closing.")
                browser.close()
        else:
            browser.close()


if __name__ == "__main__":
    main()