from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from os import environ

from src import hive, aws

BUCKET_NAME = "prod-hive-regression-tracking"
S3_PREFIX_ASSETS = environ.get("S3_PREFIX_ASSETS", "assets/")
S3_PREFIX_RESULTS = "results/"
S3_PREFIX_RESULTS_RAW = "raw/"
S3_PREFIX_RESULTS_FLATTENED = "flattened/"
# Add any specific asset_final_ids to skip here, e.g. assets too large for the current project
ASSET_IDS_TO_SKIP = set([])
# 'asset_final_id' (default): standard regression assets named '<asset_final_id>.flac'.
# 'filename': oneoff sets mixing flac/mp3/wav, tracked as '{name}_{ext}'
# (e.g. '10000.flac' -> '10000_flac') to match the HIVE_INVESTIGATION_* tables.
# The mode is also the name of the tracking field added to each result row.
FILE_ID_MODE = environ.get("FILE_ID_MODE", "asset_final_id")
if FILE_ID_MODE not in ("asset_final_id", "filename"):
    raise ValueError(f"FILE_ID_MODE must be 'asset_final_id' or 'filename', got {FILE_ID_MODE!r}")

AUDIO_EXTENSIONS = ("flac", "mp3", "wav")


def get_file_id_from_path(path: str) -> str:
    """Extract the tracking id from an asset path, per FILE_ID_MODE."""
    if not isinstance(path, str) or not path:
        raise ValueError("path must be a non-empty string")

    filename = path.rsplit("/", 1)[-1]
    name, _, ext = filename.rpartition(".")
    ext = ext.lower()
    if not name:
        raise ValueError(f"missing name in {filename!r}")

    if FILE_ID_MODE == "asset_final_id":
        if ext != "flac":
            raise ValueError(f"expected a .flac file, got {filename!r}")
        return name

    if ext not in AUDIO_EXTENSIONS:
        raise ValueError(f"unsupported audio file: {filename!r}")

    return f"{name}_{ext}"


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 file ids found in the S3 results folder.
    """
    # Key completion on the flattened CSV: it is written after the raw JSON, so
    # its existence implies the asset fully processed. Keying on raw/ would
    # permanently skip assets that failed between the two writes.
    prefix = f"{S3_PREFIX_RESULTS}{run_id}/{S3_PREFIX_RESULTS_FLATTENED}"
    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: .../flattened/{file_id}_hive_response_flattened.csv
        filename = key.split("/")[-1]
        if filename.endswith("_hive_response_flattened.csv"):
            file_id = filename.removesuffix("_hive_response_flattened.csv")
            if file_id:
                processed_ids.add(file_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")


def process_asset(file_key: str, run_id: str) -> None:
    """Process a single asset: call HIVE API and write results to S3."""
    file_id = get_file_id_from_path(file_key)
    print(f"Processing {file_id}...")

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

    # Tag the response with its tracking id for reporting
    hive_response[FILE_ID_MODE] = file_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}{file_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}{file_id}_hive_response_flattened.csv"
    )


if __name__ == "__main__":
    if not environ.get("HIVE_CREDENTIALS"):
        raise RuntimeError("HIVE_CREDENTIALS is not set")

    max_workers = int(environ.get("MAX_WORKERS", "5"))

    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)}")

    # Pre-warm S3 client on the main thread before spawning workers
    aws.get_s3_client()

    # Filter to only unprocessed, valid assets upfront
    assets_to_process = []

    for file_key in files:
        try:
            file_id = get_file_id_from_path(file_key)
        except ValueError:
            # Skip files that don't match expected naming (e.g. folder placeholders)
            continue
        if file_id not in processed_ids and file_id not in ASSET_IDS_TO_SKIP:
            assets_to_process.append(file_key)

    print(f"Assets to process: {len(assets_to_process)} (workers: {max_workers})")

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(process_asset, fk, run_id): fk for fk in assets_to_process}
        for future in as_completed(futures):
            file_key = futures[future]
            try:
                future.result()
            except Exception as e:
                print(f"ERROR processing {file_key}: {e}")