app.py

"""
Gradio UI for the Textract + Comprehend Medical PDF redaction workflow.

Phase 1 (Review): Upload a PDF → runs Textract + Comprehend Medical →
  returns a highlighted review PDF and an editable redactions JSON.

Phase 2 (Redact): Optionally edit the JSON to mark false positives as
  excluded, then apply permanent redactions and download the final PDF.
"""

import importlib.util
import json
import logging
import os
import queue
import tempfile
import threading
import time
import uuid
import webbrowser
from pathlib import Path

import boto3
import fitz  # PyMuPDF — already a transitive dep via the redaction module
import gradio as gr

# ── Load the core redaction module (hyphenated filename) ────────────────────────
_spec = importlib.util.spec_from_file_location(
    "redaction",
    Path(__file__).parent / "comprehend-medical-redaction-text.py",
)
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)

get_textract_blocks = _mod.get_textract_blocks
build_page_data = _mod.build_page_data
collect_redactions = _mod.collect_redactions
generate_review_pdf = _mod.generate_review_pdf
apply_final_redactions = _mod.apply_final_redactions

# ── Usage logging ───────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [USAGE] %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%SZ",
)
_log = logging.getLogger("usage")


def _event(name: str, session: str = "", **kwargs):
    """Log a structured usage event. No PII, filenames, or AWS identifiers."""
    parts = [f"event={name}"]
    if session:
        parts.append(f"session={session}")
    parts.extend(f"{k}={v}" for k, v in kwargs.items() if v is not None)
    _log.info(" ".join(parts))


# ── Configuration ───────────────────────────────────────────────────────────────
REGION = "us-east-1"
POLL_INTERVAL = 15
DEFAULT_MAX_PAGES = 20

# Set SHOW_AUTH=false in the App Runner environment variables for single-account
# deployments where the task role provides credentials. Set to true (default)
# for multi-account deployments where users authenticate via IAM Identity Center.
SHOW_AUTH = os.environ.get("SHOW_AUTH", "true").lower() != "false"


def _parse_s3_url(s3_url: str) -> tuple[str, str]:
    """Parse 's3://bucket/prefix/' into ('bucket', 'prefix'). Prefix may be empty."""
    s3_url = s3_url.strip().rstrip("/")
    if not s3_url.startswith("s3://"):
        raise ValueError("Expected format: s3://bucket-name/optional-prefix/")
    path = s3_url[5:]
    bucket, _, prefix = path.partition("/")
    if not bucket:
        raise ValueError("Bucket name cannot be empty")
    return bucket, prefix


# ── AWS session helpers ─────────────────────────────────────────────────────────


def _make_boto3_session(creds: dict) -> boto3.Session:
    """Build a boto3 Session from stored credentials state."""
    if creds.get("type") == "static":
        return boto3.Session(
            aws_access_key_id=creds["key_id"],
            aws_secret_access_key=creds["secret_key"],
            aws_session_token=creds.get("session_token") or None,
            region_name=creds.get("region") or REGION,
        )
    if creds.get("type") == "profile":
        return boto3.Session(
            profile_name=creds.get("profile") or None,
            region_name=creds.get("region") or REGION,
        )
    return boto3.Session(region_name=creds.get("region") or REGION)


def _make_clients(creds: dict):
    session = _make_boto3_session(creds)
    return (
        session.client("s3"),
        session.client("textract"),
        session.client("comprehendmedical"),
    )


def _caller_identity(creds: dict) -> str:
    """Return the caller ARN string, or raise on failure."""
    session = _make_boto3_session(creds)
    sts = session.client("sts", region_name=creds.get("region") or REGION)
    return sts.get_caller_identity()["Arn"]


# ── Auth actions ────────────────────────────────────────────────────────────────


def check_connection(creds: dict):
    try:
        arn = _caller_identity(creds)
        return f"**Status:** Connected as `{arn}`", creds
    except Exception as exc:
        return f"**Status:** Not connected — {exc}", creds


