import gradio as gr
import torch
import torchaudio
import logging
import tempfile
import os
import sys
from pathlib import Path
import numpy as np
from typing import Optional, Tuple
import time
import traceback

# Add hf_AC to path
current_dir = Path(__file__).parent
hf_ac_path = current_dir / "hf_AC"
if hf_ac_path.exists():
    sys.path.insert(0, str(hf_ac_path))

# Configuration for HF Space
EXAMPLE_PROMPTS = [
    "Crackling fireplace with gentle flames",
    "Ocean waves crashing on rocky shore", 
    "Forest ambience with bird songs",
    "Keyboard typing sounds",
    "Footsteps on wooden floor",
    "Rain on metal roof"
]

USAGE_TIPS = """
### 💡 Usage Tips

**Basic Settings:**
- **Video Quality**: Use clear, well-lit videos, recommended 1-15 seconds
- **Reference Audio**: Provide clear audio clips as timbre reference
- **CFG Strength**: Between 1-8, higher values follow description more closely

**Advanced Features:**
- **mask_away_clip**: Enable when video content differs significantly from desired audio
- **Fine-grained Control**: Use reference audio for precise timbre and style control
- **Zero-shot Generation**: Generate novel sound combinations without training

**Application Scenarios:**
- Film post-production audio
- Game sound effect creation
- Music composition assistance
- Sound design experimentation
"""

# Check and install missing dependencies
def check_dependencies():
    """Check if all required packages are available"""
    missing_packages = []
    required_packages = [
        'torch', 'torchaudio', 'numpy', 'scipy', 'librosa', 
        'torchdiffeq', 'einops', 'hydra', 'tensordict', 'av'
    ]
    
    for package in required_packages:
        try:
            if package == 'hydra':
                __import__('hydra')
            elif package == 'av':
                __import__('av')
            else:
                __import__(package)
        except ImportError:
            missing_packages.append(package)
    
    return missing_packages

# Import hf_AC modules with error handling
try:
    # First check basic dependencies
    missing_deps = check_dependencies()
    if missing_deps:
        print(f"Warning: Missing dependencies: {missing_deps}")
        print("Some dependencies may be installing in the background...")
    
    from hf_AC.mmaudio.eval_utils import (ModelConfig, all_model_cfg, generate, load_video, 
                                          setup_eval_logging)
    from hf_AC.mmaudio.model.flow_matching import FlowMatching
    from hf_AC.mmaudio.model.networks import MMAudio, get_my_mmaudio
    from hf_AC.mmaudio.model.utils.features_utils import FeaturesUtils
    from hf_AC.inf import Audio
    
    # Setup logging
    setup_eval_logging()
    log = logging.getLogger()
    HF_AC_AVAILABLE = True
    print("✅ hf_AC modules loaded successfully!")
    
except ImportError as e:
    print(f"Warning: hf_AC modules not available: {e}")
    print("This may be due to missing dependencies. Please wait for installation to complete.")
    log = logging.getLogger()
    HF_AC_AVAILABLE = False

