"""VoiceGuard - AI Voice Detection App for HuggingFace Spaces."""

import gradio as gr
import numpy as np
import torch
import librosa
from transformers import pipeline
import warnings
import tempfile
import os

warnings.filterwarnings("ignore")

# Model Configuration
MODEL_NAME = "abhishtagatya/hubert-base-960h-itw-deepfake"

# Global classifier (loaded once)
classifier = None

def load_model():
    """Load the deepfake detection model."""
    global classifier
    if classifier is None:
        print("🔄 Loading model...")
        device = 0 if torch.cuda.is_available() else -1
        classifier = pipeline(
            "audio-classification",
            model=MODEL_NAME,
            device=device
        )
        print(f"✅ Model loaded on {'GPU' if device == 0 else 'CPU'}")
    return classifier

def detect_deepfake(audio_input):
    """
    Detect if audio is AI-generated or human.
    
    Args:
        audio_input: Tuple of (sample_rate, audio_array) from Gradio
        
    Returns:
        Dictionary with detection results
    """
    if audio_input is None:
        return "❌ Please upload or record an audio file."
    
    try:
        # Load model
        model = load_model()
        
        # Handle Gradio audio input
        sample_rate, audio_data = audio_input
        
        # Convert to float32 and normalize
        if audio_data.dtype == np.int16:
            audio_data = audio_data.astype(np.float32) / 32768.0
        elif audio_data.dtype == np.int32:
            audio_data = audio_data.astype(np.float32) / 2147483648.0
        
        # Convert stereo to mono
        if len(audio_data.shape) > 1:
            audio_data = np.mean(audio_data, axis=1)
        
        # Resample to 16kHz if needed
        if sample_rate != 16000:
            audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=16000)
        
        # Save to temp file for pipeline
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            import soundfile as sf
            sf.write(f.name, audio_data, 16000)
            temp_path = f.name
        
        # Run detection
        results = model(temp_path)
        
        # Clean up
        os.unlink(temp_path)
        
        # Parse results
        scores = {r["label"]: r["score"] for r in results}
        best = max(results, key=lambda x: x["score"])
        label = best["label"].lower()
        confidence = best["score"]
        
        # Determine classification with 98% threshold
        if "spoof" in label or "fake" in label:
            if confidence >= 0.98:
                classification = "🤖 AI-GENERATED"
                emoji = "🚨"
            else:
                classification = "👤 LIKELY HUMAN"
                emoji = "✅"
        else:
            classification = "👤 HUMAN"
            emoji = "✅"
        
        # Format output
        result = f"""
## {emoji} Detection Result

### Classification: {classification}
### Confidence: {confidence*100:.2f}%

---

### Raw Scores:
"""
        for label, score in scores.items():
            bar = "█" * int(score * 20) + "░" * (20 - int(score * 20))
            result += f"- **{label}**: {bar} {score*100:.1f}%\n"
        
        result += f"""
---

### Model Used:
`{MODEL_NAME}`

### Device:
{'🖥️ GPU (CUDA)' if torch.cuda.is_available() else '💻 CPU'}
"""
        
        return result
        
    except Exception as e:
        return f"❌ Error processing audio: {str(e)}"

# Create Gradio Interface
with gr.Blocks(title="VoiceGuard - AI Voice Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🛡️ VoiceGuard - AI Voice Detection
    
    Detect whether an audio sample is **AI-generated (deepfake)** or **genuine human speech**.
    
    ### Supported Languages
    🇮🇳 Tamil | English | Hindi | Malayalam | Telugu
    
    ---
    """)
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                label="Upload or Record Audio",
                sources=["upload", "microphone"],
                type="numpy"
            )
            
            detect_btn = gr.Button("🔍 Analyze Voice", variant="primary", size="lg")
            
            gr.Markdown("""
            ### Tips:
            - Upload MP3, WAV, or other audio formats
            - Record directly using your microphone
            - Minimum 1 second of audio recommended
            """)
        
        with gr.Column():
            output = gr.Markdown(label="Detection Result")
    
    detect_btn.click(
        fn=detect_deepfake,
        inputs=[audio_input],
        outputs=[output]
    )
    
    gr.Markdown("""
    ---
    
    ### About
    
    VoiceGuard uses the **HuBERT** model fine-tuned for deepfake audio detection.
    
    **Model**: `abhishtagatya/hubert-base-960h-itw-deepfake`
    
    **India AI Impact Summit Buildathon 2026**
    """)

# Launch
if __name__ == "__main__":
    # Pre-load model
    load_model()
    demo.launch()