def apply_static_creds(key_id: str, secret: str, token: str, region: str, creds: dict, session_id: str = ""):
    if not key_id.strip() or not secret.strip():
        return (
            "Access Key ID and Secret Access Key are required.",
            creds,
            "**Status:** Not connected",
        )
    new_creds = {
        "type": "static",
        "key_id": key_id.strip(),
        "secret_key": secret.strip(),
        "session_token": token.strip() or None,
        "region": region.strip() or REGION,
    }
    try:
        arn = _caller_identity(new_creds)
        _event("logged_in", session_id, method="static")
        msg = f"**Status:** Connected as `{arn}`"
        return msg, new_creds, msg
    except Exception as exc:
        return f"Invalid credentials: {exc}", creds, "**Status:** Not connected"


_AUTH_LINK_EMPTY = ""
_AUTH_LINK_TMPL = """
<div style="margin:8px 0;padding:12px 16px;background:#eff6ff;border-radius:6px;
            border-left:4px solid #2563eb;font-family:sans-serif;">
  <strong>Authorization required</strong><br><br>
  <a href="{url}" target="_blank"
     style="display:inline-block;padding:8px 16px;background:#2563eb;color:#fff;
            border-radius:4px;text-decoration:none;font-weight:600;">
    Authorize in AWS &rarr;
  </a>
  <br><br>
  <small style="color:#555;">
    If the button above doesn't work, copy this URL into your browser:<br>
    <code style="word-break:break-all;">{url}</code>
  </small>
</div>
"""