class AudioFoleyModel:
    def __init__(self):
        self.device = 'cpu'
        if torch.cuda.is_available():
            self.device = 'cuda'
        elif torch.backends.mps.is_available():
            self.device = 'mps'
        
        self.dtype = torch.bfloat16
        self.model = None
        self.net = None
        self.fm = None
        self.feature_utils = None
        
    def load_model(self, variant='large_44k', model_path=None):
        """Load the hf_AC model with progress updates"""
        global model_loading_status
        
        try:
            if not HF_AC_AVAILABLE:
                return "❌ hf_AC modules not available. Please install the hf_AC package."
            
            if variant not in all_model_cfg:
                available_variants = list(all_model_cfg.keys()) if all_model_cfg else []
                return f"❌ Unknown model variant: {variant}. Available: {available_variants}"
            
            # Step 1: Initialize model config
            model_loading_status = "🔧 Initializing model configuration..."
            log.info(f"Loading model variant: {variant}")
            self.model: ModelConfig = all_model_cfg[variant]
            
            # Step 2: Download model components
            model_loading_status = "📥 Downloading model components..."
            try:
                self.model.download_if_needed()
            except Exception as e:
                log.warning(f"Could not download model components: {e}")
            
            # Step 3: Download main model weights
            model_loading_status = "📥 Downloading main model weights..."
            if not hasattr(self.model, 'model_path') or not self.model.model_path or not Path(self.model.model_path).exists():
                try:
                    from huggingface_hub import hf_hub_download
                    log.info("Downloading main model weights from HuggingFace...")
                    
                    # Create weights directory
                    weights_dir = Path("weights")
                    weights_dir.mkdir(exist_ok=True)
                    
                    # Download model.pth from HuggingFace
                    model_file = hf_hub_download(
                        repo_id="FF2416/AC-Foley",
                        filename="model.pth",
                        local_dir=str(weights_dir),
                        local_dir_use_symlinks=False
                    )
                    self.model.model_path = Path(model_file)
                    log.info(f"✅ Downloaded model weights to {model_file}")
                    
                except Exception as e:
                    log.warning(f"Could not download main model weights: {e}")
                    log.info("Will proceed with available components only")
            
            # Set custom model path if provided
            if model_path and os.path.exists(model_path):
                self.model.model_path = Path(model_path)
                log.info(f"Using custom model path: {model_path}")
            
            # Step 4: Load neural network
            model_loading_status = "🧠 Loading neural network..."
            self.net: MMAudio = get_my_mmaudio(self.model.model_name).to(self.device, self.dtype).eval()
            
            # Step 5: Load weights
            model_loading_status = "⚖️ Loading model weights..."
            if hasattr(self.model, 'model_path') and self.model.model_path and Path(self.model.model_path).exists():
                try:
                    weights = torch.load(self.model.model_path, map_location=self.device, weights_only=True)
                    self.net.load_weights(weights['weights'])
                    log.info(f'✅ Loaded weights from {self.model.model_path}')
                except Exception as e:
                    log.error(f"Failed to load weights: {e}")
                    model_loading_status = f"❌ Failed to load model weights: {e}"
                    return model_loading_status
            else:
                log.warning('⚠️ No model weights found, using default initialization')
                model_loading_status = "⚠️ 模型组件已加载，但主权重不可用。某些功能可能受限。"
                return model_loading_status
            
            # Step 6: Initialize flow matching
            model_loading_status = "🌊 Initializing flow matching..."
            self.fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=25)
            
            # Step 7: Initialize feature utils
            model_loading_status = "🔧 Initializing feature utilities..."
            try:
                self.feature_utils = FeaturesUtils(
                    tod_vae_ckpt=self.model.vae_path,
                    synchformer_ckpt=self.model.synchformer_ckpt,
                    enable_conditions=True,
                    mode=self.model.mode,
                    bigvgan_vocoder_ckpt=self.model.bigvgan_16k_path,
                    need_vae_encoder=True
                )
                self.feature_utils = self.feature_utils.to(self.device, self.dtype).eval()
            except Exception as e:
                log.error(f"Failed to initialize feature utils: {e}")
                model_loading_status = f"❌ Failed to initialize feature utilities: {e}"
                return model_loading_status
            
            # Step 8: Complete
            model_loading_status = "✅ Model loaded successfully! Ready to generate audio."
            return model_loading_status
            
        except Exception as e:
            error_msg = f"❌ Model loading error: {str(e)}"
            log.error(error_msg)
            model_loading_status = error_msg
            return error_msg
    
    def generate_audio(self, video_file, prompt: str, negative_prompt: str = "", 
                      duration: float = 8.0, cfg_strength: float = 4.5, 
                      seed: int = 42, reference_audio: str = None, mask_away_clip: bool = False) -> Tuple[Optional[str], str]:
        """Generate audio from video and text prompt"""
        try:
            # Validation checks
            if not HF_AC_AVAILABLE:
                return None, "❌ hf_AC modules not available."
            
            if self.net is None or self.feature_utils is None:
                return None, "❌ Model not loaded. Please load the model first."
            
            if video_file is None:
                return None, "❌ Please upload a video file."
            
            log.info(f'🎬 Processing video: {video_file}')
            if prompt.strip():
                log.info(f'📝 Prompt: "{prompt}"')
            else:
                log.info('📝 No prompt provided - will generate based on video content')
            if reference_audio:
                log.info(f'🎵 Reference audio: {reference_audio}')
            
            # Load and process reference audio if provided
            reference_audio_tensor = None
            if reference_audio and os.path.exists(reference_audio):
                try:
                    # Use the same Audio class from hf_AC
                    SAMPLE_RATE = 44100
                    audio_processor = Audio([reference_audio], SAMPLE_RATE)
                    audio_list = audio_processor.load_audio()
                    if audio_list:
                        reference_audio_tensor = audio_list[0]
                        log.info(f'🎵 Reference audio loaded: {reference_audio_tensor.shape}')
                except Exception as e:
                    log.warning(f"Failed to load reference audio: {e}")
                    reference_audio_tensor = None
            
            # Load and process video
            try:
                video_path = Path(video_file)
                if not video_path.exists():
                    return None, f"❌ Video file not found: {video_file}"
                
                video_info = load_video(video_path, duration)
                clip_frames = video_info.clip_frames
                sync_frames = video_info.sync_frames
                duration_sec = video_info.duration_sec
                
                log.info(f'📹 Video loaded: {duration_sec:.2f}s duration')
                
            except Exception as e:
                return None, f"❌ Failed to load video: {str(e)}"
            
            # Prepare frames
            if mask_away_clip:
                clip_frames = None  # Mask away clip frames when video and audio don't match well
                log.info("🎭 Using mask_away_clip: ignoring visual features")
            else:
                clip_frames = clip_frames.unsqueeze(0) if clip_frames is not None else None
            sync_frames = sync_frames.unsqueeze(0)
            
            # Update model sequence configuration
            try:
                self.model.seq_cfg.duration = duration_sec
                # Set audio sample count based on reference audio or default
                if reference_audio_tensor is not None:
                    self.model.seq_cfg.audio_num_sample = reference_audio_tensor.shape[0]
                else:
                    self.model.seq_cfg.audio_num_sample = 89088  # Default for 44kHz
                    
                self.net.update_seq_lengths(
                    self.model.seq_cfg.latent_seq_len, 
                    self.model.seq_cfg.clip_seq_len, 
                    self.model.seq_cfg.sync_seq_len, 
                    self.model.seq_cfg.audio_seq_len
                )
            except Exception as e:
                return None, f"❌ Failed to configure model: {str(e)}"
            
            # Generate audio
            try:
                log.info('🎵 Generating audio...')
                start_time = time.time()
                
                with torch.inference_mode():
                    audios = generate(
                        clip_frames,
                        sync_frames, 
                        [prompt], 
                        reference_audio_tensor,  # Use reference audio if provided
                        negative_text=[negative_prompt] if negative_prompt.strip() else None,
                        feature_utils=self.feature_utils,
                        net=self.net,
                        fm=self.fm,
                        rng=torch.Generator(device=self.device).manual_seed(seed),
                        cfg_strength=cfg_strength
                    )
                
                generation_time = time.time() - start_time
                log.info(f'⏱️ Generation completed in {generation_time:.2f}s')
                
            except Exception as e:
                return None, f"❌ Audio generation failed: {str(e)}"
            
            # Save generated audio
            try:
                audio = audios.float().cpu()[0]
                
                # Create output filename with timestamp
                timestamp = int(time.time())
                output_filename = f"generated_audio_{timestamp}.wav"
                permanent_path = f"/tmp/{output_filename}"
                
                # Save audio file with fallback methods
                try:
                    # Try with torchaudio first
                    torchaudio.save(permanent_path, audio, self.model.seq_cfg.sampling_rate)
                except Exception as e:
                    log.warning(f"torchaudio.save failed: {e}, trying alternative method...")
                    try:
                        # Fallback: use soundfile if available
                        import soundfile as sf
                        sf.write(permanent_path, audio.numpy().T, self.model.seq_cfg.sampling_rate)
                    except ImportError:
                        try:
                            # Fallback: use scipy.io.wavfile
                            from scipy.io.wavfile import write
                            # Convert to int16 for wav format
                            audio_int16 = (audio * 32767).clamp(-32768, 32767).to(torch.int16)
                            write(permanent_path, self.model.seq_cfg.sampling_rate, audio_int16.numpy().T)
                        except Exception as e2:
                            return None, f"❌ Audio saving failed: {str(e2)}"
                
                # Verify file was created
                if not os.path.exists(permanent_path):
                    return None, "❌ Failed to save audio file"
                
                file_size = os.path.getsize(permanent_path) / 1024  # KB
                success_msg = f"✅ Audio generated successfully!\n"
                success_msg += f"📊 Duration: {duration_sec:.2f}s | "
                success_msg += f"Size: {file_size:.1f}KB | "
                success_msg += f"Generation time: {generation_time:.2f}s"
                
                return permanent_path, success_msg
                
            except Exception as e:
                return None, f"❌ Failed to save audio: {str(e)}"
                
        except Exception as e:
            error_msg = f"❌ Unexpected error: {str(e)}\n{traceback.format_exc()}"
            log.error(error_msg)
            return None, error_msg

