#/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()

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

    handle_terminate_sessions(page, run_dir)

    if page.locator("text=MFA code is invalid").count():
        raise RuntimeError("MFA rejected")

    page.wait_for_timeout(600)