def start_sso_device_flow(
    start_url: str,
    sso_region: str,
    account_id: str,
    role_name: str,
    region: str,
    creds: dict,
    session_id: str = "",
):
    """
    Generator implementing the OAuth 2.0 device authorization grant against
    AWS IAM Identity Center — no AWS CLI required.

    Yields: (log_text, auth_link_html, creds_state, auth_status_markdown)
    """
    start_url = start_url.strip()
    sso_region = sso_region.strip()
    account_id = account_id.strip()
    role_name = role_name.strip()
    region = region.strip() or REGION

    logs: list[str] = []

    def emit(*lines):
        logs.extend(lines)
        return "\n".join(logs)

    missing = [
        f
        for f, v in [
            ("SSO Start URL", start_url),
            ("Identity Center Region", sso_region),
            ("Account ID", account_id),
            ("Role Name", role_name),
        ]
        if not v
    ]
    if missing:
        yield (
            emit(f"Please fill in: {', '.join(missing)}"),
            _AUTH_LINK_EMPTY,
            creds,
            "**Status:** Not connected",
        )
        return

    try:
        oidc = boto3.client("sso-oidc", region_name=sso_region)

        # Step 1: register a public client (valid ~90 days; result could be cached)
        client_reg = oidc.register_client(
            clientName="document-redaction-agent",
            clientType="public",
        )

        # Step 2: start device authorization → get the verification URL
        device_auth = oidc.start_device_authorization(
            clientId=client_reg["clientId"],
            clientSecret=client_reg["clientSecret"],
            startUrl=start_url,
        )
    except Exception as exc:
        yield (
            emit(f"ERROR starting device authorization: {exc}"),
            _AUTH_LINK_EMPTY,
            creds,
            "**Status:** Error",
        )
        return

    verification_url = device_auth["verificationUriComplete"]
    device_code = device_auth["deviceCode"]
    interval = max(device_auth.get("interval", 5), 5)
    expires_in = device_auth.get("expiresIn", 600)
    deadline = time.time() + expires_in

    # Try to open the browser automatically (works when running locally)
    webbrowser.open(verification_url)

    link_html = _AUTH_LINK_TMPL.format(url=verification_url)
    yield (
        emit("Waiting for you to authorize in the browser ..."),
        link_html,
        creds,
        "**Status:** Waiting for authorization ...",
    )

    # Step 3: poll until the user completes auth or the code expires
    while time.time() < deadline:
        time.sleep(interval)
        try:
            token = oidc.create_token(
                clientId=client_reg["clientId"],
                clientSecret=client_reg["clientSecret"],
                grantType="urn:ietf:params:oauth:grant-type:device_code",
                deviceCode=device_code,
            )
        except oidc.exceptions.AuthorizationPendingException:
            remaining = int(deadline - time.time())
            yield (
                emit(f"    Still waiting ... ({remaining}s remaining)"),
                link_html,
                creds,
                "**Status:** Waiting for authorization ...",
            )
            continue
        except oidc.exceptions.SlowDownException:
            interval += 5
            continue
        except Exception as exc:
            yield (
                emit(f"ERROR during token poll: {exc}"),
                link_html,
                creds,
                "**Status:** Error",
            )
            return

        # Step 4: exchange the access token for temporary IAM credentials
        sso = boto3.client("sso", region_name=sso_region)
        access_token = token["accessToken"]
        try:
            role_creds = sso.get_role_credentials(
                accountId=account_id,
                roleName=role_name,
                accessToken=access_token,
            )["roleCredentials"]
        except Exception as exc:
            # Authentication succeeded but the account/role is wrong or not assigned.
            # List what this user actually has access to so they can correct the fields.
            yield (
                emit(
                    f"ERROR fetching role credentials: {exc}",
                    "",
                    "Checking what your user has access to ...",
                ),
                link_html,
                creds,
                "**Status:** Error",
            )
            try:
                accounts = sso.list_accounts(accessToken=access_token).get(
                    "accountList", []
                )
                if not accounts:
                    emit("  (no accounts found for this Identity Center user)")
                else:
                    emit("", "Your Identity Center user has access to:")
                    for acct in accounts:
                        acct_id = acct["accountId"]
                        acct_name = acct.get("accountName", acct_id)
                        emit(f"  Account: {acct_name}  (ID: {acct_id})")
                        try:
                            roles = sso.list_account_roles(
                                accessToken=access_token, accountId=acct_id
                            ).get("roleList", [])
                            for r in roles:
                                emit(f"    Permission Set: {r['roleName']}")
                        except Exception:
                            emit("    (could not list roles for this account)")
                emit(
                    "",
                    "Update the Account ID and Permission Set fields above and try again.",
                )
            except Exception as list_exc:
                emit(f"  Could not list accessible accounts: {list_exc}")
            yield (
                "\n".join(logs),
                link_html,
                creds,
                "**Status:** Error — see log for available accounts",
            )
            return

        new_creds = {
            "type": "static",
            "key_id": role_creds["accessKeyId"],
            "secret_key": role_creds["secretAccessKey"],
            "session_token": role_creds["sessionToken"],
            "region": region,
        }
        try:
            arn = _caller_identity(new_creds)
        except Exception as exc:
            yield (
                emit(f"Credentials retrieved but identity check failed: {exc}"),
                link_html,
                creds,
                "**Status:** Error",
            )
            return

        _event("logged_in", session_id, method="sso")
        status = f"**Status:** Connected as `{arn}`"
        yield emit(f"Connected as {arn}"), _AUTH_LINK_EMPTY, new_creds, status
        return

    yield (
        emit("Authorization timed out. Click the button again to restart."),
        _AUTH_LINK_EMPTY,
        creds,
        "**Status:** Timed out",
    )


# ── Queue / threading helpers ───────────────────────────────────────────────────


def _make_work_dir() -> Path:
    return Path(tempfile.mkdtemp(prefix="redact_"))


def _stream_worker(fn, log_q: queue.Queue):
    try:
        fn()
        log_q.put(("done", None))
    except Exception as exc:
        log_q.put(("error", str(exc)))


