from datetime import datetime
from os import environ

from src import hive, aws

BUCKET_NAME = "prod-hive-regression-tracking"
S3_PREFIX_ASSETS = "assets/"
S3_PREFIX_RESULTS = "results/"
S3_PREFIX_RESULTS_RAW = "raw/"
S3_PREFIX_RESULTS_FLATTENED = "flattened/"


def get_name_from_flac_path(path: str) -> str:
    """Extract {name} from a path ending with '{name}.flac'."""
    if not isinstance(path, str) or not path:
        raise ValueError("path must be a non-empty string")

    filename = path.rsplit("/", 1)[-1]
    if not filename.lower().endswith(".flac"):
        raise ValueError("path must end with '.flac'")

    name = filename[:-5]
    if not name:
        raise ValueError("missing name before '.flac'")

    return name


def get_processed_asset_ids(bucket: str, run_id: str) -> set[str]:
    """
    List all assets that have already been processed for the given run_id.
    Returns a set of asset_final_ids found in the S3 results folder.
    """
    prefix = f"{S3_PREFIX_RESULTS}{run_id}/{S3_PREFIX_RESULTS_RAW}"
    print(f"Checking for existing results in s3://{bucket}/{prefix}...")

    existing_files = aws.list_s3_objects(bucket, prefix)
    processed_ids = set()

    for key in existing_files:
        # Expected format: .../raw/{asset_final_id}_hive_response.json
        filename = key.split("/")[-1]
        if filename.endswith("_hive_response.json"):
            asset_id = filename[:-19] # Remove "_hive_response.json"
            if asset_id:
                processed_ids.add(asset_id)

    print(f"Found {len(processed_ids)} already processed assets.")
    return processed_ids


def get_run_id() -> str:
    """Generate a unique run ID."""
    provided_run_id = environ.get("RUN_ID") or None
    if provided_run_id:
        return provided_run_id
    return datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")


if __name__ == "__main__":
    run_id = get_run_id()
    print(f"Using run ID: {run_id}")

    # helper set to check for existing work
    processed_ids = get_processed_asset_ids(BUCKET_NAME, run_id)

    # list the files in the S3 bucket under the assets directory
    files = aws.list_s3_objects(BUCKET_NAME, S3_PREFIX_ASSETS)
    print(f"Total files found in assets: {len(files)}")

    for file_key in files:
        try:
            asset_final_id = get_name_from_flac_path(file_key)
        except ValueError:
            # Skip files that don't match expected naming (e.g. folder placeholders)
            continue

        if asset_final_id in processed_ids:
            # verify we're not double processing
            # print(f"Skipping {asset_final_id} (already processed)")
            continue

        print(f"Processing {asset_final_id}...")

        # Get HIVE response for the asset
        presigned_url = aws.create_presigned_url(BUCKET_NAME, file_key)
        hive_response = hive.run_task(presigned_url)

        # Append asset final ID to hive response for tracking
        hive_response["asset_final_id"] = asset_final_id

        # Save raw response to S3 under results/<run_id>/raw/
        aws.write_json_to_s3(
            hive_response,
            BUCKET_NAME,
            f"{S3_PREFIX_RESULTS}{run_id}/{S3_PREFIX_RESULTS_RAW}{asset_final_id}_hive_response.json"
        )

        # Save flattened response to S3 under results/<run_id>/flattened/
        sorted_headers, flattened_rows = hive.flatten_hive_response(hive_response)
        aws.write_csv_to_s3(
            sorted_headers,
            [list(row.get(header, "") for header in sorted_headers) for row in flattened_rows],
            BUCKET_NAME,
            f"{S3_PREFIX_RESULTS}{run_id}/{S3_PREFIX_RESULTS_FLATTENED}{asset_final_id}_hive_response_flattened.csv"
        )