import os
import re
import time
import warnings
from pathlib import Path
from playwright.sync_api import sync_playwright

import cv2
import easyocr

# ============================================================
# WARNING SUPPRESSION
# ============================================================
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
os.environ["PYTHONWARNINGS"] = "ignore"

# ============================================================
# LOGIN CONFIGURATION
# ============================================================
USERNAME = "canofy01"
PASSWORD = "Siddhartha#2026"

BASE_URL = "https://www.sezonline-ndml.co.in/User/Login.aspx"

# ============================================================
# EXPORT CONFIGURATION
# (pull these from your framework's config.py / config.yaml —
#  placeholders below just show the expected format)
# ============================================================
FROM_DATE = "01/06/2026"   # dd/mm/yyyy — must match the portal's masked date format
TO_DATE = "18/06/2026"

DOWNLOAD_DIR = Path.home() / "Downloads" / "SEZ_Reports"
DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)

# (module dropdown value, label used only for logging)
MODULES_TO_EXPORT = [
    ("17", "BILL OF ENTRY"),
    ("26", "DTA SALE"),
]

SCREENSHOT_DIR = Path(__file__).parent / "screenshots"
SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)

# Initialize EasyOCR Reader ONCE (module load, not per captcha)
ocr_reader = easyocr.Reader(["en"], gpu=False)


# ============================================================
# CUSTOM EXCEPTIONS
# ============================================================
class RateLimitError(Exception):
    """Raised when SEZ Online reports that the request rate is limited."""
    pass


class SessionExpiredError(Exception):
    """Raised when the SEZ Online server-side session has timed out."""
    pass


class CaptchaUnsolvedError(Exception):
    """Raised when the CAPTCHA could not be confidently read after retries."""
    pass


RATE_LIMIT_WAIT_SECONDS = 120

# ============================================================
# CAPTCHA CONFIGURATION
# ============================================================
CAPTCHA_MIN_LENGTH = 4
CAPTCHA_MAX_LENGTH = 8
MAX_CAPTCHA_READ_ATTEMPTS = 6
MAX_LOGIN_ATTEMPTS = 8


# ============================================================
# ACTIVE IDLE (SESSION KEEP-ALIVE) PAUSE
# ============================================================
def idle_keepalive_pause(page, total_seconds: int, chunk_seconds: int = 5):
    """Wait out a required cooldown without letting the browser tab go fully idle."""
    print(f"Idling for {total_seconds}s with periodic keep-alive scrolling...")

    elapsed = 0.0
    direction = 1

    while elapsed < total_seconds:
        chunk_start = time.time()
        remaining = total_seconds - elapsed
        this_chunk = min(chunk_seconds, remaining)

        try:
            steps = max(1, int(round(this_chunk)))
            step_delta = (300 * direction) / steps
            step_wait_ms = int((this_chunk / steps) * 1000)

            for _ in range(steps):
                page.mouse.wheel(0, step_delta)
                page.wait_for_timeout(step_wait_ms)
        except Exception:
            time.sleep(this_chunk)

        direction *= -1
        elapsed += (time.time() - chunk_start)

    print("Keep-alive idling complete.")


# ============================================================
# CAPTCHA ENGINE (unchanged from the print-automation script)
# ============================================================
def is_valid_captcha_token(token: str) -> bool:
    if not token:
        return False
    if not token.isalnum():
        return False
    if not (CAPTCHA_MIN_LENGTH <= len(token) <= CAPTCHA_MAX_LENGTH):
        return False
    return True


