"""
Tutorial: Submitting and Monitoring Jobs with the emtelligent Extraction API

This script demonstrates the complete workflow for submitting extraction jobs, monitoring job progress, and collecting output files to your S3 bucket.

Installation:
    pip install boto3 requests

Prerequisites:
    Before running this script, ensure you have the following:

    AWS Resources:
        - Input S3 bucket with appropriate permissions
        - Output S3 bucket with appropriate permissions
        - AWS access key and secret access key
        - STS Role ARNs for both input and output bucket access

    emtelligent Resources:
        - API key (provided by emtelligent)
        - Pipeline name (provided by emtelligent)

    Sample Data:
        - Sample PDF files for testing
        - Manifest CSV file listing the documents to process

Usage:
    1. Update the Configuration section with your credentials and bucket information
    2. Run the script: python tutorial.py

    The script will:
        - Upload files to your input bucket
        - Submit a job to the extraction API
        - Monitor job status until completion
        - Collect outputs to your output bucket
        - List the generated output files

"""

import time
from pathlib import Path

import boto3
import requests

# ── Configuration ─────────────────────────────────────────────────────────────

# For file uploads
# Fill in emtelligent-provided input bucket name
input_bucket_name = "your-input-bucket-name"

# Specify the directory where you extracted tutorial.zip (containing sample PDFs and manifest.csv)
file_path = "~/tutorial"
manifest_path = f"{file_path}/manifest.csv"

s3_prefix = "tutorial"
manifest_s3_key = f"{s3_prefix}/manifest.csv"

# For polling job status
FAILED_STATUSES = {"transfer_failed", "failed", "cancelled"}
# These settings assume you are submitting a small job for the tutorial
# Choose appropriate polling intervals for larger jobs
MAX_WAIT_TIME_S = 3600  # 1 hour max wait time
INITIAL_POLL_INTERVAL_S = 5
MAX_POLL_INTERVAL_S = 60
BACKOFF_FACTOR = 1.5

# For calling the API
BASE_URL = "https://extraction-api.emtelligent.com/api/v1"

# Replace with your emtelligent issued API Key and the provided pipeline name
API_KEY = "<api-key>"
pipeline_name = "<pipeline_name>"

# For this tutorial, let's extract lab results and vital signs
pipeline_params = {"extraction_type": ["labs", "vitals"]}

# Replace with the role arn to use to provide access to the input bucket
input_bkt_role_arn = "<role-arn>"
# Replace with the role arn to use to provide access to the output bucket
output_bkt_role_arn = "<role-arn>"

# For final output collection
# Fill in emtelligent-provided output bucket name
output_bucket_name = "your-output-bucket-name"
output_prefix = "tutorial_output"

# Fill in the credentials to access your s3 bucket
credentials = {
    'aws_access_key': "access_key",
    'aws_secret_key': "secret_key",
    'aws_region': "region"
}

# ── End of Configuration ──────────────────────────────────────────────────────


def get_pipeline_params(pipeline_name: str) -> dict:
    response = requests.get(
        f"{BASE_URL}/pipelines/{pipeline_name}/pipeline_params",
        headers={
            "accept": "application/json",
            "X-API-Key": API_KEY,
        },
    )
    response.raise_for_status()
    return response.json()


def submit_job(
    pipeline_name: str,
    manifest_key: str,
    aws_region: str,
    user_credentials: dict,
    pipeline_parameters: dict,
) -> dict:
    response = requests.post(
        f"{BASE_URL}/jobs",
        headers={
            "accept": "application/json",
            "X-API-Key": API_KEY,
            "Content-Type": "application/json",
        },
        json={
            "pipeline_name": pipeline_name,
            "manifest_key": manifest_key,
            "aws_region": aws_region,
            "user_credentials": user_credentials,
            "pipeline_parameters": pipeline_parameters,
        },
    )
    response.raise_for_status()
    return response.json()


def get_job_status(job_id: str) -> dict:
    response = requests.get(
        f"{BASE_URL}/jobs/{job_id}/status",
        headers={
            "accept": "application/json",
            "X-API-Key": API_KEY,
        },
    )
    response.raise_for_status()
    return response.json()


def collect_job(
    job_id: str,
    aws_region: str,
    output_prefix: str,
    user_credentials: dict,
) -> dict:
    response = requests.post(
        f"{BASE_URL}/jobs/{job_id}/collect",
        headers={
            "accept": "application/json",
            "X-API-Key": API_KEY,
            "Content-Type": "application/json",
        },
        json={
            "aws_region": aws_region,
            "output_prefix": output_prefix,
            "user_credentials": user_credentials,
        },
    )
    response.raise_for_status()
    return response.json()


