main.py

"""
Production FastAPI Server with S3 Support
Works on: Hugging Face Space (testing) → AWS (production)
"""

import os
import tempfile
from pathlib import Path
from typing import Optional
import time
import uuid

import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException, status, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles

from kid_coach_pipeline import EnhancedPublicSpeakingCoach

# Try to import boto3 (for AWS S3)
try:
    import boto3
    S3_AVAILABLE = True
except ImportError:
    S3_AVAILABLE = False
    print("⚠️  boto3 not available - S3 uploads disabled")

# ================= CONFIGURATION =================

app = FastAPI(
    title="Aurator Speech Coach API",
    description="AI-powered speech analysis with personalized coaching",
    version="4.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Audio output directory (local fallback)
AUDIO_OUTPUT_DIR = "/tmp/audio_outputs"
os.makedirs(AUDIO_OUTPUT_DIR, exist_ok=True)

# Mount for local testing
app.mount("/audio", StaticFiles(directory=AUDIO_OUTPUT_DIR), name="audio")

# AWS S3 Configuration (optional - for production)
USE_S3 = os.getenv("USE_S3", "false").lower() == "true" and S3_AVAILABLE
S3_BUCKET = os.getenv("S3_BUCKET_NAME", "aurator-audio-outputs")
S3_REGION = os.getenv("AWS_REGION", "us-east-1")

if USE_S3:
    s3_client = boto3.client('s3', region_name=S3_REGION)
    print(f"✅ S3 enabled: {S3_BUCKET}")
else:
    print("📁 Using local file storage")

coach_engine: Optional[EnhancedPublicSpeakingCoach] = None

SUPPORTED_FORMATS = {'.wav', '.mp3', '.m4a', '.flac', '.ogg', '.aac', '.webm'}
MAX_FILE_SIZE = 50 * 1024 * 1024


# ================= S3 HELPERS =================

def upload_to_s3(file_path: str, file_type: str) -> str:
    """Upload file to S3 and return public URL"""
    if not USE_S3:
        # Return local URL for HF testing
        filename = os.path.basename(file_path)
        return f"/audio/{filename}"
    
    try:
        # Generate S3 key with date structure
        from datetime import datetime
        now = datetime.now()
        file_uuid = str(uuid.uuid4())
        s3_key = f"{now.year}/{now.month:02d}/{now.day:02d}/{file_type}_{file_uuid}.wav"
        
        # Upload to S3
        s3_client.upload_file(
            file_path,
            S3_BUCKET,
            s3_key,
            ExtraArgs={'ContentType': 'audio/wav', 'ACL': 'public-read'}
        )
        
        # Return public URL
        return f"https://{S3_BUCKET}.s3.{S3_REGION}.amazonaws.com/{s3_key}"
    
    except Exception as e:
        print(f"❌ S3 upload failed: {e}")
        # Fallback to local URL
        filename = os.path.basename(file_path)
        return f"/audio/{filename}"


# ================= STARTUP =================

@app.on_event("startup")
async def startup_event():
    global coach_engine
    
    print("\n" + "="*60)
    print("🚀 AURATOR SPEECH COACH API")
    print("="*60)
    
    try:
        openai_key = os.getenv("OPENAI_API_KEY")
        
        print("\n📦 Loading models...")
        coach_engine = EnhancedPublicSpeakingCoach(
            whisper_model_size="base",
            enable_tts=True,
            openai_api_key=openai_key
        )
        print("✅ Engine ready!")
        print(f"   Storage: {'S3' if USE_S3 else 'Local'}")
        print(f"   OpenAI: {'Enabled' if openai_key else 'Fallback mode'}")
        print("\n" + "="*60 + "\n")
        
    except Exception as e:
        print(f"\n❌ STARTUP FAILED: {e}")
        coach_engine = None


# ================= ENDPOINTS =================

@app.get("/")
async def root():
    """API info"""
    return {
        "service": "Aurator Speech Coach API",
        "version": "4.0.0",
        "status": "online" if coach_engine else "degraded",
        "storage": "s3" if USE_S3 else "local",
        "endpoints": {
            "analyze": "POST /api/analyze",
            "health": "GET /api/health"
        }
    }


@app.get("/api/health")
async def health_check():
    """Health check for AWS load balancer"""
    return {
        "status": "healthy" if coach_engine else "unhealthy",
        "engine_loaded": coach_engine is not None,
        "timestamp": time.time()
    }


@app.post("/api/analyze")
async def analyze_speech(
    audio_file: UploadFile = File(...),
    avatar_gender: str = Form('male')
):
    """
    Main endpoint: Analyze speech and return results
    
    Request:
    - audio_file: Audio file (wav/mp3/m4a/flac/ogg/aac/webm)
    - avatar_gender: "male" or "female" (default: male)
    
    Response:
    {
        "success": true,
        "data": {
            "overall_score": 8.6,
            "transcription": {...},
            "analysis": {...},
            "coaching": {
                "tips": [...],
                "tips_audio_url": "https://...",
                "improved_audio_url": "https://..."
            }
        },
        "processing_time_ms": 3250,
        "timestamp": "2025-12-16T..."
    }
    """
    start_time = time.time()
    
    # Validate engine
    if coach_engine is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Engine not initialized"
        )
    
    # Validate file
    if not audio_file or not audio_file.filename:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="No audio file provided"
        )
    
    file_ext = Path(audio_file.filename).suffix.lower()
    if file_ext not in SUPPORTED_FORMATS:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Unsupported format: {file_ext}"
        )
    
    temp_file = None
    
    try:
        # Save uploaded file temporarily
        content = await audio_file.read()
        
        if len(content) > MAX_FILE_SIZE:
            raise HTTPException(
                status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
                detail="File too large (max 50MB)"
            )
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as temp:
            temp.write(content)
            temp_file = temp.name
        
        print(f"\n🎤 Analyzing: {audio_file.filename} ({len(content)/1024:.1f} KB)")
        
        # Run analysis
        result = coach_engine.analyze_speech(
            temp_file,
            output_dir=AUDIO_OUTPUT_DIR,
            enable_tts=True,
            avatar_gender=avatar_gender
        )
        
        if "error" in result:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=result["error"]
            )
        
        # Upload audio files to S3 (if enabled) or use local URLs
        tips_audio_path = None
        improved_audio_path = None
        
        if result.get("tips_audio_url"):
            local_path = os.path.join(AUDIO_OUTPUT_DIR, os.path.basename(result["tips_audio_url"]))
            if os.path.exists(local_path):
                tips_audio_url = upload_to_s3(local_path, "tips")
                result["tips_audio_url"] = tips_audio_url
        
        if result.get("avatar_audio_url"):
            local_path = os.path.join(AUDIO_OUTPUT_DIR, os.path.basename(result["avatar_audio_url"]))
            if os.path.exists(local_path):
                improved_audio_url = upload_to_s3(local_path, "improved")
                result["avatar_audio_url"] = improved_audio_url
        
        processing_time = int((time.time() - start_time) * 1000)
        
        print(f"✅ Complete in {processing_time}ms")
        
        # Format response for React Native
        response = {
            "success": True,
            "data": {
                "overall_score": result.get("overall_score", 0),
                "duration_seconds": result.get("duration_seconds", 0),
                "word_count": result.get("word_count", 0),
                
                "transcription": {
                    "text": result.get("transcription", ""),
                    "improved_text": result.get("improved_transcript", "")
                },
                
                "analysis": {
                    "pacing": result.get("pacing", {}),
                    "prosody": result.get("prosody", {}),
                    "fillers": result.get("filler_words", {}),
                    "silences": result.get("silence_detection", {}),
                    "sentiment": result.get("sentiment_analysis", {}),
                    "vocabulary": result.get("vocabulary", {}),
                    "flow": result.get("logical_flow", {}),
                    "coherence": result.get("coherence", {}),
                    "persuasion": result.get("persuasion", {})
                },
                
                "coaching": {
                    "tips": result.get("personalized_tips", []),
                    "tips_audio_url": result.get("tips_audio_url"),
                    "improved_audio_url": result.get("avatar_audio_url")
                }
            },
            "processing_time_ms": processing_time,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        }
        
        return JSONResponse(content=response)
    
    except HTTPException:
        raise
    
    except Exception as e:
        import traceback
        traceback.print_exc()
        
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Analysis failed: {str(e)}"
        )
    
    finally:
        # Cleanup
        if temp_file and os.path.exists(temp_file):
            try:
                os.remove(temp_file)
            except:
                pass


@app.get("/audio/{filename}")
async def serve_audio(filename: str):
    """Serve audio files (for local/HF testing)"""
    file_path = os.path.join(AUDIO_OUTPUT_DIR, filename)
    
    if not os.path.exists(file_path):
        raise HTTPException(404, "Audio file not found")
    
    return FileResponse(file_path, media_type="audio/wav")


# ================= RUN =================

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