#!/usr/bin/env python3
"""
Download audio files for records with cover_ids from the training dataset.
Samples records, downloads original and cover audio files, and converts to MP3.
"""

import argparse
import json
import os
import random
import subprocess
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional
import logging
from tqdm import tqdm
import concurrent.futures
from threading import Lock


def setup_logging(output_dir: Path = None):
    """Set up logging configuration."""
    if output_dir:
        output_dir.mkdir(parents=True, exist_ok=True)
        log_file = output_dir / f"download_covers_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"

        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s",
            handlers=[logging.FileHandler(log_file), logging.StreamHandler()],
        )
    else:
        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

    return logging.getLogger(__name__)


class CoverAudioDownloader:
    """Download and convert audio files for records with covers."""

    def __init__(self, dataset_file: str, output_dir: str, sample_size: int = 10, max_workers: int = 4):
        """
        Initialize downloader.

        Args:
            dataset_file: Path to dataset JSONL
            output_dir: Output directory for downloaded files
            sample_size: Number of records to sample
            max_workers: Number of parallel download workers
        """
        self.dataset_file = Path(dataset_file)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.sample_size = sample_size
        self.max_workers = max_workers

        self.logger = setup_logging(self.output_dir)

        # Base path for audio files
        self.audio_base_path = Path("/app2/suno/data/auk_v0")

        # Track statistics
        self.stats = {
            "total_scanned": 0,
            "records_with_covers": 0,
            "records_sampled": 0,
            "audio_downloaded": 0,
            "audio_converted": 0,
            "errors": 0,
        }

        # Thread lock for stats
        self.stats_lock = Lock()

    def find_records_with_covers(self, max_scan: int = 1000000) -> List[Dict]:
        """
        Find records that have cover_ids.

        Args:
            max_scan: Maximum records to scan

        Returns:
            List of records with covers
        """
        self.logger.info(f"Scanning for records with cover_ids (max {max_scan:,} records)...")

        records_with_covers = []

        with open(self.dataset_file, "r") as f:
            for line_num, line in enumerate(tqdm(f, desc="Scanning", total=max_scan), 0):
                if line_num >= max_scan:
                    break

                self.stats["total_scanned"] += 1

                if not line.strip():
                    continue

                try:
                    record = json.loads(line.strip())

                    # Check if record has cover_ids
                    if record.get("cover_ids") and len(record["cover_ids"]) > 0:
                        records_with_covers.append(record)
                        self.stats["records_with_covers"] += 1

                except json.JSONDecodeError:
                    continue
                except Exception as e:
                    self.logger.warning(f"Error processing line {line_num}: {e}")

        self.logger.info(f"Found {len(records_with_covers):,} records with covers")
        return records_with_covers

    def sample_records(self, records: List[Dict]) -> List[Dict]:
        """
        Sample records for downloading.

        Args:
            records: List of all records with covers

        Returns:
            Sampled records
        """
        if len(records) <= self.sample_size:
            sampled = records
        else:
            # Random sampling
            sampled = random.sample(records, self.sample_size)

        self.stats["records_sampled"] = len(sampled)
        self.logger.info(f"Sampled {len(sampled)} records for download")

        return sampled

    def get_audio_path(self, record: Dict) -> Optional[Path]:
        """
        Get audio file path for a record using local_filepath.

        Args:
            record: Record dict

        Returns:
            Path to audio file or None
        """
        # Use local_filepath like audioloader.py
        local_filepath = record.get("local_filepath")
        if local_filepath:
            return Path(local_filepath)

        # Fallback to s3_filepath if local not available
        s3_filepath = record.get("s3_filepath")
        if s3_filepath:
            # Extract filename from s3 path and construct local path
            filename = Path(s3_filepath).name
            local_path = self.audio_base_path / "audio" / filename
            if local_path.exists():
                return local_path

        return None

    def convert_audio_to_mp3(self, input_path: Path, output_path: Path) -> bool:
        """
        Convert audio file to MP3 using ffmpeg.

        Args:
            input_path: Input audio file
            output_path: Output MP3 file

        Returns:
            True if successful
        """
        try:
            # Ensure output directory exists
            output_path.parent.mkdir(parents=True, exist_ok=True)

            # FFmpeg command for conversion
            cmd = [
                "ffmpeg",
                "-i",
                str(input_path),
                "-codec:a",
                "libmp3lame",
                "-b:a",
                "192k",
                "-y",  # Overwrite if exists
                str(output_path),
            ]

            # Run conversion
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)

            if result.returncode == 0:
                with self.stats_lock:
                    self.stats["audio_converted"] += 1
                return True
            else:
                self.logger.warning(f"FFmpeg error: {result.stderr}")
                return False

        except subprocess.TimeoutExpired:
            self.logger.error(f"Timeout converting {input_path}")
            return False
        except Exception as e:
            self.logger.error(f"Error converting {input_path}: {e}")
            return False

    def download_record_audio(self, record: Dict) -> Dict:
        """
        Download and convert audio for a single record.

        Args:
            record: Record to process

        Returns:
            Download results
        """
        record_id = record.get("id", "unknown")
        record_dir = self.output_dir / record_id
        record_dir.mkdir(parents=True, exist_ok=True)

        results = {"id": record_id, "original": False, "covers": [], "errors": []}

        # Save record metadata
        metadata_file = record_dir / "metadata.json"
        with open(metadata_file, "w", encoding="utf-8") as f:
            json.dump(record, f, indent=2, ensure_ascii=False)

        # Download original audio - first try direct path
        original_path = Path(f"/app2/suno/data/raw_audio_opus_v0/{record_id}.opus")
        if not original_path.exists():
            # Fallback to using record's filepath
            original_path = self.get_audio_path(record)

        if original_path and original_path.exists():
            output_mp3 = record_dir / "original.mp3"
            if self.convert_audio_to_mp3(original_path, output_mp3):
                results["original"] = True
                with self.stats_lock:
                    self.stats["audio_downloaded"] += 1
                self.logger.info(f"Downloaded original audio for {record_id}")
            else:
                results["errors"].append("Failed to convert original audio")
        else:
            results["errors"].append(f"Original audio not found: {original_path}")

        # Download cover audio files (limit to 3)
        cover_ids = record.get("cover_ids", [])
        max_covers = 3
        for i, cover_id in enumerate(cover_ids[:max_covers]):
            # First try direct path for cover
            cover_path = Path(f"/app2/suno/data/raw_audio_opus_v0/{cover_id}.opus")

            # No fallback needed since we're using direct path

            if cover_path and cover_path.exists():
                output_mp3 = record_dir / f"cover_{cover_id}.mp3"
                if self.convert_audio_to_mp3(cover_path, output_mp3):
                    results["covers"].append(cover_id)
                    with self.stats_lock:
                        self.stats["audio_downloaded"] += 1
                    self.logger.info(f"Downloaded cover {cover_id} for {record_id}")
                else:
                    results["errors"].append(f"Failed to convert cover {cover_id}")
            else:
                self.logger.warning(f"Cover audio file not found: {cover_id} (tried {cover_path})")
                results["errors"].append(f"Cover audio file not found: {cover_id}")

        if results["errors"]:
            with self.stats_lock:
                self.stats["errors"] += len(results["errors"])

        return results

    def run(self):
        """Run the download process."""
        self.logger.info("=" * 80)
        self.logger.info("COVER AUDIO DOWNLOADER")
        self.logger.info("=" * 80)
        self.logger.info(f"Dataset: {self.dataset_file}")
        self.logger.info(f"Output: {self.output_dir}")
        self.logger.info(f"Sample size: {self.sample_size}")
        self.logger.info(f"Workers: {self.max_workers}")

        start_time = time.time()

        # Find records with covers
        records_with_covers = self.find_records_with_covers()

        if not records_with_covers:
            self.logger.warning("No records with covers found!")
            return

        # Sample records
        sampled_records = self.sample_records(records_with_covers)

        # Download audio files in parallel
        self.logger.info(f"\nDownloading audio files...")

        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all tasks
            futures = {
                executor.submit(self.download_record_audio, record): record for record in sampled_records
            }

            # Process completed tasks with progress bar
            with tqdm(total=len(futures), desc="Downloading") as pbar:
                for future in concurrent.futures.as_completed(futures):
                    try:
                        result = future.result()
                        pbar.update(1)

                        # Log results
                        if result["original"] or result["covers"]:
                            self.logger.info(
                                f"Record {result['id']}: "
                                f"original={'✓' if result['original'] else '✗'}, "
                                f"covers={len(result['covers'])}"
                            )

                    except Exception as e:
                        self.logger.error(f"Download error: {e}")
                        pbar.update(1)

        # Generate summary report
        elapsed = time.time() - start_time
        self.generate_report(elapsed)

    def generate_report(self, elapsed_time: float):
        """Generate and save download report."""
        report = {
            "timestamp": datetime.now().isoformat(),
            "configuration": {
                "dataset_file": str(self.dataset_file),
                "output_dir": str(self.output_dir),
                "sample_size": self.sample_size,
                "max_workers": self.max_workers,
            },
            "statistics": {
                "total_scanned": self.stats["total_scanned"],
                "records_with_covers": self.stats["records_with_covers"],
                "records_sampled": self.stats["records_sampled"],
                "audio_downloaded": self.stats["audio_downloaded"],
                "audio_converted": self.stats["audio_converted"],
                "errors": self.stats["errors"],
                "processing_time": round(elapsed_time, 2),
            },
        }

        # Save report
        report_file = self.output_dir / "download_report.json"
        with open(report_file, "w") as f:
            json.dump(report, f, indent=2)

        # Print summary
        self.logger.info("\n" + "=" * 80)
        self.logger.info("DOWNLOAD COMPLETE")
        self.logger.info("=" * 80)
        self.logger.info(f"Records scanned: {self.stats['total_scanned']:,}")
        self.logger.info(f"Records with covers: {self.stats['records_with_covers']:,}")
        self.logger.info(f"Records sampled: {self.stats['records_sampled']}")
        self.logger.info(f"Audio files downloaded: {self.stats['audio_downloaded']}")
        self.logger.info(f"Audio files converted: {self.stats['audio_converted']}")
        self.logger.info(f"Errors: {self.stats['errors']}")
        self.logger.info(f"Processing time: {elapsed_time:.1f} seconds")
        self.logger.info(f"\nOutput directory: {self.output_dir}")


def main():
    """Main execution."""
    parser = argparse.ArgumentParser(description="Download audio files for records with covers")

    parser.add_argument(
        "--dataset-file",
        default="/app2/suno/data/auk_v0/metas_v6_tr.jsonl",
        help="Path to dataset JSONL file",
    )

    parser.add_argument(
        "--output-dir",
        default="/home/vibert/tmp/cover_audio_samples",
        help="Output directory for downloaded audio",
    )

    parser.add_argument(
        "--sample-size", type=int, default=10, help="Number of records to sample and download"
    )

    parser.add_argument(
        "--max-scan", type=int, default=1000000, help="Maximum records to scan for covers"
    )

    parser.add_argument("--workers", type=int, default=4, help="Number of parallel download workers")

    args = parser.parse_args()

    # Create downloader
    downloader = CoverAudioDownloader(
        dataset_file=args.dataset_file,
        output_dir=args.output_dir,
        sample_size=args.sample_size,
        max_workers=args.workers,
    )

    # Run download
    downloader.run()


if __name__ == "__main__":
    main()