"""S3-upload trigger that fans one PDF out to both extraction engines.

An ``s3:ObjectCreated`` event on the input bucket invokes this handler. For
each PDF record it mirrors the inbound key into an output prefix and starts
whichever engines have not yet produced results there:

- the Fargate table job (``quber table``), exactly as before, and
- the RunPod parse job, by minting presigned URLs for the source PDF, the
  three parse artifacts, and a completion marker, then POSTing them to the
  serverless endpoint. The worker holds no AWS credentials; the URLs carry
  the authorization.

Deduplication is per engine: a re-dropped file re-runs only the engine whose
results are missing, which is also the re-drive path after a failure.

For every parse job submitted, a one-shot EventBridge schedule invokes the
join Lambda at the endpoint's execution timeout plus margin. That check is
the safety net for jobs that die without writing anything — a worker killed
mid-run, an execution timeout, or an endpoint idle-decayed to zero workers.

Path mirroring:
    inbound/<company>/<sub>/<file>.pdf  ->  <company>/<sub>/<file>/

Required environment (set by Terraform):
    ECS_CLUSTER, TASK_DEFINITION, CONTAINER_NAME, SUBNETS, SECURITY_GROUPS,
    OUTPUT_BUCKET, RUNPOD_ENDPOINT_ID, RUNPOD_API_KEY_PARAM,
    JOIN_LAMBDA_ARN, SCHEDULER_ROLE_ARN;
    optional INPUT_PREFIX (default "inbound/"), DEADLINE_MINUTES (default 12).
"""

from __future__ import annotations

import json
import logging
import os
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone
from typing import Any

import boto3
from botocore.config import Config

log = logging.getLogger()
log.setLevel(logging.INFO)

ecs = boto3.client("ecs")
s3 = boto3.client("s3", config=Config(signature_version="s3v4"))
ssm = boto3.client("ssm")
scheduler = boto3.client("scheduler")

CLUSTER = os.environ["ECS_CLUSTER"]
TASK_DEFINITION = os.environ["TASK_DEFINITION"]
CONTAINER_NAME = os.environ["CONTAINER_NAME"]
SUBNETS = [s for s in os.environ["SUBNETS"].split(",") if s]
SECURITY_GROUPS = [s for s in os.environ["SECURITY_GROUPS"].split(",") if s]
OUTPUT_BUCKET = os.environ["OUTPUT_BUCKET"]
INPUT_PREFIX = os.environ.get("INPUT_PREFIX", "inbound/")

RUNPOD_ENDPOINT_ID = os.environ["RUNPOD_ENDPOINT_ID"]
RUNPOD_API_KEY_PARAM = os.environ["RUNPOD_API_KEY_PARAM"]
JOIN_LAMBDA_ARN = os.environ["JOIN_LAMBDA_ARN"]
SCHEDULER_ROLE_ARN = os.environ["SCHEDULER_ROLE_ARN"]
DEADLINE_MINUTES = int(os.environ.get("DEADLINE_MINUTES", "12"))

URL_EXPIRY_SECONDS = 2 * 60 * 60  # queue wait + cold start + parse, with slack

runpod_api_key: str | None = None


def get_runpod_api_key() -> str:
    """SSM-stored API key, fetched once per warm container."""
    global runpod_api_key
    if runpod_api_key is None:
        response = ssm.get_parameter(Name=RUNPOD_API_KEY_PARAM, WithDecryption=True)
        runpod_api_key = str(response["Parameter"]["Value"])
    return runpod_api_key


def output_prefix(key: str) -> str:
    """Map an inbound key to its output prefix (path-mirrored, extension dropped)."""
    rel = key[len(INPUT_PREFIX) :] if key.startswith(INPUT_PREFIX) else key
    stem = rel.rsplit(".", 1)[0]
    return f"{stem}/"


def object_exists(key: str) -> bool:
    try:
        s3.head_object(Bucket=OUTPUT_BUCKET, Key=key)
        return True
    except s3.exceptions.ClientError:
        return False


