import os, io, json, base64, time, zipfile, datetime
import requests, boto3
from google.cloud import storage
from google.auth.transport.requests import Request
import google.auth
from google.auth.transport.requests import AuthorizedSession
import base64
import functions_framework
import pathlib
import textwrap
import json
from datetime import datetime, timedelta, timezone



# Config
PROJECT_ID = "pidtoo-fhir"
LOCATION = "us-east4"
DATASET_ID = "isc"
FHIR_STORE_ID = "fhir-omop"
GCS_EXPORT_BUCKET = "fhir-export-bucket"
AWS_BUCKET = "intersystems-fhir2omop"
AWS_REGION = "us-east-2"
# Trigger FHIR export
def trigger_incremental_export(export_time_iso):
    client = storage.Client()
    bucket = client.bucket("fhir-export-bucket")

    blobs = bucket.list_blobs()
    for blob in blobs:
        print(f"Deleting: {blob.name}")
        blob.delete()
    
    credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
    authed_session = AuthorizedSession(credentials)

    export_uri = f"gs://{GCS_EXPORT_BUCKET}/fhir-export-{int(time.time())}/"
    export_uri = f"gs://{GCS_EXPORT_BUCKET}/"
    url = (
        f"https://healthcare.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/"
        f"datasets/{DATASET_ID}/fhirStores/{FHIR_STORE_ID}:export"
    )

    body = {
        "gcsDestination": {"uriPrefix": export_uri},
        "since": export_time_iso
    }

    response = authed_session.post(url, json=body)
    print(f"Export response: {response.status_code} - {response.text}")
    return export_uri if response.ok else None
# Poll GCS for export results
def wait_for_ndjson_files(export_uri_prefix):
    client = storage.Client()
    bucket_name = export_uri_prefix.split("/")[2]
    prefix = "/".join(export_uri_prefix.split("/")[3:])
    print(bucket_name)
    print(prefix)

    bucket = client.bucket(bucket_name)
    for _ in range(20):  # Wait up to ~5 mins
        blobs = list(bucket.list_blobs(prefix=prefix))
        if any(blob.name.endswith("Organization") for blob in blobs):
            return [blob for blob in blobs if blob.name.endswith("Organization")]
        time.sleep(5)
    raise TimeoutError("Export files did not appear in GCS within timeout window")

# Zip .ndjsons into flat ZIP file
def create_zip_from_blobs(blobs, zip_path):
    client = storage.Client()
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for blob in blobs:
            data = blob.download_as_bytes()
            fname = os.path.basename(blob.name)
            zipf.writestr(fname + ".ndjson", data)

# Upload ZIP to AWS S3
def upload_to_s3(zip_path, s3_key):
    s3 = boto3.client('s3', region_name=AWS_REGION)
    s3.upload_file(zip_path, AWS_BUCKET, "from_gcp_to_omop" + s3_key)
    print(f"Uploaded {zip_path} to s3://{AWS_BUCKET}/from_gcp_to_omop/{s3_key}")


#@functions_framework.cloud_event
#def mit_grandhack(cloud_event):
    # Print out the data from Pub/Sub, to prove that it worked
#    print(base64.b64decode(cloud_event.data["message"]["data"]))
#    question = base64.b64decode(cloud_event.data["message"]["data"]).decode()
@functions_framework.cloud_event
def receive_pubsub(cloud_event):
    #envelope = request.get_json()
    print(cloud_event)
    data = base64.b64decode(cloud_event.data["message"]["data"]).decode()
    data = cloud_event.data
    print(data)
    print(type(data))
    if not data:
        return "No data", 400
    #payload = data # json.loads(data)
    #method = payload.get("protoPayload", {}).get("methodName", "")
    method = data['message']['attributes']['action']
    #resource_name = payload.get("protoPayload", {}).get("resourceName", "")
    resource_name = data['message']['attributes']['resourceType']
    #timestamp = payload.get("timestamp", "")
    timestamp = data['message']['publishTime']
    # Input datetime string
    # Parse the string to a datetime object
    dt = datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%S.%fZ").replace(tzinfo=timezone.utc)

    # Subtract 5 minutes
    five_minutes_ago = dt - timedelta(minutes=5)

    # Convert back to ISO 8601 string format with 'Z'
    timestamp = five_minutes_ago.isoformat().replace('+00:00', 'Z')

    print(method)
    print(resource_name)
    print(timestamp)

    if "CreateResource" in method and "Organization" in resource_name:
        print(f"New Organization detected at {timestamp}")
        export_uri = trigger_incremental_export(timestamp)
        if not export_uri:
            return "Export failed", 500
        blobs = wait_for_ndjson_files(export_uri)
        zip_file_path = "/tmp/fhir_export.zip"
        create_zip_from_blobs(blobs, zip_file_path)
        s3_key = f"/export-{int(time.time())}.zip"
        upload_to_s3(zip_file_path, s3_key)
        return "Exported and uploaded", 200
    return "No relevant event", 204

{'attributes': {'specversion': '1.0', 'id': '13999883936448345', 'source': '//pubsub.googleapis.com/projects/pidtoo-fhir/topics/fhir-omop-topic', 'type': 'google.cloud.pubsub.topic.v1.messagePublished', 'datacontenttype': 'application/json', 'time': '2025-05-13T20:13:20.339Z'}, 'data': {'message': {'attributes': {'action': 'CreateResource', 'lastUpdatedTime': 'Tue, 13 May 2025 20:13:20 UTC', 'payloadType': 'FullResource', 'resourceType': 'Organization', 'storeName': 'projects/pidtoo-fhir/locations/us-east4/datasets/isc/fhirStores/fhir-omop', 'versionId': 'MTc0NzE2NzIwMDEwNzczODAwMA'}, 'data': 'ewogICJhZGRyZXNzIjogWwogICAgewogICAgICAiY2l0eSI6IC', 'messageId': '13999883936448345', 'message_id': '13999883936448345', 'publishTime': '2025-05-13T20:13:20.339Z', 'publish_time': '2025-05-13T20:13:20.339Z'}, 'subscription': 'projects/pidtoo-fhir/subscriptions/eventarc-us-east4-fhir2omop-trigger-sub-855'}}