def resolve_captcha(image_path: Path) -> str:
    """Multi-variant OCR with confidence-weighted consensus voting."""
    image = cv2.imread(str(image_path))

    if image is None:
        raise FileNotFoundError(f"Could not load CAPTCHA image from {image_path}")

    resized = cv2.resize(image, None, fx=2.5, fy=2.5, interpolation=cv2.INTER_CUBIC)
    gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)

    variants = {}

    filtered = cv2.bilateralFilter(gray, 9, 75, 75)
    variants["bilateral"] = filtered

    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    variants["enhanced"] = clahe.apply(filtered)

    _, otsu = cv2.threshold(filtered, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    variants["otsu"] = otsu

    variants["adaptive"] = cv2.adaptiveThreshold(
        filtered, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
    )

    allowlist = (
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz"
        "0123456789"
    )

    candidates = []

    for variant_name, variant_image in variants.items():
        try:
            results = ocr_reader.readtext(
                variant_image, detail=1, paragraph=False, allowlist=allowlist
            )
        except Exception as exc:
            print(f"OCR failed for {variant_name}: {exc}")
            continue

        if not results:
            continue

        try:
            results = sorted(results, key=lambda item: min(point[0] for point in item[0]))
        except Exception:
            pass

        raw_text = "".join(str(item[1]) for item in results if len(item) >= 2)
        cleaned_text = re.sub(r"[^A-Za-z0-9]", "", raw_text)

        if not cleaned_text:
            continue

        confidences = [float(item[2]) for item in results if len(item) >= 3]
        avg_confidence = (sum(confidences) / len(confidences)) if confidences else 0.0

        candidates.append({
            "variant": variant_name,
            "text": cleaned_text,
            "confidence": avg_confidence,
        })

    if not candidates:
        return ""

    grouped = {}
    for candidate in candidates:
        token = candidate["text"]
        bucket = grouped.setdefault(token, {"text": token, "confidence_sum": 0.0, "count": 0})
        bucket["confidence_sum"] += candidate["confidence"]
        bucket["count"] += 1

    grouped_candidates = list(grouped.values())
    for candidate in grouped_candidates:
        candidate["average_confidence"] = candidate["confidence_sum"] / candidate["count"]
        candidate["score"] = candidate["average_confidence"] + (0.05 * (candidate["count"] - 1))

    grouped_candidates.sort(key=lambda item: item["score"], reverse=True)
    best = grouped_candidates[0]

    case_groups = {}
    for candidate in candidates:
        key = candidate["text"].lower()
        case_groups.setdefault(key, []).append(candidate)

    compatible_group = case_groups.get(best["text"].lower(), [])
    best_text = best["text"]

    if compatible_group:
        max_length = max(len(item["text"]) for item in compatible_group)
        final_chars = []

        for index in range(max_length):
            char_votes = {}
            for candidate in compatible_group:
                text = candidate["text"]
                if index >= len(text):
                    continue
                char = text[index]
                char_votes[char] = char_votes.get(char, 0.0) + candidate["confidence"]

            if char_votes:
                final_chars.append(max(char_votes, key=char_votes.get))

        consensus_text = "".join(final_chars)
        if consensus_text:
            best_text = consensus_text

    cleaned_token = re.sub(r"[^A-Za-z0-9]", "", best_text)

    print(f"CAPTCHA OCR candidates: {[(c['variant'], c['text'], round(c['confidence'], 3)) for c in candidates]}")
    print(f"CAPTCHA consensus token: '{cleaned_token}'")

    return cleaned_token


def try_refresh_captcha_image(page) -> bool:
    candidate_selectors = [
        "#captchaControlRevamp_btnRefreshCaptcha",
        "#captchaControlRevamp_imgRefresh",
        "a[id*='Refresh' i][id*='aptcha' i]",
        "img[id*='Refresh' i][id*='aptcha' i]",
        "a[href*='refresh' i][id*='aptcha' i]",
        "[title*='refresh' i][id*='aptcha' i]",
        "[title*='reload' i][id*='aptcha' i]",
    ]

    for selector in candidate_selectors:
        try:
            control = page.locator(selector).first
            if control.count() > 0 and control.is_visible():
                control.click(timeout=3000)
                page.wait_for_timeout(500)
                return True
        except Exception:
            continue

    return False


def get_captcha_token(page, captcha_image_locator) -> str:
    raw_captcha_path = SCREENSHOT_DIR / "raw_captcha.png"
    best_fallback = ""

    for attempt in range(1, MAX_CAPTCHA_READ_ATTEMPTS + 1):
        try:
            captcha_image_locator.wait_for(state="visible", timeout=15000)
            captcha_image_locator.screenshot(path=str(raw_captcha_path))
        except Exception as exc:
            print(f"Captcha screenshot attempt {attempt} failed: {exc}")
            continue

        token = resolve_captcha(raw_captcha_path)

        if token and len(token) > len(best_fallback):
            best_fallback = token

        if is_valid_captcha_token(token):
            print(f"Captcha read accepted on attempt {attempt}: '{token}'")
            return token

        print(
            f"Captcha read attempt {attempt}/{MAX_CAPTCHA_READ_ATTEMPTS} "
            f"rejected (raw='{token}'). Trying to get a fresh captcha image..."
        )

        refreshed = try_refresh_captcha_image(page)
        if not refreshed:
            break

    if best_fallback:
        print(f"No structurally valid captcha read; using best guess: '{best_fallback}'")
        return best_fallback

    raise CaptchaUnsolvedError("Could not obtain any usable CAPTCHA reading.")


# ============================================================
# PERMANENT POPUP & OVERLAY IMMUNIZATION
# ============================================================
def immunize_page_against_popups(context):
    """
    Disables driver.js (welcome tour) and injects style blockades for stray
    overlays, plus a lightweight interval that keeps stripping the
    'driver-active' / 'driver-fade' body/html classes and known overlay
    elements if they reappear.

    IMPORTANT: this intentionally does NOT use a MutationObserver watching
    the whole document (subtree + childList). That was tried first and
    caused a worse regression: on this portal's page (which does its own
    AJAX/DOM churn via postbacks and update panels), a document-wide
    observer meant every unrelated DOM change anywhere on the page fired
    our callback, which itself mutates the DOM (removing elements),
    re-triggering the observer -- a feedback loop that pegged the page's
    JS thread and made the whole page look "stuck loading" and
    unresponsive, even before any of our own script's actions ran. A
    plain, modestly-scoped interval avoids that feedback loop while still
    catching the overlay if it reappears.
    """
    context.add_init_script("""
        window.Driver = function() {
            return {
                start: function() {},
                reset: function() {},
                destroy: function() {},
                defineSteps: function() {},
                showNextStep: function() {},
                preventMove: function() {}
            };
        };

        const style = document.createElement('style');
        style.innerHTML = `
            .driver-overlay, .driver-popover, .driver-popover-tip,
            .driver-popover-footer, .driver-popover-description,
            .driver-popover-title, .driver-stage-no-animation,
            .modal-backdrop, .overlay {
                display: none !important;
                visibility: hidden !important;
                pointer-events: none !important;
                width: 0 !important;
                height: 0 !important;
                opacity: 0 !important;
            }
            body, html {
                overflow: auto !important;
                pointer-events: auto !important;
            }
        `;
        document.head ? document.head.appendChild(style) : document.addEventListener('DOMContentLoaded', () => document.head.appendChild(style));

        const stripOverlayState = () => {
            if (document.body && document.body.classList.contains('driver-active')) {
                document.body.classList.remove('driver-active', 'driver-fade');
                document.body.style.overflow = 'auto';
                document.body.style.pointerEvents = 'auto';
            }
            if (document.documentElement && document.documentElement.classList.contains('driver-active')) {
                document.documentElement.classList.remove('driver-active', 'driver-fade');
            }
            const strayOverlays = document.querySelectorAll(
                '.driver-overlay, .driver-popover, .driver-popover-tip, .driver-popover-footer, .modal-backdrop, .overlay'
            );
            if (strayOverlays.length > 0) {
                strayOverlays.forEach((el) => el.remove());
            }
        };

        const startPolling = () => {
            stripOverlayState();
            // Plain interval only -- no MutationObserver. See docstring above.
            setInterval(stripOverlayState, 400);
        };

        if (document.body) {
            startPolling();
        } else {
            document.addEventListener('DOMContentLoaded', startPolling);
        }
    """)


def clean_popups_now(page):
    try:
        page.evaluate("""
            document.querySelectorAll('.driver-overlay, .driver-popover, .driver-popover-tip, .driver-popover-footer, .modal-backdrop, .overlay').forEach(el => el.remove());
            if (document.body) {
                document.body.classList.remove('driver-active', 'driver-fade');
                document.body.style.overflow = 'auto';
                document.body.style.pointerEvents = 'auto';
            }
            if (document.documentElement) {
                document.documentElement.classList.remove('driver-active', 'driver-fade');
                document.documentElement.style.overflow = 'auto';
            }
        """)
    except Exception:
        pass


def handle_welcome_popup(page):
    """Belt-and-braces close of the Welcome Tour, mirroring the framework's dashboard page."""
    try:
        page.wait_for_timeout(1000)
        popup = page.locator(".driver-popover")
        close_btn = page.locator(".driver-popover-close-btn")
        if popup.count() > 0 and popup.is_visible():
            print("Welcome Tour detected, closing...")
            close_btn.evaluate("(el) => el.click()")
            popup.wait_for(state="hidden", timeout=5000)
            print("Welcome Tour closed.")
    except Exception as exc:
        print(f"Unable to close Welcome Tour (continuing): {exc}")

    clean_popups_now(page)


def safe_click(page, locator_or_selector):
    clean_popups_now(page)

    loc = (
        locator_or_selector
        if hasattr(locator_or_selector, "click")
        else page.locator(locator_or_selector)
    )

    try:
        loc.scroll_into_view_if_needed(timeout=5000)
    except Exception:
        pass

    try:
        loc.click(timeout=10000, force=True)
    except Exception as first_exc:
        try:
            loc.click(timeout=10000)
        except Exception as second_exc:
            raise RuntimeError(f"Could not click portal control: {second_exc}") from first_exc


def check_rate_limit(page) -> None:
    try:
        current_url = page.url
    except Exception:
        current_url = ""

    if "SessionTimeOut.aspx" in current_url or "sessionexpired=true" in current_url.lower():
        raise SessionExpiredError(
            "SEZ Online session has expired (redirected to SessionTimeOut.aspx)."
        )

    try:
        body_text = page.locator("body").inner_text(timeout=2000)
    except Exception:
        return

    normalized = re.sub(r"\s+", " ", body_text or "").strip().lower()

    if "session has expired" in normalized or "session expired" in normalized:
        raise SessionExpiredError(
            "SEZ Online session has expired ('Your Session has expired' banner detected)."
        )

    if "rate limit exceeded" in normalized or "try again in 120 seconds" in normalized:
        raise RateLimitError(
            "SEZ Online rate limit exceeded. The portal requested a retry after 120 seconds."
        )


def recover_from_session_expiry(page):
    print()
    print("=" * 60)
    print("SESSION EXPIRED — RE-AUTHENTICATING")
    print("=" * 60)
    print("Current URL:", page.url)

    try:
        here_link = page.get_by_role("link", name="Here")
        if here_link.count() > 0:
            print("Clicking 'Here' to return to login...")
            here_link.first.click(timeout=5000)
            page.wait_for_load_state("domcontentloaded", timeout=15000)
            page.wait_for_timeout(500)
    except Exception as exc:
        print(f"Could not click 'Here' link (continuing anyway): {exc}")

    login(page)

    print("Re-authentication after session expiry: SUCCESSFUL.")
    print("=" * 60)


def login_failed_reason(page) -> str:
    if "Login.aspx" in page.url:
        try:
            body_text = page.locator("body").inner_text(timeout=2000)
        except Exception:
            body_text = ""

        normalized = re.sub(r"\s+", " ", body_text or "").strip().lower()

        if "invalid captcha" in normalized or "captcha" in normalized and "invalid" in normalized:
            return "invalid captcha"
        if "invalid" in normalized and ("username" in normalized or "password" in normalized):
            return "invalid credentials"
        return "still on login page"

    return ""


# ============================================================
# LOGIN FLOW (same retry logic as the print-automation script)
# ============================================================
def login(page):
    print()
    print("=" * 60)
    print("SEZ ONLINE LOGIN (AUTOMATED OCR)")
    print("=" * 60)

    for attempt in range(1, MAX_LOGIN_ATTEMPTS + 1):
        print(f"\n--- Login attempt {attempt}/{MAX_LOGIN_ATTEMPTS} ---")

        print("Opening login page...")
        page.goto(BASE_URL, wait_until="domcontentloaded")
        clean_popups_now(page)

        username_input = page.get_by_role("textbox", name="User Name")
        password_input = page.get_by_role("textbox", name="Password")
        captcha_input = page.get_by_role("textbox", name="Captcha")
        login_button = page.get_by_role("button", name="Login")
        captcha_image = page.locator("#captchaControlRevamp_imgCaptcha")

        print("Entering username & password...")
        username_input.wait_for(state="visible", timeout=30000)
        username_input.fill(USERNAME)
        password_input.fill(PASSWORD)

        try:
            captcha_token = get_captcha_token(page, captcha_image)
        except CaptchaUnsolvedError as exc:
            print(f"Captcha could not be solved this attempt: {exc}")
            continue

        print(f"Entering extracted CAPTCHA: '{captcha_token}'")
        captcha_input.fill(captcha_token)

        print("Submitting login form...")
        clean_popups_now(page)
        safe_click(page, login_button)

        try:
            page.wait_for_load_state("domcontentloaded", timeout=15000)
        except Exception:
            pass

        page.wait_for_timeout(1500)
        clean_popups_now(page)
        print("Login submission completed. Current URL:", page.url)

        if "ConfirmUserCredentials.aspx" in page.url:
            print("Handling ConfirmUserCredentials screen...")
            try:
                confirm_btn = page.locator(
                    "input[type='submit'][value*='Continue'], "
                    "input[type='submit'][value*='OK'], "
                    "#btnSubmit, #btnConfirm"
                )
                if confirm_btn.count() > 0:
                    safe_click(page, confirm_btn.first)
                    page.wait_for_load_state("domcontentloaded", timeout=10000)
                    page.wait_for_timeout(1000)
            except Exception:
                pass

        failure_reason = login_failed_reason(page)
        if not failure_reason:
            print("Login SUCCESSFUL.")
            return

        print(f"Login attempt {attempt} failed ({failure_reason}). Retrying with a fresh captcha...")

    raise RuntimeError(
        f"Could not log in to SEZ Online after {MAX_LOGIN_ATTEMPTS} attempts "
        "(captcha and/or credential failures)."
    )


# ============================================================
# MASKED DATE FIELD HELPER
# ============================================================
def enter_masked_date(page, locator, date_str: str):
    """
    Clears a masked date field and types the value character-by-character,
    which is what the portal's JS date mask expects (a plain .fill() can
    leave the mask in a half-applied state).

    Routes the click through safe_click() (popup cleanup + forced click)
    instead of a raw locator.click(). A raw click here was the concrete
    cause of the "stuck / scrolling" hang: if the tour overlay reasserted
    itself right before this specific click, nothing cleared it first, so
    Playwright retried the click against a body that kept intercepting
    pointer events until the 30s timeout.
    """
    clean_popups_now(page)
    safe_click(page, locator)
    locator.evaluate("(el) => el.value = ''")
    if date_str:
        locator.type(date_str, delay=50)


# ============================================================
# EXPORT TO EXCEL (equivalent of navigate_to_export_excel)
# ============================================================
def navigate_to_export_excel(page, from_date: str, to_date: str, download_dir: Path, modules):
    """
    Reports -> Trade Data -> for each module: select module, set date range
    (first pass only), click EXPORT TO EXCEL, capture the browser download,
    save it to download_dir. Mirrors the framework's SEZDashboardPage
    method, with the same rate-limit / session-expiry checks used by the
    print-automation script layered in around each step.
    """
    downloaded_files = []

    print()
    print("=" * 60)
    print("NAVIGATE TO EXPORT EXCEL")
    print("=" * 60)

    check_rate_limit(page)

    try:
        page.wait_for_load_state("domcontentloaded", timeout=15000)
    except Exception:
        pass

    handle_welcome_popup(page)
    check_rate_limit(page)

    print("Opening Reports...")
    reports_link = page.get_by_role("link", name="Reports", exact=True)
    reports_link.wait_for(state="visible", timeout=30000)
    check_rate_limit(page)
    safe_click(page, reports_link)

    try:
        page.wait_for_load_state("domcontentloaded", timeout=15000)
    except Exception:
        pass
    page.wait_for_timeout(500)
    clean_popups_now(page)
    check_rate_limit(page)

    print("Opening Trade Data...")
    trade_data_link = page.locator("#ctl00_DefaultContent_hpTradeData")
    trade_data_link.wait_for(state="visible", timeout=30000)
    check_rate_limit(page)
    safe_click(page, trade_data_link)

    try:
        page.wait_for_load_state("domcontentloaded", timeout=15000)
    except Exception:
        pass
    page.wait_for_timeout(800)
    clean_popups_now(page)
    check_rate_limit(page)

    module_dropdown = page.locator("#ctl00_DefaultContent_ddlModuleName")
    from_date_input = page.locator("#ctl00_DefaultContent_dtFromDt_txtMaskedDate")
    to_date_input = page.locator("#ctl00_DefaultContent_dtToDate_txtMaskedDate")
    export_button = page.get_by_role("button", name="EXPORT TO EXCEL")
    logout_button = page.locator("#ctl00_LoginStatus")

    for index, (module_value, module_label) in enumerate(modules):
        print(f"Selecting module: {module_label} (value={module_value})...")
        module_dropdown.wait_for(state="visible", timeout=15000)
        check_rate_limit(page)
        module_dropdown.select_option(module_value)

        # The module dropdown triggers an ASP.NET postback on this screen,
        # so wait for the page to actually settle (not just a fixed sleep)
        # before touching anything else on it.
        try:
            page.wait_for_load_state("domcontentloaded", timeout=15000)
        except Exception:
            pass
        page.wait_for_timeout(800)
        clean_popups_now(page)

        # Dates only need to be (re)entered once; they persist across
        # module switches on this screen, matching the framework's flow.
        if index == 0:
            print(f"Setting date range: {from_date} to {to_date}...")
            from_date_input.wait_for(state="visible", timeout=15000)
            to_date_input.wait_for(state="visible", timeout=15000)
            check_rate_limit(page)
            enter_masked_date(page, from_date_input, from_date)
            enter_masked_date(page, to_date_input, to_date)
            page.wait_for_timeout(300)

        print(f"Exporting {module_label} to Excel...")
        export_button.wait_for(state="visible", timeout=15000)
        check_rate_limit(page)

        with page.expect_download(timeout=30000) as download_info:
            safe_click(page, export_button)

        download = download_info.value
        suggested_name = download.suggested_filename
        file_path = download_dir / suggested_name

        # Avoid clobbering a same-named file from a previous run/module.
        if file_path.exists():
            stem, suffix = file_path.stem, file_path.suffix
            counter = 1
            while file_path.exists():
                file_path = download_dir / f"{stem}_{counter}{suffix}"
                counter += 1

        download.save_as(str(file_path))
        print(f"Downloaded: {file_path}")
        downloaded_files.append(str(file_path))

        page.wait_for_timeout(1000)
        clean_popups_now(page)
        check_rate_limit(page)

    print("Logging out...")
    logout_button.wait_for(state="visible", timeout=30000)
    safe_click(page, logout_button)

    try:
        page.wait_for_load_state("domcontentloaded", timeout=15000)
    except Exception:
        pass
    page.wait_for_timeout(1000)

    print("Logged out successfully.")
    print("=" * 60)
    print("EXPORT COMPLETE — Files downloaded:")
    for f in downloaded_files:
        print(" -", f)
    print("=" * 60)

    return downloaded_files


# ============================================================
# MAIN ENTRY POINT
# ============================================================
def main():
    print()
    print("=" * 60)
    print("SEZ EXPORT-TO-EXCEL AUTOMATION")
    print("=" * 60)

    with sync_playwright() as p:
        browser = p.chromium.launch(
            headless=False,
            args=["--start-maximized", "--no-sandbox"],
        )

        context = browser.new_context(
            no_viewport=True,
            accept_downloads=True,
            user_agent=(
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/120.0.0.0 Safari/537.36"
            ),
        )

        immunize_page_against_popups(context)

        def auto_accept_dialogs(dialog):
            print(f"\n[Dialog Intercepted] {dialog.message}")
            dialog.accept()

        context.on("dialog", auto_accept_dialogs)

        page = context.new_page()

        try:
            # STEP 1: LOGIN (with internal captcha/login retries)
            login(page)

            # Diagnostic snapshot right after login, before anything else
            # runs. If the page hangs/looks stuck on the NEXT step, this
            # gives us a picture of exactly what state it was actually in,
            # independent of whether the browser later gets closed manually.
            try:
                page.wait_for_load_state("networkidle", timeout=15000)
            except Exception as wait_exc:
                print(f"Post-login networkidle wait did not complete (continuing): {wait_exc}")

            try:
                post_login_path = SCREENSHOT_DIR / "post_login.png"
                page.screenshot(path=str(post_login_path), full_page=True)
                print(f"Saved post-login diagnostic screenshot to: {post_login_path}")
            except Exception as shot_exc:
                print(f"Could not save post-login screenshot: {shot_exc}")

            # STEP 2: EXPORT, with rate-limit / session-expiry retry loop
            session_retries = 0

            while True:
                try:
                    check_rate_limit(page)
                    navigate_to_export_excel(page, FROM_DATE, TO_DATE, DOWNLOAD_DIR, MODULES_TO_EXPORT)
                    break

                except SessionExpiredError as session_exc:
                    session_retries += 1
                    print()
                    print("=" * 60)
                    print("SEZ ONLINE SESSION EXPIRED")
                    print("=" * 60)
                    print(f"Reason: {session_exc}")
                    print(f"Re-login attempt {session_retries}/5")

                    if session_retries > 5:
                        raise RuntimeError(
                            "Session kept expiring during export after 5 re-login attempts."
                        ) from session_exc

                    recover_from_session_expiry(page)
                    print("Retrying export after re-login...")
                    continue

                except RateLimitError as rate_exc:
                    print()
                    print("=" * 60)
                    print("SEZ ONLINE RATE LIMIT DETECTED")
                    print("=" * 60)
                    print(f"Reason: {rate_exc}")
                    print(f"Waiting {RATE_LIMIT_WAIT_SECONDS} seconds before retrying...")

                    idle_keepalive_pause(page, RATE_LIMIT_WAIT_SECONDS)

                    try:
                        page.reload(wait_until="domcontentloaded", timeout=30000)
                    except Exception as reload_exc:
                        print("Page reload warning:", reload_exc)

                    page.wait_for_timeout(1500)
                    clean_popups_now(page)

                    print("Retrying export from the beginning...")
                    continue

        except Exception as exc:
            print()
            print("=" * 60)
            print("SEZ EXPORT PIPELINE FAILED")
            print("=" * 60)
            print(f"{type(exc).__name__}: {exc}")

            failure_path = SCREENSHOT_DIR / "export_failure.png"
            page.screenshot(path=str(failure_path), full_page=True)
            print("Saved failure diagnostic screenshot to:", failure_path)

            raise

        finally:
            context.close()
            browser.close()

    print()
    print("=" * 60)
    print("SEZ EXPORT-TO-EXCEL AUTOMATION COMPLETED")
    print("=" * 60)


if __name__ == "__main__":
    main()