#!/usr/bin/env python3
"""
Save full page HTML snapshots after page loads and periodically while you browse.

This script launches a persistent Playwright profile (so you can log in),
optionally attempts automatic login, then either performs a one-shot snapshot
or continues saving snapshots as you browse. Snapshots are de-duplicated by
content hash to avoid writing identical files repeatedly.

Usage examples:
  # manual login, one-shot snapshot after reload
  python scrapers/save_page_snapshots.py --url https://app.statdx.com/ --output-dir xhr_captured

  # auto-login and continuous snapshots while you browse
  python scrapers/save_page_snapshots.py --continuous --username you --password secret --output-dir xhr_captured

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

try:
    from loguru import logger
except Exception:
    import logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    logger = logging.getLogger("save_page_snapshots")

from playwright.sync_api import sync_playwright
from dotenv import load_dotenv

# load .env
load_dotenv()


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


def write_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 write_json(path: Path, obj):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)


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


def snapshot_page(page, out_dir: Path, index_path: Path, last_hashes: dict, min_interval: float = 3.0):
    """Capture page content and write HTML + meta if content changed or enough time passed."""
    try:
        url = page.url
    except Exception:
        url = 'unknown'

    try:
        content = page.content()
    except Exception as e:
        logger.exception('Failed to get page.content(): {}', e)
        return None

    h = hashlib.sha1(content.encode('utf-8')).hexdigest()
    prev = last_hashes.get(url)
    ts = now_ts()

    # If unchanged and we saved recently, skip
    if prev and prev.get('hash') == h and (time.time() - prev.get('ts', 0)) < min_interval:
        logger.debug('Snapshot skipped (no change) for {}', url)
        return None

    # prepare filenames
    name = sanitize_filename(url.replace('https://', '').replace('http://', ''))
    short = h[:8]
    html_name = f"{name}_{short}_{ts}.html"
    meta_name = f"{name}_{short}_{ts}.json"

    html_path = out_dir / html_name
    meta_path = out_dir / meta_name

    try:
        write_text(html_path, content)
    except Exception as e:
        logger.exception('Failed to write HTML snapshot: {}', e)
        return None

    meta = {
        'url': url,
        'timestamp': ts,
        'file': str(html_path),
        'sha1': h,
        'title': (page.title() if page else None),
    }
    try:
        write_json(meta_path, meta)
    except Exception:
        logger.exception('Failed to write meta file')

    # append to index
    try:
        entry = {'url': url, 'timestamp': ts, 'html_file': str(html_path), 'meta_file': str(meta_path), 'sha1': h}
        with open(index_path, 'a', encoding='utf-8') as idx:
            idx.write(json.dumps(entry, ensure_ascii=False) + '\n')
    except Exception:
        logger.exception('Failed to append to index')

    last_hashes[url] = {'hash': h, 'ts': time.time()}
    logger.info('Saved snapshot: {} ({} bytes)', html_path.name, len(content))
    return html_path


def attempt_autologin(page, username: str, password: str, post_login_selector: str, wait_after: float = 3.0):
    if not username or not password:
        return False
    logger.info('Attempting automatic login')
    # Common selectors for STATdx and similar pages
    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 not user_sel or not pass_sel:
        logger.warning('Could not find username/password 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 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 = page.query_selector(s)
                if btn:
                    btn.click()
                    clicked = True
                    logger.info('Clicked submit button using selector: {}', s)
                    break
            except Exception:
                continue
        if not clicked:
            page.press(pass_sel, 'Enter')

        # wait for post-login selector
        try:
            if post_login_selector:
                page.wait_for_selector(post_login_selector, timeout=10000)
        except Exception:
            logger.info('Post-login selector not found within 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 capturing while you browse until interrupted')
    parser.add_argument('--capture-wait', type=float, default=6.0, help='Seconds to wait before initial snapshot')
    parser.add_argument('--poll-interval', type=float, default=5.0, help='Periodic snapshot interval when continuous')
    parser.add_argument('--min-interval', type=float, default=3.0, help='Minimum seconds between snapshots for same URL')
    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('--no-prompt', action='store_true')
    args = parser.parse_args()

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

    username = args.username or os.getenv('STATDX_USERNAME')
    password = args.password or os.getenv('STATDX_PASSWORD')

    last_hashes = {}

    with sync_playwright() as p:
        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()
        logger.info('Opening %s', args.url)
        page.goto(args.url)

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

        # wait before initial snapshot to let in-flight XHRs populate DOM
        logger.info('Waiting %.1f sec before initial snapshot', args.capture_wait)
        time.sleep(args.capture_wait)

        # initial snapshot
        snapshot_page(page, out_dir, index_path, last_hashes, min_interval=args.min_interval)

        if args.continuous:
            logger.info('Continuous snapshot mode: saving on navigation and every %.1f sec', args.poll_interval)

            # on navigation events, take a snapshot (debounced by min_interval)
            def on_load():
                try:
                    snapshot_page(page, out_dir, index_path, last_hashes, min_interval=args.min_interval)
                except Exception:
                    logger.exception('Error during on_load snapshot')

            page.on('load', lambda: on_load())
            page.on('domcontentloaded', lambda: on_load())
            page.on('framenavigated', lambda frame: on_load())

            try:
                # periodic polling snapshot
                while True:
                    time.sleep(args.poll_interval)
                    snapshot_page(page, out_dir, index_path, last_hashes, min_interval=args.min_interval)
            except KeyboardInterrupt:
                logger.info('Interrupted by user, closing')
        else:
            # one-shot: save a final snapshot after a short delay to capture dynamic updates
            try:
                time.sleep(1.0)
                snapshot_page(page, out_dir, index_path, last_hashes, min_interval=args.min_interval)
            except Exception:
                logger.exception('Error taking final snapshot')

        try:
            context.close()
        except Exception:
            pass


if __name__ == '__main__':
    main()