def _drain_queue(log_q: queue.Queue, logs: list, timeout=120):
    while True:
        try:
            kind, data = log_q.get(timeout=timeout)
        except queue.Empty:
            logs.append("(timed out waiting for response)")
            yield True, "\n".join(logs)
            return
        if kind == "log":
            logs.append(data)
            yield False, "\n".join(logs)
        elif kind == "done":
            yield True, "\n".join(logs)
            return
        elif kind == "error":
            logs.append(f"\nERROR: {data}")
            yield True, "\n".join(logs)
            return


# ── Phase 1: Review ─────────────────────────────────────────────────────────────


def run_review(pdf_upload, creds: dict, s3_url: str, max_pages: int, session_id: str = ""):
    if pdf_upload is None:
        yield "Please upload a PDF first.", None, "", "", ""
        return

    try:
        s3_bucket, s3_prefix = _parse_s3_url(s3_url)
    except ValueError as exc:
        yield f"Invalid S3 URL in Settings: {exc}", None, "", "", ""
        return

    max_pages = int(max_pages) if max_pages and int(max_pages) > 0 else None
    _event("review_started", session_id)
    unique_key = "/".join(
        filter(None, [s3_prefix, uuid.uuid4().hex[:8], Path(pdf_upload).name])
    )

    work_dir = _make_work_dir()
    pdf_path = Path(pdf_upload)
    local_pdf = work_dir / pdf_path.name
    local_pdf.write_bytes(pdf_path.read_bytes())

    textract_cache = work_dir / "_textract_cache.json"
    review_pdf = work_dir / f"{local_pdf.stem}-review.pdf"
    redactions_json = work_dir / "redactions.json"
    if max_pages:
        with fitz.open(str(local_pdf)) as _doc:
            actual_pages = min(max_pages, len(_doc))
        source_pdf = work_dir / f"_subset_{actual_pages}pages.pdf"
    else:
        source_pdf = local_pdf

    log_q: queue.Queue = queue.Queue()
    logs: list = []
    result: dict = {}

    def log(msg):
        log_q.put(("log", msg))

    def worker():
        s3, textract, cm_client = _make_clients(creds)
        blocks = get_textract_blocks(
            s3,
            textract,
            local_pdf,
            s3_bucket,
            unique_key,
            textract_cache,
            max_pages,
            POLL_INTERVAL,
            log=log,
        )
        page_data = build_page_data(blocks)
        log(f"    {len(page_data)} pages indexed.")
        redactions = collect_redactions(page_data, cm_client, log=log)
        redactions_json.write_text(
            json.dumps({"redactions": redactions}, indent=2), encoding="utf-8"
        )
        log(f"    Saved: {redactions_json}")
        generate_review_pdf(redactions, source_pdf, review_pdf, log=log)
        result["review_pdf"] = str(review_pdf)
        result["redactions_text"] = redactions_json.read_text(encoding="utf-8")
        result["work_dir"] = str(work_dir)
        result["source_pdf"] = str(source_pdf)

    thread = threading.Thread(target=_stream_worker, args=(worker, log_q))
    thread.start()

    for _done, log_text in _drain_queue(log_q, logs):
        yield (
            log_text,
            result.get("review_pdf"),
            result.get("redactions_text", ""),
            result.get("work_dir", ""),
            result.get("source_pdf", ""),
        )

    thread.join()


# ── Redactions editor: track first adjustment ────────────────────────────────────


def mark_adjusted(text: str, has_adjusted: bool, session_id: str = ""):
    """Log once per session when the user first edits the redactions JSON."""
    if not has_adjusted and text.strip():
        _event("redactions_adjusted", session_id)
        return True
    return has_adjusted


# ── Phase 2: Redact ─────────────────────────────────────────────────────────────


