# Copyright Modal Labs 2024

import json
import os
from pathlib import Path
from typing import List, Dict, Any

import modal

# Create the Modal image with required dependencies
image = modal.Image.debian_slim(python_version="3.12").pip_install("requests==2.31.0")

# Create the Modal app
app = modal.App("suno-audio-downloader", image=image)

# Create a Modal Volume to store the downloaded audio files
audio_volume = modal.Volume.from_name(
    "suno-hoot-data",
    create_if_missing=True,
    version=2,  # use VolumeFS2 for much higher write concurrency
)
model_volume = modal.Volume.from_name("suno-hoot-training")
VOLUME_PATH = Path("/audio_data")

with image.imports():
    import requests


@app.function(
    volumes={VOLUME_PATH: audio_volume},
    timeout=3600,  # 1 hour timeout for large downloads
    cpu=2,
    memory=4096,
)
def download_audio_file_from_item(item: Dict[str, Any]) -> str:
    """Download a single audio file from a JSON item."""

    # Look for URL in various possible fields
    raw_url = (
        item.get("url", "") or item.get("audio_url", "") or item.get("s3_path", "")
    )
    audio_id = item.get("id", "")

    if not raw_url or not audio_id:
        return f"Invalid data for item: {item}"

    # Only accept S3 paths in the specific format
    if raw_url.startswith("s3://suno-data-uploads/studio/uploads/"):
        # Extract filename from S3 path
        filename = raw_url.split("/")[-1]
        url = f"https://cdn1.suno.ai/{filename}"
    else:
        return f"Invalid S3 path format: {raw_url}. Expected format: s3://suno-data-uploads/studio/uploads/filename.mp3"

    # Create local file path
    local_filename = f"{audio_id}.mp3"
    local_file_path = VOLUME_PATH / local_filename
    local_file_path.parent.mkdir(parents=True, exist_ok=True)

    # Download the file
    # print(f"Downloading {url} to {local_file_path}")

    try:
        response = requests.get(url, stream=True, timeout=300)
        response.raise_for_status()

        with open(local_file_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        # Commit the volume to persist the file
        audio_volume.commit()

        return f"Successfully downloaded {local_filename}"

    except requests.RequestException as e:
        return f"Error downloading {url}: {str(e)}"
    except Exception as e:
        return f"Error saving {local_filename}: {str(e)}"


@app.function(
    volumes={VOLUME_PATH: audio_volume},
    timeout=3600,  # 1 hour timeout for large downloads
    cpu=2,
    memory=4096,
)
def download_audio_file(url: str, local_filename: str) -> str:
    """Download a single audio file from URL to the Modal volume."""

    # Create local file path
    local_file_path = VOLUME_PATH / local_filename
    local_file_path.parent.mkdir(parents=True, exist_ok=True)

    # Download the file
    print(f"Downloading {url} to {local_file_path}")

    try:
        response = requests.get(url, stream=True, timeout=300)
        response.raise_for_status()

        with open(local_file_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        # Commit the volume to persist the file
        audio_volume.commit()

        return f"Successfully downloaded {local_filename}"

    except requests.RequestException as e:
        return f"Error downloading {url}: {str(e)}"
    except Exception as e:
        return f"Error saving {local_filename}: {str(e)}"


def process_files_in_batches(
    valid_data: List[Dict[str, Any]], batch_size: int = 1000
) -> List[str]:
    """Process files in batches using Modal's .map API."""

    all_results = []

    # Process files in batches of 1000
    for i in range(0, len(valid_data), batch_size):
        batch = valid_data[i : i + batch_size]
        print(
            f"Processing batch {i // batch_size + 1}/{(len(valid_data) + batch_size - 1) // batch_size} ({len(batch)} files)"
        )

        # Use Modal's .map API to process the batch in parallel
        batch_results = list(download_audio_file_from_item.map(batch))
        all_results.extend(batch_results)

        print(f"Completed batch {i // batch_size + 1}")

    return all_results


@app.function(
    volumes={VOLUME_PATH: audio_volume},
    timeout=600,  # 10 minute timeout
)
def list_downloaded_files() -> List[str]:
    """List all files currently stored in the Modal volume."""

    files = []
    for file_path in VOLUME_PATH.rglob("*"):
        if file_path.is_file():
            files.append(str(file_path.relative_to(VOLUME_PATH)))

    return sorted(files)


@app.function(
    volumes={VOLUME_PATH: audio_volume},
)
def get_volume_stats() -> Dict[str, Any]:
    """Get statistics about the downloaded files."""

    audio_volume.reload()

    total_files = 0
    total_size = 0

    for file_path in VOLUME_PATH.rglob("*.mp3"):
        if file_path.is_file():
            total_files += 1
            total_size += file_path.stat().st_size

    return {
        "total_files": total_files,
        "total_size_bytes": total_size,
        "total_size_mb": round(total_size / (1024 * 1024), 2),
        "total_size_gb": round(total_size / (1024 * 1024 * 1024), 2),
    }


@app.function(
    volumes={VOLUME_PATH: audio_volume, "/models": model_volume},
    timeout=600,  # 10 minute timeout
    cpu=2,
    memory=4096,
)
def create_cleaned_json() -> Dict[str, Any]:
    """Create a cleaned JSON file containing only entries for which there is an actual file in the volume.

    The CDN can throw occational 403s, so we should avoid trying to load files which are not available.

    In production, use an S3 mount to ensure you get all files.
    """

    # Load the original JSON file
    with open("/models/suno_hoot_20250617_subset100k.json", "r") as f:
        original_data = json.load(f)

    print(f"Original JSON contains {len(original_data)} entries")

    # Get all existing audio files in the volume
    existing_files = set()
    for file_path in VOLUME_PATH.rglob("*.mp3"):
        if file_path.is_file():
            # Extract the ID from the filename (remove .mp3 extension)
            file_id = file_path.stem
            existing_files.add(file_id)

    print(f"Found {len(existing_files)} existing audio files in volume")

    # Filter original data to only include entries with existing files
    cleaned_data = []
    for item in original_data:
        item_id = item.get("id", "")
        if item_id in existing_files:
            cleaned_data.append(item)

    print(f"Cleaned JSON will contain {len(cleaned_data)} entries")

    # Write the cleaned JSON file
    with open("/models/suno_hoot_cleaned.json", "w") as f:
        json.dump(cleaned_data, f, indent=2)

    # Commit the model volume to persist the cleaned JSON
    model_volume.commit()

    return {
        "original_entries": len(original_data),
        "existing_files": len(existing_files),
        "cleaned_entries": len(cleaned_data),
        "success_rate": round(len(cleaned_data) / len(original_data) * 100, 2)
        if original_data
        else 0,
    }


@app.function(
    volumes={VOLUME_PATH: audio_volume, "/models": model_volume},
    timeout=14400,  # 4 hour timeout for the full orchestration
    cpu=4,
    memory=8192,
)
def orchestrate():
    """Main orchestration function that runs the entire download process."""

    with open("/models/suno_hoot_20250617_subset100k.json", "r") as f:
        data = json.load(f)

    print(f"Received {len(data)} items for processing")

    # Filter data to only include items with URLs
    valid_data = [
        item
        for item in data
        if item.get("url") or item.get("audio_url") or item.get("s3_path")
    ]
    print(f"Found {len(valid_data)} items with valid URLs")

    if not valid_data:
        print("No valid items found to download")
        return

    # Download files in batches using Modal's .map API
    print("Starting download process...")
    results = process_files_in_batches(valid_data)

    # Print results
    for result in results:
        print(result)

    # Get volume statistics
    stats = get_volume_stats.remote()
    print("\nVolume Statistics:")
    print(f"Total files: {stats['total_files']}")
    print(f"Total size: {stats['total_size_mb']} MB ({stats['total_size_gb']} GB)")

    print("\nDownload process completed!")


@app.local_entrypoint()
def main():
    """Local entrypoint that loads JSON data and runs orchestration in detached mode."""
    print("Starting orchestration...")

    # Run the orchestration function remotely (detached) with the loaded data
    orchestrate.remote()


@app.local_entrypoint()
def clean_json():
    """Local entrypoint to create a cleaned JSON file with only entries that have downloaded files."""
    print("Creating cleaned JSON file...")

    # Run the cleaning function remotely
    result = create_cleaned_json.remote()

    print(f"Cleaning completed:")
    print(f"  Original entries: {result['original_entries']}")
    print(f"  Existing files: {result['existing_files']}")
    print(f"  Cleaned entries: {result['cleaned_entries']}")
    print(f"  Success rate: {result['success_rate']}%")
    print(f"  Cleaned JSON saved to: /models/suno_hoot_cleaned.json")


if __name__ == "__main__":
    main()