#!/usr/bin/env python3
"""
Audio Analysis Example for FMUS-VID.

This example demonstrates how to use the audio analysis tools in FMUS-VID
to analyze audio characteristics in videos.
"""

import os
import sys
import argparse
from pathlib import Path
import logging
import numpy as np
import matplotlib.pyplot as plt
import tempfile
import json
import librosa
import soundfile as sf

# Add parent directory to path to allow running from examples directory
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from fmusvid.operations.analysis import AudioAnalyzer
from fmusvid import Video

# Set up logging
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def extract_audio(video_path: str, output_path: str = None) -> str:
    """
    Extract audio from a video file.

    Args:
        video_path: Path to the video file
        output_path: Path to save the extracted audio (optional)

    Returns:
        Path to the extracted audio file
    """
    if output_path is None:
        # Create temporary file
        temp_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        output_path = temp_file.name
        temp_file.close()

    logger.info(f"Extracting audio from {video_path} to {output_path}")

    # Extract audio using ffmpeg
    import subprocess
    result = subprocess.run([
        "ffmpeg", "-y", "-i", video_path,
        "-q:a", "0", "-map", "a",
        "-c:a", "pcm_s16le",
        output_path
    ], capture_output=True, text=True)

    if result.returncode != 0:
        # Check if the error is because there's no audio stream
        if "Output file is empty" in result.stderr or "Stream #0:0: not a audio" in result.stderr:
            raise ValueError("No audio stream found in video")
        else:
            raise RuntimeError(f"ffmpeg error: {result.stderr}")

    return output_path

def analyze_audio_characteristics(audio_path: str, output_dir: str = None) -> dict:
    """
    Analyze audio characteristics from an audio file.

    Args:
        audio_path: Path to the audio file
        output_dir: Directory to save visualizations (optional)

    Returns:
        Dictionary with analysis results
    """
    logger.info(f"Analyzing audio characteristics in {audio_path}")

    try:
        # Test if librosa will work (catch LLVM errors early)
        import warnings
        warnings.filterwarnings('error')
        y, sr = librosa.load(audio_path, sr=None)
        warnings.resetwarnings()
    except Exception as e:
        logger.error(f"Error loading audio file: {e}")
        # Return a minimal result structure
        return {
            "general": {"loudness": {"mean_db": -60}},
            "beats": {"count": 0, "timestamps": []},
            "spectrum": {"frequency_bands": {}}
        }

    # Create audio analyzer
    analyzer = AudioAnalyzer(sample_rate=sr)

    # Analyze audio
    analysis = analyzer.analyze_audio(y)

    # Detect beats
    beats = analyzer.detect_beats(y)

    # Analyze frequency spectrum
    spectrum = analyzer.analyze_frequency_spectrum(y)

    # Combine results
    results = {
        "general": analysis,
        "beats": {
            "count": len(beats),
            "timestamps": beats[:10]  # Show first 10 beats
        },
        "spectrum": spectrum
    }

    # Create visualizations if output directory is specified
    if output_dir:
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True, parents=True)

        # Waveform plot
        plt.figure(figsize=(12, 4))
        plt.plot(np.linspace(0, len(y) / sr, len(y)), y)
        plt.title("Audio Waveform")
        plt.xlabel("Time (s)")
        plt.ylabel("Amplitude")
        plt.grid(True)
        plt.savefig(str(output_dir / "waveform.png"), dpi=150)
        plt.close()

        # Spectrogram plot
        plt.figure(figsize=(12, 6))
        D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
        plt.subplot(2, 1, 1)
        librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='log')
        plt.colorbar(format='%+2.0f dB')
        plt.title('Spectrogram')

        # Beat plot
        plt.subplot(2, 1, 2)
        plt.plot(np.linspace(0, len(y) / sr, len(y)), y)
        for beat in beats:
            plt.axvline(x=beat, color='r', alpha=0.5)
        plt.title(f"Beats (total: {len(beats)})")
        plt.xlabel("Time (s)")
        plt.tight_layout()
        plt.savefig(str(output_dir / "spectrogram_beats.png"), dpi=150)
        plt.close()

        # Frequency bands plot
        bands = spectrum["frequency_bands"]
        plt.figure(figsize=(10, 6))

        # Bar chart of frequency band means
        band_names = list(bands.keys())
        band_means = [bands[name]["mean"] for name in band_names]

        plt.bar(band_names, band_means)
        plt.title("Frequency Band Distribution")
        plt.ylabel("Energy")
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.savefig(str(output_dir / "frequency_bands.png"), dpi=150)
        plt.close()

        # Save analysis as JSON
        with open(output_dir / "audio_analysis.json", "w") as f:
            json.dump(results, f, indent=2)

    return results

