#/apps/aroflo_connector_app/ui_automation/auth/mfa.py
import sys
from playwright.sync_api import Page

from ..core.artifacts import shot
from ..core.log import log_step, pause
from .post_login import handle_terminate_sessions



def prompt_mfa_code(arg: str) -> str:
    if arg:
        return arg.strip()
    print("Enter 6-digit MFA code: ", end="", flush=True)
    return sys.stdin.readline().strip()


def handle_mfa(page: Page, cfg, run_dir, mfa_code_arg: str) -> None:
    page.wait_for_function(
        """() => {
            const t = document.body?.innerText || '';
            const p = document.querySelector('#postLoginType');
            return t.includes('Verify Your Account') || (p && p.value === 'verifyMFA');
        }""",
        timeout=30_000,
    )

    shot(page, run_dir, "step4-mfa-blank")
    log_step("step4-mfa-blank", page)

    if cfg.pause_on_mfa:
        pause(cfg, "MFA screen detected")

    code = prompt_mfa_code(mfa_code_arg)
    if len(code) != 6 or not code.isdigit():
        raise RuntimeError("Invalid MFA code")

    inputs = page.locator('input.af-otpInput-input, input[id^="otpInput__input-"]')
    if inputs.count() < 6:
        raise RuntimeError("MFA inputs not found")

    inputs.nth(0).click()
    page.keyboard.type(code, delay=max(cfg.slow_mo_ms, 40))

    shot(page, run_dir, "step5-mfa-filled")
    log_step("step5-mfa-filled", page)

    lbl = page.locator('label[for="mfa_trust_device"]')
    if lbl.count():
        lbl.first.click(force=True)

    shot(page, run_dir, "step6-mfa-trust")
    log_step("step6-mfa-trust", page)

    btn = page.locator('button:has-text("Verify"), form#frmPostLogin button')
    if not btn.count():
        raise RuntimeError("Verify button not found")

    btn.first.click()

    # Navegación puede ocurrir; esperamos por señales sin romper el contexto.
    try:
        page.wait_for_load_state("domcontentloaded", timeout=15_000)
    except Exception:
        pass

    page.wait_for_function(
        """() => {
            const t = document.body?.innerText || '';
            const p = document.querySelector('#postLoginType')?.value || '';
            const invalid = t.includes('MFA code is invalid');
            const terminate = p === 'terminateSessions';
            const inIMS = location.href.includes('/ims/');
            return invalid || terminate || inIMS;
        }""",
        timeout=60_000,
    )

    handle_terminate_sessions(page, run_dir)

    try:
        if page.locator("text=MFA code is invalid").count():
            raise RuntimeError("MFA rejected")
    except Exception:
        # Si el contexto se destruyó por navegación, ignoramos el check puntual.
        pass

    page.wait_for_timeout(600)