import boto3
import os
import hashlib
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")
AWS_REGION = os.getenv("AWS_REGION")

# Initialize boto3 S3 client
s3 = boto3.client(
    "s3",
    aws_access_key_id=AWS_ACCESS_KEY,
    aws_secret_access_key=AWS_SECRET_KEY,
    region_name=AWS_REGION
)

def list_pdfs_in_prefix(bucket, prefix):
    """
    List all PDF files under a given S3 prefix (folder).
    Handles pagination if more than 1000 files exist.
    """
    pdf_keys = []
    continuation_token = None

    while True:
        if continuation_token:
            response = s3.list_objects_v2(
                Bucket=bucket,
                Prefix=prefix,
                ContinuationToken=continuation_token
            )
        else:
            response = s3.list_objects_v2(Bucket=bucket, Prefix=prefix)

        if "Contents" in response:
            for obj in response["Contents"]:
                key = obj["Key"]
                if key.lower().endswith(".pdf"):
                    pdf_keys.append(key)

        # Check if there are more files
        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    return pdf_keys


def compute_md5_of_file(path: str, block_size: int = 2**20) -> str:
    """
    Compute and return the lowercase hex MD5 checksum of a file.
    """
    h = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(block_size), b""):
            h.update(chunk)
    return h.hexdigest()


def download_pdf_from_s3(bucket, key, destination_path):
    """
    Download a single PDF from S3 to the given local path.
    Returns the MD5 checksum (hex lowercase) of the file.
    """
    try:
        s3.download_file(bucket, key, destination_path)
        md5 = compute_md5_of_file(destination_path)
        print(f" Downloaded: {key} → {destination_path} (md5={md5})")
        return md5
    except Exception as e:
        print(f" Error downloading {key}: {e}")
        raise