def run_redact(redactions_text: str, work_dir_str: str, source_pdf_str: str, session_id: str = ""):
    if not work_dir_str:
        yield "Run the Review phase first.", None
        return
    if not redactions_text.strip():
        yield "Redactions JSON is empty.", None
        return
    try:
        json.loads(redactions_text)
    except json.JSONDecodeError as exc:
        yield f"Invalid JSON: {exc}", None
        return

    work_dir = Path(work_dir_str)
    source_pdf = Path(source_pdf_str)
    redactions_json = work_dir / "redactions.json"
    redactions_json.write_text(redactions_text, encoding="utf-8")
    output_pdf = work_dir / f"{source_pdf.stem}-redacted.pdf"

    log_q: queue.Queue = queue.Queue()
    logs: list = []
    result: dict = {}

    def log(msg):
        log_q.put(("log", msg))

    def worker():
        t0 = time.time()
        apply_final_redactions(redactions_json, source_pdf, output_pdf, log=log)
        result["output_pdf"] = str(output_pdf)
        _event("redact_completed", session_id, duration_s=int(time.time() - t0))

    thread = threading.Thread(target=_stream_worker, args=(worker, log_q))
    thread.start()

    for _done, log_text in _drain_queue(log_q, logs):
        yield log_text, result.get("output_pdf")

    thread.join()


# ── Gradio UI ───────────────────────────────────────────────────────────────────