def main():
    # ── Step 1: Upload Files to S3 ────────────────────────────────────────────

    print("Step 1: Uploading files to input bucket...")
    s3_client = boto3.client(
        "s3",
        aws_access_key_id=credentials["aws_access_key"],
        aws_secret_access_key=credentials["aws_secret_key"],
        region_name=credentials["aws_region"],
    )

    # Upload manifest
    s3_client.upload_file(str(Path(manifest_path).expanduser()),
                          input_bucket_name, manifest_s3_key)

    # Upload all PDF files from the directory
    # Ensure the document values in your manifest match the exact S3 keys where your PDF files are stored
    for pdf_file in Path(file_path).expanduser().glob("*.pdf"):
        s3_key = f"{s3_prefix}/{pdf_file.name}"
        s3_client.upload_file(str(pdf_file), input_bucket_name, s3_key)

    # ── Step 2: Generate STS Token for Input Bucket Access ───────────────────

    print("Step 2: Generating STS Token for customer bucket access")
    sts_client = boto3.client(
        "sts",
        aws_access_key_id=credentials["aws_access_key"],
        aws_secret_access_key=credentials["aws_secret_key"],
        region_name=credentials["aws_region"],
    )

    # For testing purposes we set the expiry to 3600s but we advise to increase this for bigger jobs (more PDFs to process)
    expiry_s = 3600  # 60min

    # Assume the role
    response = sts_client.assume_role(
        RoleArn=input_bkt_role_arn,
        RoleSessionName="tutorial-temp-session",
        DurationSeconds=expiry_s,
    )

    sts_creds = response["Credentials"]

    # ── Step 3: Identify Pipeline Parameters ──────────────────────────────────

    print(
        "Step 3: Displaying pipeline parameters")
    params = get_pipeline_params(pipeline_name)
    print(params)

    # ── Step 4: Submit a Job ──────────────────────────────────────────────────

    print(
        "Step 4: Creating job via Extraction API using aws temp credentials...")
    result = submit_job(
        pipeline_name=pipeline_name,
        manifest_key="tutorial/manifest.csv",
        aws_region=credentials["aws_region"],
        user_credentials={
            "aws_access_key_id": sts_creds["AccessKeyId"],
            "aws_secret_access_key": sts_creds["SecretAccessKey"],
            "expiry": str(sts_creds["Expiration"]),
            "session_token": sts_creds["SessionToken"],
        },
        pipeline_parameters=pipeline_params,
    )
    print(result)

    # save job ID for future steps
    job_id = result.get("job_id")

    # ── Step 5: Monitor Job Status ────────────────────────────────────────────

    print(
        f"Step 5: Job is running... Polling for 'complete' status")
    start_time = time.time()
    poll_interval = INITIAL_POLL_INTERVAL_S

    while True:
        status_response = get_job_status(job_id)
        status = status_response.get("job_status")
        print(f"Job {job_id} status: {status}")

        if status == "completed":
            break
        if status in FAILED_STATUSES or status == "paused":
            raise RuntimeError(f"Job {job_id} ended with status: {status}")

        # Check if we've exceeded max wait time
        if time.time() - start_time > MAX_WAIT_TIME_S:
            raise TimeoutError(
                f"Job {job_id} did not complete within {MAX_WAIT_TIME_S}s")

        # Sleep with exponential backoff
        time.sleep(poll_interval)
        poll_interval = min(poll_interval * BACKOFF_FACTOR,
                            MAX_POLL_INTERVAL_S)

    # ── Step 6: Collect Job Outputs ───────────────────────────────────────────

    # generate STS credentials for output bucket
    response = sts_client.assume_role(
        RoleArn=output_bkt_role_arn,
        RoleSessionName="tutorial-temp-session",
        DurationSeconds=expiry_s,
    )

    sts_creds = response["Credentials"]

    print(
        "Step 6: Calling Extraction API /collect endpoint with temp credentials")
    # Use STS credentials from assume_role (see the example for the input bucket)
    result = collect_job(
        job_id=job_id,
        aws_region=credentials["aws_region"],
        output_prefix=output_prefix,
        user_credentials={
            "aws_access_key_id": sts_creds["AccessKeyId"],
            "aws_secret_access_key": sts_creds["SecretAccessKey"],
            "expiry": str(sts_creds["Expiration"]),
            "session_token": sts_creds["SessionToken"],
        },
    )

    print("Files transferring.. Polling for 'transfer_completed' status")

    start_time = time.time()
    poll_interval = INITIAL_POLL_INTERVAL_S

    while True:
        status_response = get_job_status(job_id)
        status = status_response.get("job_status")
        print(f"Job {job_id} status: {status}")
        if status == "transfer_completed":
            break
        if status in FAILED_STATUSES or status == "paused":
            raise RuntimeError(
                f"Job {job_id} transfer ended with status: {status}")

        # Check if we've exceeded max wait time
        if time.time() - start_time > MAX_WAIT_TIME_S:
            raise TimeoutError(
                f"Job {job_id} transfer did not complete within {MAX_WAIT_TIME_S}s")

        # Sleep with exponential backoff
        time.sleep(poll_interval)
        poll_interval = min(poll_interval * BACKOFF_FACTOR,
                            MAX_POLL_INTERVAL_S)

    # ── Step 7: List Output Files ─────────────────────────────────────────────

    print("Step 7: Listing output files")
    response = s3_client.list_objects_v2(
        Bucket=output_bucket_name, Prefix=output_prefix)
    for obj in response.get("Contents", []):
        print(obj["Key"])


if __name__ == "__main__":
    main()