import subprocess
import sys
import os
import urllib.request
import tarfile
import shutil
from pathlib import Path

# Install FFmpeg if not available
def setup_ffmpeg():
    """Set up FFmpeg for Hugging Face Spaces using pure Python"""
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        print("✅ FFmpeg already installed")
        return
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("📦 Installing FFmpeg...")
    
    try:
        local_bin = Path.home() / ".local" / "bin"
        local_bin.mkdir(parents=True, exist_ok=True)
        
        ffmpeg_url = "https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-amd64-static.tar.xz"
        temp_file = "/tmp/ffmpeg.tar.xz"
        
        print("⬇️ Downloading FFmpeg...")
        urllib.request.urlretrieve(ffmpeg_url, temp_file)
        
        print("📦 Extracting FFmpeg...")
        with tarfile.open(temp_file, 'r:xz') as tar:
            tar.extractall("/tmp/ffmpeg_extract")
        
        extract_dir = Path("/tmp/ffmpeg_extract")
        ffmpeg_dirs = list(extract_dir.glob("ffmpeg-*-amd64-static"))
        
        if ffmpeg_dirs:
            source_dir = ffmpeg_dirs[0]
            shutil.copy(source_dir / "ffmpeg", local_bin / "ffmpeg")
            shutil.copy(source_dir / "ffprobe", local_bin / "ffprobe")
            os.chmod(local_bin / "ffmpeg", 0o755)
            os.chmod(local_bin / "ffprobe", 0o755)
            os.environ["PATH"] = f"{local_bin}:{os.environ.get('PATH', '')}"
            os.remove(temp_file)
            shutil.rmtree("/tmp/ffmpeg_extract")
            print("✅ FFmpeg installed successfully!")
        else:
            raise Exception("Could not find extracted FFmpeg directory")
            
    except Exception as e:
        print(f"❌ FFmpeg installation failed: {e}")

setup_ffmpeg()

from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uuid
from pipeline import enhance_audio, cleanup_temp_files

app = FastAPI(title="AI Voice Enhancer API", version="1.0.0")

# Enable CORS for frontend access
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify your frontend domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Create temp directory
os.makedirs("temp", exist_ok=True)

@app.get("/")
def read_root():
    return {
        "message": "AI Voice Enhancer API",
        "version": "1.0.0",
        "authentication": "required" if os.getenv("API_KEY") else "disabled",
        "endpoints": {
            "/enhance": "POST - Enhance audio file",
            "/health": "GET - Health check"
        }
    }

@app.get("/health")
def health_check():
    return {
        "status": "healthy", 
        "ffmpeg": "installed",
        "auth_enabled": bool(os.getenv("API_KEY"))
    }

@app.post("/enhance")
async def enhance_audio_endpoint(
    file: UploadFile = File(...),
    speech_level: float = Form(80.0),
    background_level: float = Form(85.0),
    x_api_key: str = None
):
    """
    Enhance uploaded audio file.
    
    Parameters:
    - file: Audio file (WAV, MP3, M4A, FLAC, OGG, AAC)
    - speech_level: Speech enhancement level (0-100)
    - background_level: Background noise reduction level (0-100)
    - X-API-Key: Optional API key (header) - required if API_KEY env var is set
    
    Returns:
    - Enhanced audio file
    """
    
    # Check API key if authentication is enabled
    api_key = os.getenv("API_KEY")
    if api_key:
        # Get key from header
        from fastapi import Request, Header
        from typing import Optional
        
        # X-API-Key can come from header or form data
        if not x_api_key:
            raise HTTPException(
                status_code=401,
                detail="API key required. Set X-API-Key header or pass api_key parameter."
            )
        
        if x_api_key != api_key:
            raise HTTPException(
                status_code=403,
                detail="Invalid API key"
            )
    
    # Validate file extension
    valid_extensions = {'.wav', '.mp3', '.m4a', '.flac', '.ogg', '.aac'}
    file_ext = Path(file.filename).suffix.lower()
    if file_ext not in valid_extensions:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file format. Allowed: {', '.join(valid_extensions)}"
        )
    
    # Save uploaded file
    session_id = str(uuid.uuid4())[:8]
    input_path = f"temp/{session_id}_input{file_ext}"
    
    try:
        with open(input_path, "wb") as f:
            content = await file.read()
            f.write(content)
        
        # Validate parameters
        speech_enhancement = max(0.0, min(100.0, speech_level)) / 100.0
        noise_reduction = max(0.0, min(100.0, background_level)) / 100.0
        
        # Enhance audio
        enhanced_path, status = enhance_audio(
            input_path,
            speech_enhancement=speech_enhancement,
            noise_reduction=noise_reduction
        )
        
        if not enhanced_path:
            # Clean up input file
            if os.path.exists(input_path):
                os.remove(input_path)
            raise HTTPException(status_code=500, detail=status)
        
        # Return enhanced file
        response = FileResponse(
            enhanced_path,
            media_type="audio/wav",
            filename=f"enhanced_{file.filename}",
            headers={
                "X-Processing-Status": status,
                "X-Speech-Level": str(speech_level),
                "X-Background-Level": str(background_level)
            }
        )
        
        # Clean up input file
        if os.path.exists(input_path):
            os.remove(input_path)
        
        return response
        
    except HTTPException:
        raise
    except Exception as e:
        # Clean up on error
        if os.path.exists(input_path):
            os.remove(input_path)
        cleanup_temp_files()
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/cleanup")
def cleanup_endpoint():
    """Clean up old temporary files"""
    cleanup_temp_files()
    return {"status": "cleaned"}

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