with gr.Blocks(title="Document Redaction Agent") as demo:
    gr.HTML("<h1 style='text-align:center;'> PDF Redaction Agent </h1>")
    gr.Markdown(
        "<center>This agent utilizes Textract and AWS Comprehend to extract content from PDFs and "
        "then apply redactions. We use a two phase approach outlined below. Please note "
        "that comprehend is not available in all regions.</center>"
    )
    gr.Markdown(
        "**Begin:** Start by logging into your AWS account via Identity Center or Static Credentials."
        " Make sure that the user has access to S3, Textract and Comprehend and is in a region where all services are supported.  \n"
        "**Phase 1:** Upload a PDF to detect PHI and generate a highlighted review PDF.  \n"
        "**Phase 2:** Edit the redactions JSON to exclude false positives, then apply "
        "permanent redactions."
    )

    # Session state
    creds_state = gr.State({})
    work_dir_state = gr.State("")
    source_pdf_state = gr.State("")
    session_id_state = gr.State(lambda: uuid.uuid4().hex[:8])
    has_adjusted_state = gr.State(False)

    # ── AWS Authentication ──────────────────────────────────────────────────────
    with gr.Accordion("AWS Authentication", open=True, visible=SHOW_AUTH):
        auth_status = gr.Markdown("**Status:** Unknown — click Check Connection")
        check_btn = gr.Button("Check Connection", size="sm")

        with gr.Tab("SSO / Identity Center Login"):
            gr.Markdown(
                "Fill in your IAM Identity Center details and click **Login with AWS SSO**. "
                "A browser window will open for you to authorize — no AWS CLI required. "
                "Find these values in the AWS Console under **IAM Identity Center → Settings**."
            )
            with gr.Row():
                sso_start_url_in = gr.Textbox(
                    label="SSO Start URL",
                    placeholder="https://my-portal.awsapps.com/start",
                    scale=3,
                )
                sso_idc_region_in = gr.Textbox(
                    label="Identity Center Region",
                    placeholder="us-east-1",
                    scale=1,
                )
            with gr.Row():
                sso_account_id_in = gr.Textbox(
                    label="AWS Account ID",
                    placeholder="123456789012",
                    scale=1,
                )
                sso_role_name_in = gr.Textbox(
                    label="Permission Set / Role Name",
                    placeholder="AdministratorAccess",
                    scale=2,
                )
            sso_region_in = gr.Textbox(
                label="Working Region (for S3 / Textract / Comprehend)",
                value=REGION,
            )
            sso_btn = gr.Button("Login with AWS SSO", variant="primary")
            sso_auth_link = gr.HTML(_AUTH_LINK_EMPTY)
            sso_log = gr.Textbox(
                label="Login Progress",
                lines=6,
                interactive=False,
                )

        with gr.Tab("Static Credentials"):
            gr.Markdown(
                "Enter AWS credentials directly. Temporary credentials (from `aws sts "
                "assume-role` or the SSO credential file) require the Session Token."
            )
            with gr.Row():
                key_id_in = gr.Textbox(label="Access Key ID", scale=2)
                region_in = gr.Textbox(label="Region", value=REGION, scale=1)
            secret_in = gr.Textbox(label="Secret Access Key", type="password")
            token_in = gr.Textbox(label="Session Token (optional)")
            static_btn = gr.Button("Apply Credentials", variant="primary")
            static_msg = gr.Markdown("")

    # ── Settings ────────────────────────────────────────────────────────────────
    with gr.Accordion("Settings", open=True):
        s3_url_in = gr.Textbox(
            label="S3 Staging Location",
            placeholder="s3://your-bucket-name/textract-input/",
            info=(
                "The S3 folder where the PDF will be uploaded for Textract to process. "
                "Must be in the same AWS account and region as your credentials. "
                "Format: s3://bucket-name/optional-prefix/"
            ),
        )
        max_pages_in = gr.Number(
            label="Max Pages (0 = entire document)",
            value=DEFAULT_MAX_PAGES,
            minimum=0,
            precision=0,
            info="Limit processing to the first N pages. Useful for large documents or cost control.",
        )

    # ── Phase 1: Review ─────────────────────────────────────────────────────────
    with gr.Tab("Phase 1 — Review"):
        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"], type="filepath")
        review_btn = gr.Button("Run Review", variant="primary")
        review_log = gr.Textbox(
            label="Progress", lines=10, interactive=False, show_copy_button=True
        )
        review_pdf_out = gr.File(label="Review PDF (download)", interactive=False)
        redactions_editor = gr.Textbox(
            label='Redactions JSON — set "exclude": true on false positives',
            lines=20,
            interactive=True,
        )

    # ── Phase 2: Redact ─────────────────────────────────────────────────────────
    with gr.Tab("Phase 2 — Redact"):
        gr.Markdown(
            "Edit the **Redactions JSON** in Phase 1 to mark false positives "
            '(`"exclude": true`), then click **Apply Redactions**.'
        )
        redact_btn = gr.Button("Apply Redactions", variant="primary")
        redact_log = gr.Textbox(
            label="Progress", lines=8, interactive=False, show_copy_button=True
        )
        redacted_pdf_out = gr.File(label="Redacted PDF (download)", interactive=False)

    # ── Wiring ──────────────────────────────────────────────────────────────────
    check_btn.click(
        fn=check_connection,
        inputs=[creds_state],
        outputs=[auth_status, creds_state],
    )

    sso_btn.click(
        fn=start_sso_device_flow,
        inputs=[
            sso_start_url_in,
            sso_idc_region_in,
            sso_account_id_in,
            sso_role_name_in,
            sso_region_in,
            creds_state,
            session_id_state,
        ],
        outputs=[sso_log, sso_auth_link, creds_state, auth_status],
    )

    static_btn.click(
        fn=apply_static_creds,
        inputs=[key_id_in, secret_in, token_in, region_in, creds_state, session_id_state],
        outputs=[static_msg, creds_state, auth_status],
    )

    review_btn.click(
        fn=run_review,
        inputs=[pdf_input, creds_state, s3_url_in, max_pages_in, session_id_state],
        outputs=[
            review_log,
            review_pdf_out,
            redactions_editor,
            work_dir_state,
            source_pdf_state,
        ],
    )

    redactions_editor.change(
        fn=mark_adjusted,
        inputs=[redactions_editor, has_adjusted_state, session_id_state],
        outputs=[has_adjusted_state],
    )

    redact_btn.click(
        fn=run_redact,
        inputs=[redactions_editor, work_dir_state, source_pdf_state, session_id_state],
        outputs=[redact_log, redacted_pdf_out],
    )

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )