#!/usr/bin/env python3
"""
Madverse Music API
AI Music Detection Service
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks, Header, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel, HttpUrl
import torch
import librosa
import tempfile
import os
import requests
from pathlib import Path
import time
from typing import Optional, Annotated, List
import uvicorn
import asyncio

# Initialize FastAPI app
app = FastAPI(
    title="Madverse Music API",
    description="AI-powered music detection API to identify AI-generated vs human-created music",
    version="1.0.0",
    docs_url="/",
    redoc_url="/docs"
)

# API Key Configuration
API_KEY = os.getenv("MADVERSE_API_KEY", "madverse-music-api-key-2024")  # Default key for demo

# Global model variable
model = None

async def verify_api_key(x_api_key: Annotated[str | None, Header()] = None):
    """Verify API key from header"""
    if x_api_key is None:
        raise HTTPException(
            status_code=401,
            detail="Missing API key. Please provide a valid X-API-Key header."
        )
    if x_api_key != API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API key. Please provide a valid X-API-Key header."
        )
    return x_api_key

class MusicAnalysisRequest(BaseModel):
    urls: List[HttpUrl]

def check_api_key_first(request: MusicAnalysisRequest, x_api_key: Annotated[str | None, Header()] = None):
    """Check API key before processing request"""
    if x_api_key is None:
        raise HTTPException(
            status_code=401,
            detail="Missing API key. Please provide a valid X-API-Key header."
        )
    if x_api_key != API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API key. Please provide a valid X-API-Key header."
        )
    return request
    
class FileAnalysisResult(BaseModel):
    url: str
    success: bool
    classification: Optional[str] = None  # "Real" or "Fake"
    confidence: Optional[float] = None    # 0.0 to 1.0
    probability: Optional[float] = None   # Raw sigmoid probability
    raw_score: Optional[float] = None     # Raw model output
    duration: Optional[float] = None      # Audio duration in seconds
    message: str
    processing_time: Optional[float] = None
    error: Optional[str] = None

class MusicAnalysisResponse(BaseModel):
    success: bool
    total_files: int
    successful_analyses: int
    failed_analyses: int
    results: List[FileAnalysisResult]
    total_processing_time: float
    message: str

class ErrorResponse(BaseModel):
    success: bool
    error: str
    message: str

@app.on_event("startup")
async def load_model():
    """Load the AI model on startup"""
    global model
    try:
        from sonics import HFAudioClassifier
        print("🔄 Loading Madverse Music AI model...")
        model = HFAudioClassifier.from_pretrained("awsaf49/sonics-spectttra-alpha-120s")
        model.eval()
        print("✅ Model loaded successfully!")
    except Exception as e:
        print(f"❌ Failed to load model: {e}")
        raise

def cleanup_file(file_path: str):
    """Background task to cleanup temporary files"""
    try:
        if os.path.exists(file_path):
            os.unlink(file_path)
    except:
        pass

def download_audio(url: str, max_size_mb: int = 100) -> str:
    """Download audio file from URL with size validation"""
    try:
        # Check if URL is accessible
        response = requests.head(str(url), timeout=10)
        
        # Check content size
        content_length = response.headers.get('Content-Length')
        if content_length and int(content_length) > max_size_mb * 1024 * 1024:
            raise HTTPException(
                status_code=413, 
                detail=f"File too large. Maximum size: {max_size_mb}MB"
            )
        
        # Download file
        response = requests.get(str(url), timeout=30, stream=True)
        response.raise_for_status()
        
        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.tmp') as tmp_file:
            downloaded_size = 0
            for chunk in response.iter_content(chunk_size=8192):
                downloaded_size += len(chunk)
                if downloaded_size > max_size_mb * 1024 * 1024:
                    os.unlink(tmp_file.name)
                    raise HTTPException(
                        status_code=413,
                        detail=f"File too large. Maximum size: {max_size_mb}MB"
                    )
                tmp_file.write(chunk)
            
            return tmp_file.name
            
    except requests.exceptions.RequestException as e:
        raise HTTPException(
            status_code=400,
            detail=f"Failed to download audio: {str(e)}"
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error downloading file: {str(e)}"
        )

def classify_audio(file_path: str) -> dict:
    """Classify audio file using the AI model"""
    try:
        # Load audio (model uses 16kHz sample rate)
        audio, sr = librosa.load(file_path, sr=16000)
        
        # Convert to tensor and add batch dimension
        audio_tensor = torch.FloatTensor(audio).unsqueeze(0)
        
        # Get prediction
        with torch.no_grad():
            output = model(audio_tensor)
            
            # Convert logit to probability using sigmoid
            prob = torch.sigmoid(output).item()
            
            # Classify: prob < 0.5 = Real, prob >= 0.5 = Fake
            if prob < 0.5:
                classification = "Real"
                confidence = (1 - prob) * 2  # Convert to 0-1 scale
            else:
                classification = "Fake"
                confidence = (prob - 0.5) * 2  # Convert to 0-1 scale
        
        return {
            "classification": classification,
            "confidence": min(confidence, 1.0),  # Cap at 1.0
            "probability": prob,
            "raw_score": output.item(),
            "duration": len(audio) / sr
        }
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error analyzing audio: {str(e)}"
        )

async def process_single_url(url: str) -> FileAnalysisResult:
    """Process a single URL and return result"""
    start_time = time.time()
    
    try:
        # Download audio file
        temp_file = download_audio(url)
        
        # Classify audio
        result = classify_audio(temp_file)
        
        # Calculate processing time
        processing_time = time.time() - start_time
        
        # Cleanup file in background
        try:
            os.unlink(temp_file)
        except:
            pass
        
        # Prepare response
        emoji = "🎤" if result["classification"] == "Real" else "🤖"
        message = f'{emoji} Detected as {result["classification"].lower()} music'
        
        return FileAnalysisResult(
            url=str(url),
            success=True,
            classification=result["classification"],
            confidence=result["confidence"],
            probability=result["probability"],
            raw_score=result["raw_score"],
            duration=result["duration"],
            message=message,
            processing_time=processing_time
        )
        
    except Exception as e:
        processing_time = time.time() - start_time
        error_msg = str(e)
        
        return FileAnalysisResult(
            url=str(url),
            success=False,
            message=f"❌ Failed to process: {error_msg}",
            processing_time=processing_time,
            error=error_msg
        )

@app.post("/analyze", response_model=MusicAnalysisResponse)
async def analyze_music(
    request: MusicAnalysisRequest = Depends(check_api_key_first)
):
    """
    Analyze music from URL(s) to detect if it's AI-generated or human-created
    
    - **urls**: Array of direct URLs to audio files (MP3, WAV, FLAC, M4A, OGG)
    - Returns classification results for each file
    - Processes files concurrently for better performance when multiple URLs provided
    """
    start_time = time.time()
    
    if not model:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded. Please try again later."
        )
    
    if len(request.urls) > 50:  # Limit processing
        raise HTTPException(
            status_code=400,
            detail="Too many URLs. Maximum 50 files per request."
        )
    
    if len(request.urls) == 0:
        raise HTTPException(
            status_code=400,
            detail="At least one URL is required."
        )
    
    try:
        # Process all URLs concurrently with limited concurrency
        semaphore = asyncio.Semaphore(5)  # Limit to 5 concurrent downloads
        
        async def process_with_semaphore(url):
            async with semaphore:
                return await process_single_url(str(url))
        
        # Create tasks for all URLs
        tasks = [process_with_semaphore(url) for url in request.urls]
        
        # Wait for all tasks to complete
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Process results and handle any exceptions
        processed_results = []
        successful_count = 0
        failed_count = 0
        
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                # Handle exception case
                processed_results.append(FileAnalysisResult(
                    url=str(request.urls[i]),
                    success=False,
                    message=f"❌ Processing failed: {str(result)}",
                    error=str(result)
                ))
                failed_count += 1
            else:
                processed_results.append(result)
                if result.success:
                    successful_count += 1
                else:
                    failed_count += 1
        
        # Calculate total processing time
        total_processing_time = time.time() - start_time
        
        # Prepare summary message
        total_files = len(request.urls)
        if total_files == 1:
            # Single file message
            if successful_count == 1:
                message = processed_results[0].message
            else:
                message = processed_results[0].message
        else:
            # Multiple files message
            if successful_count == total_files:
                message = f"✅ Successfully analyzed all {total_files} files"
            elif successful_count > 0:
                message = f"⚠️ Analyzed {successful_count}/{total_files} files successfully"
            else:
                message = f"❌ Failed to analyze any files"
        
        return MusicAnalysisResponse(
            success=successful_count > 0,
            total_files=total_files,
            successful_analyses=successful_count,
            failed_analyses=failed_count,
            results=processed_results,
            total_processing_time=total_processing_time,
            message=message
        )
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Internal server error during processing: {str(e)}"
        )

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": model is not None,
        "service": "Madverse Music API"
    }

@app.get("/info")
async def get_info():
    """Get API information"""
    return {
        "name": "Madverse Music API",
        "version": "1.0.0",
        "description": "AI-powered music detection to identify AI-generated vs human-created music",
        "model": "SpecTTTra-α (120s)",
        "accuracy": {
            "f1_score": 0.97,
            "sensitivity": 0.96,
            "specificity": 0.99
        },
        "supported_formats": ["MP3", "WAV", "FLAC", "M4A", "OGG"],
        "max_file_size": "100MB",
        "max_duration": "120 seconds",
        "authentication": {
            "required": True,
            "type": "API Key",
            "header": "X-API-Key",
            "example": "X-API-Key: your-api-key-here"
        },
        "usage": {
            "curl_example": "curl -X POST 'http://localhost:8000/analyze' -H 'X-API-Key: your-api-key' -H 'Content-Type: application/json' -d '{\"url\":\"https://example.com/song.mp3\"}'"
        }
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)