#!/usr/bin/env python3
"""Capture a fresh SWIFTLY_BEARER_TOKEN from luckysupermarkets.com.

Uses seleniumbase in stealth CDP mode to drive the live site, intercept
window.fetch + XMLHttpRequest, and capture the Authorization header on
the first prod.swiftlyapi.net request after a store is selected. The
captured JWT is validated (iss + exp), then written to the env file.

Runs on the HOST (not inside docker) since seleniumbase needs a real
Chrome/Chromium binary. Requires:
    pip install seleniumbase
    seleniumbase install chromedriver   # one-time

Usage:
    python scripts/refresh_swiftly_token.py
    python scripts/refresh_swiftly_token.py --debug
    python scripts/refresh_swiftly_token.py --restart-backend
    python scripts/refresh_swiftly_token.py --env-file .env --zip 94806

Exits 0 on success, non-zero with diagnostics otherwise.
"""
from __future__ import annotations

import argparse
import base64
import json
import re
import subprocess
import sys
import time
from pathlib import Path

try:
    from seleniumbase import SB
except ImportError:
    print(
        "[refresh_swiftly_token] seleniumbase not installed.\n"
        "  pip install seleniumbase\n"
        "  seleniumbase install chromedriver",
        file=sys.stderr,
    )
    sys.exit(2)


SWIFTLY_HOST = "swiftlyapi"
EXPECTED_ISS = "https://securetoken.google.com/swiftly-lu-prod"
DEFAULT_ENV_FILE = Path(".env.test")
DEFAULT_ZIP = "94806"
DEFAULT_STORE_NUMBER = "757"
ENV_LINE_RE = re.compile(r"^SWIFTLY_BEARER_TOKEN=.*$", re.MULTILINE)

# JS that wraps window.fetch and XMLHttpRequest to capture the first
# Bearer header sent to a Swiftly API host. Re-applied after every
# navigation since SPAs may not always reset window state.
FETCH_HOOK_JS = r"""
(function(){
  if (window.__swiftlyHookInstalled) return;
  window.__swiftlyHookInstalled = true;
  window.__capturedSwiftlyAuth = null;

  const HOST = 'swiftlyapi';
  const origFetch = window.fetch;
  window.fetch = function() {
    try {
      const req = arguments[0];
      const init = arguments[1] || {};
      const url = typeof req === 'string' ? req : (req && req.url) || '';
      let auth = null;
      if (init.headers) {
        const h = init.headers;
        if (h instanceof Headers) auth = h.get('Authorization') || h.get('authorization');
        else if (h && typeof h === 'object') auth = h['Authorization'] || h['authorization'];
      }
      if (!auth && req && req.headers && req.headers.get) {
        auth = req.headers.get('Authorization') || req.headers.get('authorization');
      }
      if (auth && auth.indexOf('Bearer ') === 0 && url.indexOf(HOST) !== -1) {
        window.__capturedSwiftlyAuth = auth.slice(7);
      }
    } catch (e) {}
    return origFetch.apply(this, arguments);
  };

  const origOpen = XMLHttpRequest.prototype.open;
  XMLHttpRequest.prototype.open = function(m, u) {
    this.__url = u;
    return origOpen.apply(this, arguments);
  };
  const origSet = XMLHttpRequest.prototype.setRequestHeader;
  XMLHttpRequest.prototype.setRequestHeader = function(name, value) {
    try {
      if (/^authorization$/i.test(name) && value && value.indexOf('Bearer ') === 0
          && this.__url && this.__url.indexOf(HOST) !== -1) {
        window.__capturedSwiftlyAuth = value.slice(7);
      }
    } catch (e) {}
    return origSet.apply(this, arguments);
  };
})();
"""


def install_hook(sb) -> None:
    sb.cdp.evaluate(FETCH_HOOK_JS)


def get_captured(sb) -> str | None:
    val = sb.cdp.evaluate("window.__capturedSwiftlyAuth")
    if isinstance(val, str) and val:
        return val
    return None


def decode_jwt_payload(token: str) -> dict:
    parts = token.split(".")
    if len(parts) != 3:
        raise ValueError("invalid JWT shape (expected 3 segments)")
    pad = "=" * (-len(parts[1]) % 4)
    return json.loads(base64.urlsafe_b64decode(parts[1] + pad))


def validate_token(token: str) -> dict:
    payload = decode_jwt_payload(token)
    iss = payload.get("iss")
    if iss != EXPECTED_ISS:
        raise ValueError(f"unexpected JWT iss: {iss!r} (expected {EXPECTED_ISS!r})")
    exp = int(payload.get("exp", 0))
    if exp <= time.time():
        raise ValueError(f"captured token is already expired (exp={exp})")
    return payload