# Global model instance - initialized once
audio_model = None
model_loading_status = "Not initialized"

def initialize_model():
    """Initialize model once at startup"""
    global audio_model, model_loading_status
    
    if audio_model is None:
        try:
            model_loading_status = "Initializing model..."
            audio_model = AudioFoleyModel()
            load_result = audio_model.load_model()
            model_loading_status = load_result
            return load_result
        except Exception as e:
            model_loading_status = f"❌ Model initialization failed: {str(e)}"
            return model_loading_status
    else:
        return "✅ Model already loaded"

def generate_audio_interface(video_file, audio_file, prompt, duration, cfg_strength, mask_away_clip):
    """Interface function for generating audio"""
    global audio_model, model_loading_status
    
    # Check if model is loaded
    if audio_model is None or audio_model.net is None:
        return None, "❌ Model not loaded. Please wait for initialization to complete or refresh the page."
    
    # Use fixed seed for consistency in HF Space
    seed = 42
    negative_prompt = ""  # Simplified interface
    
    audio_path, message = audio_model.generate_audio(
        video_file, prompt, negative_prompt, duration, cfg_strength, seed, audio_file, mask_away_clip
    )
    return audio_path, message

def get_model_status():
    """Get current model loading status"""
    global model_loading_status
    return model_loading_status