def detect_silent_segments(video_path: str, output_dir: str = None) -> list:
    """
    Detect silent segments in a video.

    Args:
        video_path: Path to the video file
        output_dir: Directory to save visualizations (optional)

    Returns:
        List of silent segments as (start_time, end_time) tuples
    """
    logger.info(f"Detecting silent segments in {video_path}")

    # Create audio analyzer
    analyzer = AudioAnalyzer()

    # Detect silent segments
    silent_segments = analyzer.detect_silence(
        video_path,
        threshold=0.05,  # Adjust threshold based on your video
        min_duration=0.5  # Minimum silence duration in seconds
    )

    # Log results
    logger.info(f"Detected {len(silent_segments)} silent segments")
    for i, (start, end) in enumerate(silent_segments):
        logger.info(f"Silence {i+1}: {start:.2f}s - {end:.2f}s (duration: {end-start:.2f}s)")

    # Create visualization if output directory is specified
    if output_dir and silent_segments:
        output_dir = Path(output_dir)
        output_dir.mkdir(exist_ok=True, parents=True)

        # Create visualization of silent segments
        plt.figure(figsize=(12, 4))

        # Plot timeline
        video = Video(video_path)
        duration = video.duration

        for i, (start, end) in enumerate(silent_segments):
            plt.axvspan(start, end, alpha=0.3, color='red', label='Silence' if i == 0 else "")

        plt.xlim(0, duration)
        plt.xlabel("Time (seconds)")
        plt.yticks([])
        plt.title(f"Silent Segments in {Path(video_path).name}")
        plt.grid(True, axis='x')
        if silent_segments:
            plt.legend()
        plt.savefig(str(output_dir / "silent_segments.png"), dpi=150)
        plt.close()

        # Create annotated video with silence indicators
        output_video_path = output_dir / "silence_annotated.mp4"

        logger.info(f"Creating annotated video with silence indicators: {output_video_path}")
        video_with_indicators = video

        # Add silence indicators as text overlays
        for i, (start, end) in enumerate(silent_segments):
            video_with_indicators = video_with_indicators.add_text(
                f"SILENCE {i+1}",
                position=(0.05, 0.1),  # Top-left position (relative)
                size=30,
                color=(255, 0, 0),
                start_time=start,
                end_time=end,
                relative_position=True
            )

        # Save annotated video
        video_with_indicators.save(str(output_video_path))

    return silent_segments

def create_audio_html_report(video_path: str, output_dir: str):
    """
    Create a comprehensive audio analysis HTML report.

    Args:
        video_path: Path to the video file
        output_dir: Directory to save the report
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True, parents=True)

    # Extract audio for analysis
    try:
        audio_path = extract_audio(video_path)
    except (ValueError, RuntimeError) as e:
        logger.warning(f"Could not extract audio from video: {e}")
        logger.info("Creating a simple HTML report indicating no audio found")
        # Create a simple HTML report indicating no audio
        report_path = output_dir / "audio_analysis_report.html"
        video = Video(video_path)
        with open(report_path, "w") as f:
            f.write(f"""<!DOCTYPE html>
