import requests
import os
from typing import List, Dict, Optional


def download_audio(url: str, filepath: str) -> bool:
    """
    Download audio file from URL to specified filepath
    """
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()

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

        print(f"Downloaded: {filepath}")
        return True
    except requests.exceptions.RequestException as e:
        print(f"Failed to download {url}: {e}")
        return False


def organize_separation_results(run_ids: List[int], output_dir: str) -> Optional[Dict]:
    """
    Retrieve and organize multiple audio separation results
    Args:
        run_ids: List of run IDs to fetch results for
        output_dir: Directory to save organized results
    Returns:
        Dictionary containing organized results or None if request fails
    """
    try:
        # Make API request
        response = requests.post(
            "https://client.camb.ai/apis/audio-separation-results",
            headers={"x-api-key": "YOUR_API_KEY_HERE", "Content-Type": "application/json"},
            json={"run_ids": run_ids},
        )
        response.raise_for_status()
        results = response.json()

        # Create main output directory
        os.makedirs(output_dir, exist_ok=True)

        organized_results = {}

        # Process each result (key is run_id as string, value is the result data)
        for run_id_str, result_data in results.items():
            run_id = int(run_id_str)

            # Create directory for this separation
            run_dir = os.path.join(output_dir, f"run_{run_id}")
            os.makedirs(run_dir, exist_ok=True)

            # Download foreground audio
            foreground_url = result_data["foreground_audio_url"]
            foreground_path = os.path.join(run_dir, "foreground.flac")
            foreground_success = download_audio(foreground_url, foreground_path)

            # Download background audio
            background_url = result_data["background_audio_url"]
            background_path = os.path.join(run_dir, "background.flac")
            background_success = download_audio(background_url, background_path)

            # Store organized result info
            organized_results[run_id] = {
                "directory": run_dir,
                "foreground_url": foreground_url,
                "background_url": background_url,
                "foreground_path": foreground_path if foreground_success else None,
                "background_path": background_path if background_success else None,
                "download_success": foreground_success and background_success,
            }

            status = "✓" if organized_results[run_id]["download_success"] else "✗"
            print(f"{status} Organized separation results for run_id {run_id}")

        print(f"\nProcessed {len(organized_results)} separation results")
        print(f"Results saved to: {output_dir}")

        return organized_results

    except requests.exceptions.HTTPError as e:
        print(f"HTTP Error {e.response.status_code}: {e.response.text}")
        return None
    except requests.exceptions.RequestException as e:
        print(f"Request failed: {e}")
        return None
    except Exception as e:
        print(f"Unexpected error: {e}")
        return None


def print_summary(organized_results: Dict) -> None:
    """
    Print a summary of organized results
    """
    print("\n" + "=" * 50)
    print("AUDIO SEPARATION RESULTS SUMMARY")
    print("=" * 50)

    successful_downloads = 0
    failed_downloads = 0

    for run_id, result in organized_results.items():
        status = "SUCCESS" if result["download_success"] else "FAILED"
        print(f"Run ID {run_id}: {status}")
        print(f"  Directory: {result['directory']}")

        if result["download_success"]:
            successful_downloads += 1
            print(f"  Foreground: {result['foreground_path']}")
            print(f"  Background: {result['background_path']}")
        else:
            failed_downloads += 1
            print("Download failed for one or both audio files")
        print()

    print(f"Total: {len(organized_results)} results")
    print(f"Successful: {successful_downloads}")
    print(f"Failed: {failed_downloads}")


# Example usage
if __name__ == "__main__":
    # Example run IDs from your response
    run_ids_to_process = [112233, 456789]
    output_directory = "./audio_separation_results"

    # Organize the results
    results = organize_separation_results(run_ids_to_process, output_directory)

    if results:
        print_summary(results)

        # Example: Get all foreground audio paths for further processing
        foreground_paths = [
            result["foreground_path"]
            for result in results.values()
            if result["foreground_path"]
        ]

        if foreground_paths:
            print("\nForeground audio files ready for processing:")
            for path in foreground_paths:
                print(f"  - {path}")