# Create Gradio interface
with gr.Blocks(title="hf_AC Audio Foley Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎵 AC-Foley: Reference-Audio-Guided Video-to-Audio Synthesis
    
    ## 📖 About
    AC-Foley is a reference-audio-guided video-to-audio synthesis model that enables precise fine-grained sound synthesis. Unlike traditional text-dependent methods, AC-Foley directly leverages reference audio to achieve precise control over generated sounds, addressing the ambiguity of textual descriptions in micro-acoustic features.
    
    ## ✨ Key Features
    - **Fine-grained Sound Synthesis**: Generate footsteps with distinct timbres (wood, marble, gravel, etc.)
    - **Timbre Transfer**: Transform violin melodies into bright, piercing suona tones
    - **Zero-shot Generation**: Create unique sound effects without specialized training
    - **Visual-Audio Alignment**: Automatically generate matching audio from video content
    
    *Based on paper: [AC-Foley: Reference-Audio-Guided Video-to-Audio Synthesis with Acoustic Transfer](https://openreview.net/forum?id=URPXhnWdBF)*
    """)
    
    # Model status display - will be updated automatically
    model_status = gr.Textbox(
        label="Model Status", 
        value=model_loading_status,
        interactive=False
    )
    
    # Add a refresh button for status
    refresh_status_btn = gr.Button("🔄 Refresh Status", size="sm")
    refresh_status_btn.click(
        fn=get_model_status,
        outputs=model_status
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            # Required inputs
            gr.Markdown("### 📹 Required Input")
            video_input = gr.Video(
                label="Video File - Upload video for audio generation",
                format="mp4"
            )
            
            # Optional inputs
            gr.Markdown("### 🎛️ Optional Inputs")
            audio_input = gr.Audio(
                label="Reference Audio - Provide timbre, style, rhythm reference (fine-grained control)",
                type="filepath",
                sources=["upload"],
                format="wav"
            )
            
            prompt_input = gr.Textbox(
                label="Text Prompt - Describe desired audio type (leave empty for auto-generation from video)",
                placeholder="e.g., 'footsteps', 'metal clang', 'bird chirping'",
                lines=2
            )
            
            # Advanced options
            with gr.Accordion("🔧 Advanced Options", open=False):
                with gr.Row():
                    duration_slider = gr.Slider(
                        minimum=1.0,
                        maximum=15.0,
                        value=8.0,
                        step=0.5,
                        label="Duration (seconds)"
                    )
                    
                    cfg_strength_slider = gr.Slider(
                        minimum=1.0,
                        maximum=8.0,
                        value=4.5,
                        step=0.1,
                        label="CFG Strength"
                    )
                
                mask_away_clip = gr.Checkbox(
                    label="Ignore Visual Features (mask_away_clip) - Enable when video and reference audio differ significantly",
                    value=False
                )
        
        with gr.Column(scale=1):
            # Usage guide
            gr.Markdown("### 📋 Usage Guide")
            gr.Markdown("""
            **Four Generation Modes:**
            
            1️⃣ **Video Only**: Upload video only
            - Auto-generate audio from visual content
            
            2️⃣ **Video + Reference Audio**: Upload video + audio
            - Use reference audio's timbre and style
            - Achieve fine-grained timbre control
            
            3️⃣ **Video + Text**: Upload video + text
            - Generate specified audio type from text description
            
            4️⃣ **Complete Mode**: Video + Audio + Text
            - Most precise control method
            - Combine visual, timbral, and semantic guidance
            """)
            
            # Example prompts
            gr.Markdown("### 🎯 Example Prompts")
            example_buttons = []
            for prompt in EXAMPLE_PROMPTS[:4]:
                btn = gr.Button(prompt, size="sm")
                example_buttons.append(btn)
                btn.click(
                    fn=lambda p=prompt: p,
                    outputs=prompt_input
                )
    
    # Generate button
    generate_btn = gr.Button("🎵 Generate Audio", variant="primary", size="lg")
    
    # Output area
    gr.Markdown("### 🎧 Generated Results")
    audio_output = gr.Audio(
        label="Generated Audio",
        type="filepath",
        format="wav",
        autoplay=False
    )
    
    generation_status = gr.Textbox(
        label="Generation Status", 
        interactive=False,
        lines=2
    )
    
    # 绑定生成事件
    generate_btn.click(
        fn=generate_audio_interface,
        inputs=[
            video_input, audio_input, prompt_input, 
            duration_slider, cfg_strength_slider, mask_away_clip
        ],
        outputs=[audio_output, generation_status]
    )
    
    with gr.Accordion("💡 Detailed Information", open=False):
        gr.Markdown(USAGE_TIPS)
        
        gr.Markdown("""
        ### 🎬 Application Examples
        
        **Fine-grained Sound Synthesis:**
        - "Footsteps on wooden floor" + reference audio → Specific timbre footsteps
        - "Metal collision" + different reference audio → Iron vs. copper distinction
        
        **Timbre Transfer:**
        - Piano melody video + violin reference audio → Violin playing same melody
        - Human humming + instrument reference → Instrumental version
        
        **Creative Sound Effects:**
        - Sci-fi scene video + real sound reference → Unique sci-fi effects
        - Animation video + real sound effects → Cartoon-reality hybrid effects
        
        ### 📚 Technical Details
        - Model based on diffusion models and audio conditioning mechanisms
        - Supports 44.1kHz high-quality audio generation
        - Achieves visual-audio-text multimodal alignment
        """)
    
    # Auto-initialize model on startup
    demo.load(
        fn=initialize_model,
        outputs=[model_status]
    )

# Initialize model when module is imported (for HF Space)
if HF_AC_AVAILABLE:
    print("🚀 Starting model initialization...")
    initialize_model()
    print(f"📊 Model status: {model_loading_status}")

if __name__ == "__main__":
    # HF Space will handle the server configuration
    demo.launch()