def select_store(sb, zip_code: str, store_number: str) -> None:
    """Open the locator drawer, type a zip, click the matching store."""
    print(f"[refresh] opening store locator at zip {zip_code}, target store #{store_number}")
    sb.cdp.get("https://luckysupermarkets.com/categories?showStoreLocator=true")
    install_hook(sb)
    time.sleep(2)

    # Type zip via JS so we don't depend on exact CSS selectors.
    typed = sb.cdp.evaluate(f"""
      (function() {{
        const inputs = Array.from(document.querySelectorAll('input'));
        const c = inputs.find(i =>
          /zip|search|locat|store/i.test(i.placeholder || '') ||
          /zip|search|locat|store/i.test(i.name || '') ||
          /zip|search|locat|store/i.test(i.id || '') ||
          i.type === 'search'
        );
        if (!c) return null;
        const setter = Object.getOwnPropertyDescriptor(window.HTMLInputElement.prototype, 'value').set;
        setter.call(c, '{zip_code}');
        c.dispatchEvent(new Event('input', {{ bubbles: true }}));
        c.dispatchEvent(new Event('change', {{ bubbles: true }}));
        c.focus();
        return 'typed in: ' + (c.placeholder || c.name || c.id || c.tagName);
      }})();
    """)
    if not typed:
        raise RuntimeError("could not find zip-code input on store-locator page")
    print(f"[refresh] {typed}")
    time.sleep(3)

    # Click any element whose text contains the store number.
    deadline = time.time() + 10
    clicked = None
    while time.time() < deadline:
        clicked = sb.cdp.evaluate(f"""
          (function() {{
            const target = '{store_number}';
            const candidates = Array.from(document.querySelectorAll(
              'button, a, [role=button], [data-store-id], [data-store], li, div'
            ));
            const match = candidates.find(el => {{
              const txt = (el.textContent || '').trim();
              return txt.length > 0 && txt.length < 400 && txt.includes(target);
            }});
            if (!match) return null;
            match.click();
            return 'clicked: ' + match.tagName + ' -> ' + (match.textContent || '').slice(0, 100);
          }})();
        """)
        if clicked:
            break
        time.sleep(0.5)
    if not clicked:
        raise RuntimeError(f"could not click an element containing store #{store_number}")
    print(f"[refresh] {clicked}")
    time.sleep(3)


def trigger_api_call(sb) -> None:
    print("[refresh] navigating to a category page to trigger /search/api/v1 call")
    sb.cdp.get("https://luckysupermarkets.com/categories/Product%2Fmeat_seafood")
    install_hook(sb)


def wait_for_token(sb, timeout_s: int) -> str:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        token = get_captured(sb)
        if token:
            return token
        time.sleep(0.5)
    raise TimeoutError(
        f"no Swiftly Bearer token captured within {timeout_s}s — "
        "the site may have blocked the request, or the UI flow may have changed"
    )


def write_token_to_env(env_path: Path, token: str) -> None:
    if not env_path.exists():
        raise FileNotFoundError(f"env file not found: {env_path}")
    content = env_path.read_text()
    new_line = f"SWIFTLY_BEARER_TOKEN={token}"
    if ENV_LINE_RE.search(content):
        new_content = ENV_LINE_RE.sub(new_line, content)
    else:
        suffix = "" if content.endswith("\n") else "\n"
        new_content = content + suffix + new_line + "\n"
    tmp = env_path.with_suffix(env_path.suffix + ".tmp")
    tmp.write_text(new_content)
    tmp.replace(env_path)


def restart_backend(env_file: Path) -> None:
    cmd = [
        "docker", "compose", "--env-file", str(env_file),
        "up", "-d", "--force-recreate", "backend",
    ]
    print(f"[refresh] $ {' '.join(cmd)}")
    subprocess.run(cmd, check=True)


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--env-file", type=Path, default=DEFAULT_ENV_FILE,
                   help=f"env file to update (default: {DEFAULT_ENV_FILE})")
    p.add_argument("--zip", default=DEFAULT_ZIP,
                   help=f"zip code for store search (default: {DEFAULT_ZIP})")
    p.add_argument("--store", default=DEFAULT_STORE_NUMBER,
                   help=f"target store number to click (default: {DEFAULT_STORE_NUMBER})")
    p.add_argument("--debug", action="store_true",
                   help="run with a visible Chrome window")
    p.add_argument("--restart-backend", action="store_true",
                   help="restart the docker backend after writing the token")
    p.add_argument("--timeout", type=int, default=30,
                   help="seconds to wait for token capture (default: 30)")
    args = p.parse_args()

    headless = not args.debug
    print(
        f"[refresh] env={args.env_file} headless={headless} "
        f"zip={args.zip} store={args.store} timeout={args.timeout}s"
    )

    with SB(uc=True, headless=headless, test=True) as sb:
        sb.activate_cdp_mode("https://luckysupermarkets.com")
        install_hook(sb)
        select_store(sb, args.zip, args.store)
        trigger_api_call(sb)
        token = wait_for_token(sb, args.timeout)

    print(f"[refresh] captured Bearer JWT (length={len(token)})")
    payload = validate_token(token)
    expires_in_s = int(payload["exp"]) - int(time.time())
    print(
        f"[refresh] valid: iss={payload.get('iss')} "
        f"exp_in={expires_in_s}s ({expires_in_s/60:.1f} min)"
    )

    write_token_to_env(args.env_file, token)
    print(f"[refresh] wrote SWIFTLY_BEARER_TOKEN to {args.env_file}")

    if args.restart_backend:
        restart_backend(args.env_file)
        print("[refresh] backend restarted")

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (RuntimeError, TimeoutError, ValueError, FileNotFoundError) as e:
        print(f"[refresh] FAILED: {e}", file=sys.stderr)
        sys.exit(1)