def run_table_task(src: str, dst: str) -> list[str]:
    resp = ecs.run_task(
        cluster=CLUSTER,
        taskDefinition=TASK_DEFINITION,
        launchType="FARGATE",
        count=1,
        networkConfiguration={
            "awsvpcConfiguration": {
                "subnets": SUBNETS,
                "securityGroups": SECURITY_GROUPS,
                "assignPublicIp": "ENABLED",
            }
        },
        overrides={
            "containerOverrides": [
                {
                    "name": CONTAINER_NAME,
                    "command": ["table", src, "-o", dst, "--review", "--llm-backend", "api"],
                }
            ]
        },
    )
    failures = resp.get("failures", [])
    if failures:
        log.error("RunTask failures for %s: %s", src, failures)
    return [t["taskArn"] for t in resp.get("tasks", [])]


def presign_put(key: str) -> str:
    return s3.generate_presigned_url(
        "put_object", Params={"Bucket": OUTPUT_BUCKET, "Key": key}, ExpiresIn=URL_EXPIRY_SECONDS
    )


def submit_parse_job(bucket: str, key: str, prefix: str, base: str) -> str:
    """POST the presigned-URL contract to the RunPod endpoint.

    log_url receives the worker's log every few seconds while the job runs;
    the same key is overwritten in place, so reading it at any moment gives
    the full log so far.
    """
    payload = {
        "input": {
            "base": base,
            "pdf_url": s3.generate_presigned_url(
                "get_object", Params={"Bucket": bucket, "Key": key}, ExpiresIn=URL_EXPIRY_SECONDS
            ),
            "artifact_urls": {
                "document": presign_put(f"{prefix}{base}.docling.json"),
                "confidence": presign_put(f"{prefix}{base}.confidence.json"),
                "cells": presign_put(f"{prefix}{base}.cells.json"),
            },
            "marker_url": presign_put(f"{prefix}{base}.parse.complete.json"),
            "log_url": presign_put(f"{prefix}{base}.parse.log"),
        }
    }
    request = urllib.request.Request(
        f"https://api.runpod.ai/v2/{RUNPOD_ENDPOINT_ID}/run",
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {get_runpod_api_key()}",
            "Content-Type": "application/json",
        },
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        job = json.load(response)
    log.info("runpod job %s (%s) for %s", job.get("id"), job.get("status"), key)
    return job["id"]


def schedule_deadline_check(prefix: str, base: str, job_id: str) -> None:
    """One-shot schedule that asks the join Lambda to verify the parse marker.

    ActionAfterCompletion=DELETE makes the schedule self-cleaning; the name
    embeds the job id so re-drops of the same document never collide.
    """
    fire_at = datetime.now(timezone.utc) + timedelta(minutes=DEADLINE_MINUTES)
    scheduler.create_schedule(
        Name=f"quber-parse-deadline-{job_id}"[:64],
        ScheduleExpression=f"at({fire_at.strftime('%Y-%m-%dT%H:%M:%S')})",
        FlexibleTimeWindow={"Mode": "OFF"},
        ActionAfterCompletion="DELETE",
        Target={
            "Arn": JOIN_LAMBDA_ARN,
            "RoleArn": SCHEDULER_ROLE_ARN,
            "Input": json.dumps({"deadline_check": {"prefix": prefix, "base": base, "job_id": job_id}}),
        },
    )


def handler(event: dict[str, Any], context: object) -> dict[str, list[str]]:
    started: dict[str, list[str]] = {"table": [], "parse": []}
    for record in event.get("Records", []):
        bucket = record["s3"]["bucket"]["name"]
        key = urllib.parse.unquote_plus(record["s3"]["object"]["key"])

        if not key.lower().endswith(".pdf"):
            log.info("skip non-pdf: %s", key)
            continue

        prefix = output_prefix(key)
        base = prefix.rstrip("/").rsplit("/", 1)[-1]
        src = f"s3://{bucket}/{key}"
        dst = f"s3://{OUTPUT_BUCKET}/{prefix}"

        # Per-engine dedup: results-presence, not marker-presence, so a
        # failed half (which leaves no results) re-runs on re-drop while a
        # completed half is never repeated.
        if object_exists(f"{prefix}{base}.tables.json"):
            log.info("skip table (results exist): %s", key)
        else:
            arns = run_table_task(src, dst)
            log.info("started table task(s) for %s: %s", key, arns)
            started["table"].extend(arns)

        if object_exists(f"{prefix}{base}.docling.json"):
            log.info("skip parse (results exist): %s", key)
        else:
            job_id = submit_parse_job(bucket, key, prefix, base)
            schedule_deadline_check(prefix, base, job_id)
            started["parse"].append(job_id)

    return started