<html>
<head>
    <title>Audio Analysis Report: {Path(video_path).name}</title>
    <style>
        body {{ font-family: Arial, sans-serif; line-height: 1.6; margin: 0; padding: 20px; }}
        .container {{ max-width: 1200px; margin: 0 auto; }}
        .warning {{ background-color: #fff3cd; border: 1px solid #ffc107; padding: 15px; border-radius: 5px; }}
    </style>
</head>
<body>
    <div class="container">
        <h1>Audio Analysis Report</h1>
        <div class="warning">
            <h2>No Audio Found</h2>
            <p>The video file "{Path(video_path).name}" does not contain an audio stream or the audio could not be extracted.</p>
            <p>Duration: {video.duration:.2f} seconds</p>
            <p>Resolution: {video.width}x{video.height}</p>
        </div>
    </div>
</body>
</html>
""")
        logger.info(f"Audio analysis report saved to {report_path}")
        return

    try:
        # Analyze audio characteristics
        characteristics_dir = output_dir / "characteristics"
        characteristics = analyze_audio_characteristics(audio_path, characteristics_dir)

        # Detect silent segments
        silence_dir = output_dir / "silence"
        silent_segments = detect_silent_segments(video_path, silence_dir)

        # Create HTML report
        report_path = output_dir / "audio_analysis_report.html"

        # Basic video info
        video = Video(video_path)

        info = {
            "filename": Path(video_path).name,
            "duration": video.duration,
            "silent_segments": len(silent_segments),
            "beats_detected": characteristics["beats"]["count"],
            "mean_loudness": characteristics["general"]["loudness"]["mean_db"]
        }

        with open(report_path, "w") as f:
            f.write(f"""<!DOCTYPE html>
<html>
<head>
    <title>Audio Analysis Report: {info['filename']}</title>
    <style>
        body {{ font-family: Arial, sans-serif; line-height: 1.6; margin: 0; padding: 20px; }}
        h1, h2, h3 {{ color: #333; }}
        .container {{ max-width: 1200px; margin: 0 auto; }}
        .info-table {{ border-collapse: collapse; width: 100%; margin-bottom: 20px; }}
        .info-table th, .info-table td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
        .info-table th {{ background-color: #f2f2f2; }}
        .section {{ margin-bottom: 40px; }}
        .visualization {{ margin: 20px 0; text-align: center; }}
        .visualization img {{ max-width: 100%; }}
        pre {{ background-color: #f5f5f5; padding: 15px; overflow: auto; }}
    </style>
</head>
<body>
    <div class="container">
        <h1>Audio Analysis Report</h1>

        <div class="section">
            <h2>Video Information</h2>
            <table class="info-table">
                <tr><th>Filename</th><td>{info['filename']}</td></tr>
                <tr><th>Duration</th><td>{info['duration']:.2f} seconds</td></tr>
                <tr><th>Silent Segments</th><td>{info['silent_segments']}</td></tr>
                <tr><th>Beats Detected</th><td>{info['beats_detected']}</td></tr>
                <tr><th>Mean Loudness</th><td>{info['mean_loudness']:.2f} dB</td></tr>
            </table>
        </div>

        <div class="section">
            <h2>Audio Waveform & Spectrum</h2>
            <div class="visualization">
                <img src="characteristics/waveform.png" alt="Audio Waveform">
            </div>
            <div class="visualization">
                <img src="characteristics/spectrogram_beats.png" alt="Spectrogram & Beats">
            </div>
        </div>

        <div class="section">
            <h2>Frequency Analysis</h2>
            <div class="visualization">
                <img src="characteristics/frequency_bands.png" alt="Frequency Bands">
            </div>
            <h3>Frequency Bands Data</h3>
            <pre>{json.dumps(characteristics["spectrum"]["frequency_bands"], indent=4)}</pre>
        </div>
""")

            if silent_segments:
                f.write(f"""
        <div class="section">
            <h2>Silent Segments</h2>
            <div class="visualization">
                <img src="silence/silent_segments.png" alt="Silent Segments">
            </div>
            <h3>Silence Data</h3>
            <table class="info-table">
                <tr><th>#</th><th>Start Time (s)</th><th>End Time (s)</th><th>Duration (s)</th></tr>
""")

                for i, (start, end) in enumerate(silent_segments):
                    f.write(f"""                <tr>
                    <td>{i+1}</td>
                    <td>{start:.2f}</td>
                    <td>{end:.2f}</td>
                    <td>{end-start:.2f}</td>
                </tr>
""")

                f.write(f"""
            </table>
            <p>An annotated video with silence indicators has been created: <a href="silence/silence_annotated.mp4">silence_annotated.mp4</a></p>
        </div>
""")

            f.write(f"""
    </div>
</body>
</html>
""")

        logger.info(f"Audio analysis report saved to {report_path}")

    finally:
        # Clean up temporary audio file
        try:
            if audio_path and Path(audio_path).exists() and "NamedTemporaryFile" in audio_path:
                os.unlink(audio_path)
        except:
            pass

def main():
    """Main function to run the example."""
    # Hardcoded sample videos to try
    sample_video_paths = [
        os.path.join("test_media", "test_video.mp4"),
        os.path.join("test_media", "sample.mp4"),
        os.path.join("test_media", "sample2.mp4")
    ]

    # Find the first existing sample video
    video_path = None
    for path in sample_video_paths:
        if os.path.exists(path):
            video_path = path
            break

    # If no sample videos found, create a test video with audio
    if video_path is None:
        logger.warning("No sample videos found. Creating a test video with audio.")

        # Create test_media directory if it doesn't exist
        os.makedirs("test_media", exist_ok=True)

        # Create a test video with audio using ffmpeg
        test_video_path = os.path.join("test_media", "generated_test_audio.mp4")

        # Generate a test video with audio using ffmpeg
        import subprocess
        result = subprocess.run([
            "ffmpeg", "-f", "lavfi", "-i", "color=c=blue:s=640x360:d=5",
            "-f", "lavfi", "-i", "sine=frequency=1000:duration=5",
            "-c:v", "libx264", "-c:a", "aac", "-shortest",
            "-y", test_video_path
        ], capture_output=True, text=True)

        if result.returncode == 0:
            logger.info(f"Created test video with audio: {test_video_path}")
            video_path = test_video_path
        else:
            logger.error(f"Failed to create test video: {result.stderr}")
            return 1

    # Set output directory
    output_dir = "audio_analysis_output"

    try:
        # Create analysis report
        logger.info(f"Analyzing video: {video_path}")
        create_audio_html_report(video_path, output_dir)
        logger.info(f"Audio analysis completed successfully. Results saved to {output_dir}")
        logger.info(f"Open {output_dir}/audio_analysis_report.html to view the report")
        return 0
    except Exception as e:
        logger.exception(f"Error during audio analysis: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())
