import os
import zipfile
import boto3
import sys
import time
from datetime import datetime, timezone
import subprocess

session = boto3.Session(profile_name="staging")
glue_client = session.client("glue", region_name="us-east-2")
s3_client = session.client("s3", region_name="us-east-2")


def zip_code(source_files, output_filename):
    with zipfile.ZipFile(output_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
        for source_file in source_files:
            if os.path.isdir(source_file):
                for root, dirs, files in os.walk(source_file):
                    for file in files:
                        zipf.write(
                            os.path.join(root, file),
                            os.path.relpath(
                                os.path.join(root, file),
                                os.path.join(source_file, ".."),
                            ),
                        )
            else:
                zipf.write(source_file, os.path.basename(source_file))
    print(f"File {output_filename} zipped successfully")


def upload_to_s3(file_name, bucket, object_name=None):
    if object_name is None:
        object_name = file_name
    try:
        response = s3_client.upload_file(file_name, bucket, object_name)
        print(f"File {output_filename} uploaded successfully")
    except Exception as e:
        print(f"Error uploading file: {e}")
        return False
    return True


def create_glue_job(job_name, script_location, role):
    try:
        response = glue_client.create_job(
            Name=job_name,
            Role=role,
            Command={
                "Name": "glueetl",
                "ScriptLocation": script_location,
                "PythonVersion": "3",
            },
            DefaultArguments={
                "--job-language": "python",
                "--TempDir": "s3://testing-glue-table/temporary/",
                "--enable-metrics": "true",
                "--enable-spark-ui": "true",
                "--enable-glue-datacatalog": "true",
                "--spark-event-logs-path": "s3://testing-glue-table/sparkHistoryLogs/",
                "--enable-job-insights": "true",
                "--enable-observability-metrics": "true",
                "--enable-continuous-cloudwatch-log": "true",
                "--job-bookmark-option": "job-bookmark-disable",
                "--enable-auto-scaling": "true",
                "--extra-py-files": "s3://testing-glue-table/suno-glue.zip",
            },
            WorkerType="G.1X",
            NumberOfWorkers=2,
            Timeout=2880,
            GlueVersion="4.0",
            MaxRetries=0,
            Tags={"Project": "suno-glue"},
            Connections={
                "Connections": ["analystic-database-connection"],
            },
        )
        print(f"Created Glue job: {response['Name']}")
    except glue_client.exceptions.AlreadyExistsException:
        print(f"Glue job '{job_name}' already exists.")
    except Exception as e:
        print(f"Error creating Glue job: {e}")
        return False
    return True


def start_glue_job(job_name, job_arguments):
    try:
        response = glue_client.start_job_run(JobName=job_name, Arguments=job_arguments)
        print(f"Started Glue job with run ID: {response['JobRunId']}")
        return response["JobRunId"]
    except Exception as e:
        print(f"Error starting Glue job: {e}")
        return False


def glue_job_exists(job_name):
    try:
        glue_client.get_job(JobName=job_name)
        return True
    except glue_client.exceptions.EntityNotFoundException:
        return False
    except Exception as e:
        print(f"Error checking Glue job existence: {e}")
        return False


def print_with_timestamp(message: str):
    print(f"{time.strftime('%H:%M:%S')} {message}")


def poll_glue_job_status(job_name, run_id):
    total_time = 0  # Bail after 5 mins
    while total_time < 300:
        try:
            job = glue_client.get_job_run(JobName=job_name, RunId=run_id)

            job_run_state = job.get("JobRun", {}).get("JobRunState", None)
            match job_run_state:  # 'STARTING'|'RUNNING'|'STOPPING'|'STOPPED'|'SUCCEEDED'|'FAILED'|'TIMEOUT'|'ERROR'|'WAITING'|'EXPIRED'
                case "SUCCEEDED":
                    print_with_timestamp(
                        f"Job run succeeded at time {job.get('JobRun', {}).get('CompletedOn')}"
                    )
                    return True  # TODO: s3 get_object to get and print the dumped data
                case (
                    "STOPPING" | "STOPPED" | "FAILED" | "TIMEOUT" | "ERROR" | "EXPIRED"
                ):
                    print_with_timestamp(
                        f"Job run status {job_run_state}. Error message is {job.get('JobRun', {}).get('ErrorMessage')} - check AWS Glue console for more information."
                    )
                    return False
                case _:
                    print_with_timestamp(f"Job status is {job_run_state}")

            # If still STARTING | RUNNING | WAITING, sleep and poll
            time.sleep(20)
            total_time += 20

        except Exception as e:
            print_with_timestamp(f"Error polling for job status: {e}")
            return False

    if total_time >= 300:
        print_with_timestamp(
            "Job still running after 5 minutes. Check the AWS Console for more info."
        )
        return False


def download_and_print_object(key, skip_cleanup, dir=None):
    # Call S3 GetObject
    object = s3_client.get_object(Bucket=bucket_name, Key=key)

    # Download file from S3
    dir = dir or "tmp"
    file_name = os.path.basename(key)
    download_path = os.path.join(os.getcwd(), dir, file_name)
    with open(download_path, "wb") as f:
        f.write(object["Body"].read())

    # Show data
    subprocess.run(["parquet-tools", "show", download_path])

    # Maybe clean up tmp file
    if not skip_cleanup:
        os.remove(download_path)


def get_data_dump(table_name, skip_cleanup, p_date=None, p_hour=None):
    curr_time = datetime.now(timezone.utc)
    p_hour = (
        curr_time.hour - 1
    ) % 24  # Glue job dumps data from previous FULL hour (x:00-y:00)
    p_date = p_date or curr_time.strftime("%Y-%m-%d")

    try:
        print_with_timestamp(
            f"Getting results from s3 at output/{table_name}/pdate={p_date}/phour={p_hour}"
        )
        result = s3_client.list_objects_v2(
            Bucket=bucket_name,
            Prefix=f"output/{table_name}/pdate={p_date}/phour={p_hour}",
        )
        if result and result.get("Contents"):
            sorted_contents = sorted(
                result["Contents"], key=lambda object: object["LastModified"]
            )
            last_modified_key = sorted_contents[-1]["Key"]
            print_with_timestamp(
                f"Last modified object is {last_modified_key} - Starting download."
            )
            download_and_print_object(last_modified_key, skip_cleanup=skip_cleanup)

        else:
            print_with_timestamp(
                f"Did not find any data in s3. Are there entries in your table from the last full hour ({p_hour} - {curr_time.hour} UTC)?"
            )
            return False
    except Exception as e:
        print_with_timestamp(f"Error getting data from S3 {e}")
        return False


if __name__ == "__main__":
    skip_cleanup = False
    if len(sys.argv) != 3:
        if len(sys.argv) == 4 and sys.argv[3] == "--skip-cleanup":
            skip_cleanup = True
        else:
            print("Usage: uv run test_upload.py <glue_job_name> <table_name>")
            sys.exit(1)

    glue_job_name = sys.argv[1]
    table_name = sys.argv[2]

    # Define source directory and output file
    source_files = ["jobs"]
    output_filename = "suno-glue.zip"

    # Zip the code
    zip_code(source_files, output_filename)

    # Define S3 bucket name and AWS profile name
    bucket_name = "testing-glue-table"

    # Upload the ZIP file to S3
    upload_result = upload_to_s3(output_filename, bucket_name)
    upload_result = upload_to_s3("main.py", bucket_name)

    if upload_result:
        os.remove(output_filename)
        print(f"Deleted local file: {output_filename}")

        # Define script location in S3
        script_location = f"s3://{bucket_name}/main.py"

        role = "testing-rds-glue-role"

        # Check if the Glue job exists
        if not glue_job_exists(glue_job_name):
            # Create the Glue job if it doesn't exist
            create_glue_job(glue_job_name, script_location, role)

        # Define Glue job arguments
        job_arguments = {
            "--JOB_NAME": glue_job_name,
            "--job_class": f"jobs.{glue_job_name}.GlueJob",
            "--env": "staging",
        }

        # Start the Glue job
        run_id = start_glue_job(glue_job_name, job_arguments)

        # Poll for results
        if not run_id:
            print("Did not successfully start glue job")
        else:
            success = poll_glue_job_status(glue_job_name, run_id)
            if success:
                get_data_dump(
                    table_name,
                    skip_cleanup,
                )
    else:
        print(f"Failed to upload file: {output_filename}")