main.py

from __future__ import annotations

import io
import os
import tempfile
import time
import warnings
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from contextlib import asynccontextmanager
import re

import librosa
import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor

# Suppress audio processing warnings
warnings.filterwarnings('ignore', category=UserWarning, module='librosa')
warnings.filterwarnings('ignore', category=FutureWarning, module='librosa')

# ML models and alignment
try:
    from phoneme_features import extract_phoneme_features, get_feature_names
    from phoneme_scorer import PronunciationScorer
    ML_SCORER_AVAILABLE = True
except ImportError:
    ML_SCORER_AVAILABLE = False
    print("[STARTUP] ML scorer modules not available, using rule-based scoring")

# Character-level aligner (primary method - production ready)
from character_aligner import (
    align_phonemes_character_level, 
    align_phonemes_ctc,
    calculate_pronunciation_accuracy
)

# ============================================================================
# GLOBAL MODEL CACHE (Loaded at startup for low latency)
# ============================================================================
ASR_MODEL: Optional[Wav2Vec2ForCTC] = None
ASR_PROCESSOR: Optional[Wav2Vec2Processor] = None
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Phoneme scoring model (ML-based - PyTorch)
PHONEME_SCORER = None
USE_ML_SCORER = False


# ============================================================================
# LIFESPAN: Modern FastAPI startup/shutdown handler
# ============================================================================
@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Modern FastAPI lifespan context manager for startup/shutdown events.
    Replaces deprecated @app.on_event("startup") and @app.on_event("shutdown").
    """
    # STARTUP
    global ASR_MODEL, ASR_PROCESSOR, PHONEME_SCORER, USE_ML_SCORER
    
    print(f"[STARTUP] Loading Arabic ASR model on device: {DEVICE}")
    model_name = "jonatasgrosman/wav2vec2-large-xlsr-53-arabic"
    
    try:
        ASR_PROCESSOR = Wav2Vec2Processor.from_pretrained(model_name)
        ASR_MODEL = Wav2Vec2ForCTC.from_pretrained(model_name)
        ASR_MODEL.to(DEVICE)
        ASR_MODEL.eval()  # Set to evaluation mode
        print(f"[STARTUP] ✓ ASR model loaded successfully ({sum(p.numel() for p in ASR_MODEL.parameters()) / 1e6:.1f}M parameters)")
    except Exception as e:
        print(f"[STARTUP] ✗ Failed to load ASR model: {e}")
        raise RuntimeError("ASR model is required for the service to function")
    
    # Load phoneme scoring model (PyTorch-based - optional)
    if ML_SCORER_AVAILABLE:
        try:
            PHONEME_SCORER = PronunciationScorer(models_dir="models")
            if PHONEME_SCORER.load_model():
                USE_ML_SCORER = True
                print(f"[STARTUP] ✓ ML phoneme scorer loaded successfully")
                print(f"[STARTUP]   Model type: {PHONEME_SCORER.model_type}")
                print(f"[STARTUP]   Device: {PHONEME_SCORER.device}")
            else:
                print(f"[STARTUP] ⚠ ML scorer not available, using rule-based scoring")
                PHONEME_SCORER = None
        except Exception as e:
            print(f"[STARTUP] ⚠ Error initializing ML scorer: {e}")
            print("[STARTUP] Using rule-based scoring")
            PHONEME_SCORER = None
    
    print(f"[STARTUP] ✓ Using CTC-based phoneme alignment (frame-accurate)")
    print(f"[STARTUP] ✓ Service ready on port 8000")
    print(f"[STARTUP]   Alignment: CTC segmentation with fallback")
    print(f"[STARTUP]   Scoring: {'ML-based' if USE_ML_SCORER else 'Rule-based'}")
    print(f"[STARTUP]   Device: {DEVICE}")
    
    yield  # Application is running
    
    # SHUTDOWN (cleanup if needed)
    print("[SHUTDOWN] Cleaning up resources...")
    if ASR_MODEL is not None:
        del ASR_MODEL
    if ASR_PROCESSOR is not None:
        del ASR_PROCESSOR
    if PHONEME_SCORER is not None:
        del PHONEME_SCORER
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    print("[SHUTDOWN] ✓ Resources cleaned up")


app = FastAPI(
    title="Arabic Speech Analysis - ASR + CTC Alignment",
    lifespan=lifespan  # Modern lifespan handler
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def root():
    return {
        "status": "ok",
        "service": "SmartArabicSpeechTherapy AI",
        "docs": "/docs",
        "endpoints": ["/analyze", "/health"],
    }


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "device": DEVICE,
        "asrModelLoaded": ASR_MODEL is not None,
        "mlScorerEnabled": USE_ML_SCORER,
    }


# ============================================================================
# AUDIO PREPROCESSING
# ============================================================================
def convert_audio_to_wav(audio_bytes: bytes, target_sr: int = 16000, filename: str = None) -> np.ndarray:
    """
    Convert any audio format (WebM, MP3, WAV, etc.) to 16kHz mono WAV.
    
    Args:
        audio_bytes: Raw audio file bytes
        target_sr: Target sample rate (16kHz for wav2vec2)
        filename: Original filename (optional, helps with format detection)
    
    Returns:
        numpy array of audio samples
    """
    if not audio_bytes or len(audio_bytes) == 0:
        raise ValueError("No audio data received")
    
    # Determine file extension
    if filename:
        ext = os.path.splitext(filename)[1].lower()
        if not ext:
            ext = '.wav'
    else:
        ext = '.wav'
    
    # Use temporary file for reliable format support
    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp_file:
        tmp_file.write(audio_bytes)
        tmp_path = tmp_file.name
    
    try:
        # Try soundfile first (faster, more reliable for standard formats)
        try:
            audio_array, sr = sf.read(tmp_path, dtype='float32')
            # Convert to mono if stereo
            if len(audio_array.shape) > 1:
                audio_array = audio_array.mean(axis=1)
            # Resample if needed
            if sr != target_sr:
                audio_array = librosa.resample(audio_array, orig_sr=sr, target_sr=target_sr)
        except Exception:
            # Fallback to librosa (supports more formats via ffmpeg/audioread)
            audio_array, sr = librosa.load(
                tmp_path,
                sr=target_sr,
                mono=True,
                res_type='kaiser_best'
            )
        
        # Validate audio was loaded
        if audio_array is None or len(audio_array) == 0:
            raise ValueError("Audio file is empty or unreadable")
        
        # Normalize audio to [-1, 1]
        max_val = np.abs(audio_array).max()
        if max_val > 0:
            if max_val > 1.0:
                audio_array = audio_array / max_val
        else:
            raise ValueError("Audio contains only silence")
        
        return audio_array
    
    except Exception as e:
        raise ValueError(f"Failed to convert audio file '{filename or 'unknown'}': {str(e)}")
    
    finally:
        # Clean up temporary file
        if os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except:
                pass


def validate_audio_duration(audio_array: np.ndarray, sr: int = 16000) -> bool:
    """
    Validate audio duration is within acceptable range.
    
    Args:
        audio_array: Audio samples
        sr: Sample rate
    
    Returns:
        True if valid, False otherwise
    """
    duration = len(audio_array) / sr
    
    # Speech therapy typically 1-15 seconds per utterance
    if duration < 0.5:
        raise ValueError(f"Audio too short: {duration:.2f}s (minimum 0.5s)")
    if duration > 20:
        raise ValueError(f"Audio too long: {duration:.2f}s (maximum 20s)")
    
    return True


# ============================================================================
# ARABIC ASR (wav2vec2)
# ============================================================================
def transcribe_audio(audio_array: np.ndarray, sr: int = 16000, return_ctc_data: bool = False):
    """
    Transcribe audio using wav2vec2 Arabic model.
    
    Args:
        audio_array: Audio samples (16kHz mono)
        sr: Sample rate
        return_ctc_data: If True, return logits and predicted_ids for CTC alignment
    
    Returns:
        If return_ctc_data=False: Tuple of (transcript, confidence_scores_per_frame)
        If return_ctc_data=True: Tuple of (transcript, confidence_scores, logits, predicted_ids)
    """
    if ASR_MODEL is None or ASR_PROCESSOR is None:
        raise RuntimeError("ASR model not loaded")
    
    try:
        # Prepare input for model
        inputs = ASR_PROCESSOR(
            audio_array,
            sampling_rate=sr,
            return_tensors="pt",
            padding=True
        )
        
        # Move to same device as model
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
        
        # Run inference
        with torch.no_grad():
            logits = ASR_MODEL(**inputs).logits
        
        # Get predicted IDs
        predicted_ids = torch.argmax(logits, dim=-1)
        
        # Decode to text
        transcription = ASR_PROCESSOR.batch_decode(predicted_ids)[0]
        
        # Get confidence scores (softmax of logits)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        confidence_scores = torch.max(probs, dim=-1)[0].cpu().numpy()[0]
        
        if return_ctc_data:
            return transcription.strip(), confidence_scores, logits, predicted_ids
        else:
            return transcription.strip(), confidence_scores
    
    except Exception as e:
        raise RuntimeError(f"ASR transcription failed: {str(e)}")


# ============================================================================
# FORCED ALIGNMENT (Simplified Character-Level)
# ============================================================================
ARABIC_DIACRITICS = re.compile(r"[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06ED]")


def normalize_arabic(text: str) -> str:
    """Normalize Arabic text by removing diacritics and tatweel."""
    if not text:
        return ""
    text = ARABIC_DIACRITICS.sub("", text)
    text = text.replace("\u0640", "")  # tatweel
    text = text.replace("ة", "ه")  # normalize ta marbuta
    text = text.replace("ا", "ه")  # normalize alef to ha
    return "".join(text.split()).strip()

def align_phonemes_simple(
    audio_array: np.ndarray,
    expected_text: str,
    transcribed_text: str,
    confidence_scores: np.ndarray,
    sr: int = 16000
) -> List[Dict[str, Any]]:
    """
    DEPRECATED: Legacy character-level alignment for Arabic text.
    
    Kept for backward compatibility with old tests.
    For production, use character_aligner.align_phonemes_character_level() instead.
    
    Args:
        audio_array: Audio samples
        expected_text: What child should say
        transcribed_text: What ASR heard
        confidence_scores: Frame-level confidence from ASR
        sr: Sample rate
    
    Returns:
        List of phoneme dictionaries with alignment info
    """
    # Clean texts (remove diacritics, tatweel, whitespace)
    expected_clean = normalize_arabic(expected_text)
    transcribed_clean = normalize_arabic(transcribed_text)
    
    phoneme_results = []
    audio_duration = len(audio_array) / sr
    
    # Simple character-by-character comparison
    max_len = max(len(expected_clean), len(transcribed_clean))
    
    for i, expected_char in enumerate(expected_clean):
        # Calculate timing (evenly distribute across audio duration)
        start_time = (i / max_len) * audio_duration if max_len > 0 else 0
        end_time = ((i + 1) / max_len) * audio_duration if max_len > 0 else audio_duration
        duration = end_time - start_time
        
        # Get corresponding transcribed character
        transcribed_char = transcribed_clean[i] if i < len(transcribed_clean) else None
        
        # Character match check
        is_correct = (expected_char == transcribed_char)
        
        # Calculate confidence (use mean of relevant frames)
        frame_start = int((start_time / audio_duration) * len(confidence_scores))
        frame_end = int((end_time / audio_duration) * len(confidence_scores))
        frame_start = max(0, min(frame_start, len(confidence_scores) - 1))
        frame_end = max(frame_start + 1, min(frame_end, len(confidence_scores)))
        
        char_confidence = float(np.mean(confidence_scores[frame_start:frame_end]))
        
        # Detect error type
        error_type = None
        if not is_correct:
            if transcribed_char is None:
                error_type = "deletion"
            elif transcribed_char != expected_char:
                error_type = "substitution"
        
        # Apply duration penalty (very short/long = suspicious)
        duration_penalty = 0.0
        expected_duration = 0.08  # ~80ms per character (rough estimate)
        if duration < expected_duration * 0.5:
            duration_penalty = 0.15  # Too fast
        elif duration > expected_duration * 2.0:
            duration_penalty = 0.10  # Too slow
        
        final_confidence = max(0.0, char_confidence - duration_penalty)
        
        phoneme_results.append({
            "symbol": expected_char,
            "expected": is_correct,
            "confidence": round(final_confidence, 3),
            "errorType": error_type,
            "duration": round(duration, 3),
            "timestamp": round(start_time, 3)
        })
    
    # Handle insertions (extra characters in transcription)
    if len(transcribed_clean) > len(expected_clean):
        for i in range(len(expected_clean), len(transcribed_clean)):
            phoneme_results.append({
                "symbol": transcribed_clean[i],
                "expected": False,
                "confidence": 0.3,
                "errorType": "insertion",
                "duration": 0.05,
                "timestamp": audio_duration * (i / len(transcribed_clean))
            })
    
    return phoneme_results


# ============================================================================
# PHONEME SCORING ALGORITHM
# ============================================================================
def score_phonemes_ml(phoneme_results: List[Dict[str, Any]], 
                     audio_array: np.ndarray) -> List[Dict[str, Any]]:
    """
    Score phonemes using trained ML model.
    
    Args:
        phoneme_results: List of PyTorch ML model.
    
    Args:
        phoneme_results: List of phoneme dictionaries from alignment
        audio_array: Full audio samples for feature extraction
        
    Returns:
        Phoneme results with ML-based confidence scores
    """
    if not USE_ML_SCORER or PHONEME_SCORER is None:
        return phoneme_results
    
    try:
        # Extract features and score using batch processing
        features_list = []
        
        for phoneme_info in phoneme_results:
            # Extract features for this phoneme
            features = extract_phoneme_features(audio_array, phoneme_info, sr=16000, n_mfcc=13)
            features_list.append(features)
        
        # Batch score all phonemes
        results = PHONEME_SCORER.score_batch(features_list)
        
        # Update phoneme results with ML scores
        updated_results = []
        for phoneme_info, (ml_score, feedback) in zip(phoneme_results, results):
            # Get alignment error detection
            error_type = phoneme_info.get('errorType', None)
            alignment_expected = phoneme_info.get('expected', True)
            
            # PRIORITY: Alignment error detection takes precedence over ML
            # If alignment detected substitution/deletion/insertion, it's WRONG regardless of ML
            if error_type in ['substitution', 'deletion', 'insertion']:
                # Alignment detected error - override ML feedback
                phoneme_info['expected'] = False
                phoneme_info['ml_feedback'] = 'incorrect'
                # Use lower confidence for errors
                phoneme_info['confidence'] = min(round(ml_score, 3), 0.6)
                phoneme_info['ml_scored'] = True
            elif not alignment_expected:
                # Alignment says unexpected, respect it
                phoneme_info['expected'] = False
                phoneme_info['ml_feedback'] = feedback if feedback == 'incorrect' else 'uncertain'
                phoneme_info['confidence'] = round(ml_score, 3)
                phoneme_info['ml_scored'] = True
            else:
                # No error detected by alignment, use ML assessment
                phoneme_info['confidence'] = round(ml_score, 3)
                phoneme_info['ml_scored'] = True
                phoneme_info['ml_feedback'] = feedback
                
                # Only mark as expected if BOTH alignment AND ML agree
                if feedback == 'correct':
                    phoneme_info['expected'] = True
                elif feedback == 'incorrect':
                    phoneme_info['expected'] = False
                # Leave 'uncertain' as-is from alignment
            
            updated_results.append(phoneme_info)
        
        return updated_results
    
    except Exception as e:
        print(f"[WARNING] ML scoring failed: {e}, using original scores")
        import traceback
        traceback.print_exc()
        return phoneme_results


def calculate_overall_score(phoneme_results: List[Dict[str, Any]]) -> int:
    """
    Calculate overall pronunciation score (0-100) from phoneme-level results.
    
    Args:
        phoneme_results: List of phoneme dictionaries with confidence scores
    
    Returns:
        Integer score 0-100
    """
    if not phoneme_results:
        return 0
    
    # Weight by confidence and correctness
    total_score = 0.0
    total_weight = 0.0
    
    for phoneme in phoneme_results:
        confidence = phoneme.get("confidence", 0.0)
        is_expected = phoneme.get("expected", False)
        error_type = phoneme.get("errorType")
        
        # Calculate phoneme score
        if is_expected:
            phoneme_score = confidence * 100
        else:
            # Penalize errors - be strict for therapeutic feedback
            if error_type == "deletion":
                phoneme_score = 0  # Missing phoneme = worst
            elif error_type == "substitution":
                phoneme_score = confidence * 25  # Wrong phoneme, strong penalty
            elif error_type == "insertion":
                phoneme_score = 15  # Extra phoneme, penalty
            else:
                phoneme_score = confidence * 40  # Unknown error
        
        total_score += phoneme_score
        total_weight += 1.0
    
    # Average score
    avg_score = total_score / total_weight if total_weight > 0 else 0
    
    # Round to integer
    return max(0, min(100, int(round(avg_score))))


def _tone_from_score(score: int) -> str:
    """Map score to feedback tone."""
    if score >= 90:
        return "gold"
    if score >= 75:
        return "blue"
    return "gray"


# ============================================================================
# FALLBACK MODE (if ASR fails)
# ============================================================================
def analyze_fallback(expected_text: str, audio_bytes: bytes) -> Dict[str, Any]:
    """Fallback to deterministic scoring if ASR pipeline fails."""
    clean = "".join(expected_text.split())[:10]
    letters = list(clean) if clean else ["?"]
    
    base = 70 + min(20, len(letters) * 2)
    score = max(0, min(100, base))
    
    phonemes = []
    for idx, ch in enumerate(letters):
        ok = (idx % 3) != 0  # 2/3 correct
        phonemes.append({
            "symbol": ch,
            "expected": ok,
            "confidence": 0.80 if ok else 0.50,
            "errorType": None if ok else "substitution",
            "duration": 0.1,
            "timestamp": idx * 0.1
        })
    
    return {
        "score": score,
        "feedbackTone": _tone_from_score(score),
        "phonemes": phonemes,
        "meta": {
            "fallback": True,
            "reason": "ASR pipeline unavailable",
            "audioBytes": len(audio_bytes)
        }
    }


# ============================================================================
# MAIN ANALYSIS ENDPOINT
# ============================================================================
@app.get("/health")
def health():
    """Health check endpoint with system status."""
    return {
        "ok": True,
        "asr_loaded": ASR_MODEL is not None,
        "alignment_method": "character-level",
        "ml_scorer_available": USE_ML_SCORER,
        "device": DEVICE,
        "production_ready": ASR_MODEL is not None
    }


@app.get("/status")
def status():
    """Detailed system status for debugging."""
    return {
        "service": "Arabic Speech Therapy AI Service",
        "version": "2.0.0",
        "components": {
            "asr_model": {
                "loaded": ASR_MODEL is not None,
                "model": "wav2vec2-large-xlsr-53-arabic" if ASR_MODEL else None,
                "device": DEVICE,
                "parameters": f"{sum(p.numel() for p in ASR_MODEL.parameters()) / 1e6:.1f}M" if ASR_MODEL else None
            },
            "alignment": {
                "method": "character-level",
                "description": "ASR-based character-level phoneme alignment",
                "features": [
                    "Fast processing (2-3 seconds)",
                    "Works with all audio quality",
                    "Handles mispronunciations",
                    "Uses ASR confidence scores",
                    "Production-ready"
                ]
            },
            "scoring": {
                "ml_scorer_enabled": USE_ML_SCORER,
                "method": "ML-based" if USE_ML_SCORER else "Rule-based",
                "description": "Phoneme-level pronunciation scoring"
            }
        },
        "production_ready": ASR_MODEL is not None,
        "performance": {
            "typical_latency": "2-3 seconds",
            "audio_formats": ["WAV", "MP3", "WebM", "M4A"],
            "sample_rate": "16kHz (auto-converted)"
        }
    }


@app.post("/test-upload")
async def test_upload(
    audio: UploadFile = File(...),
    text: str = Form("test")
):
    """Simple test endpoint to verify file uploads work."""
    try:
        audio_bytes = await audio.read()
        return {
            "ok": True,
            "filename": audio.filename,
            "size": len(audio_bytes),
            "text": text,
            "content_type": audio.content_type
        }
    except Exception as e:
        return {
            "ok": False,
            "error": str(e)
        }


@app.post("/check-audio")
async def check_audio(
    expected_text: str = Form(""),
    audio: UploadFile = File(...),
):
    """
    Diagnostic endpoint to check audio file and dictionary coverage.
    Returns detailed information about what's wrong without performing full analysis.
    """
    try:
        # Read audio file
        audio_bytes = await audio.read()
        
        # Step 1: Check audio can be loaded
        try:
            audio_array = convert_audio_to_wav(audio_bytes, target_sr=16000, filename=audio.filename)
            duration = len(audio_array) / 16000
            
            # Check audio stats
            max_amplitude = np.abs(audio_array).max()
            rms = np.sqrt(np.mean(audio_array**2))
            
        except Exception as e:
            return {
                "ok": False,
                "error": "audio_loading",
                "message": f"Failed to load audio: {str(e)}",
                "filename": audio.filename
            }
        
        # Step 2: Try ASR transcription
        transcription = None
        asr_confidence = 0.0
        if ASR_MODEL is not None:
            try:
                transcription, confidence_scores = transcribe_audio(audio_array, sr=16000)
                asr_confidence = float(np.mean(confidence_scores)) if len(confidence_scores) > 0 else 0.0
            except Exception as e:
                transcription = f"ASR Error: {str(e)}"
        
        return {
            "ok": True,
            "audio": {
                "filename": audio.filename,
                "size_bytes": len(audio_bytes),
                "duration_seconds": round(duration, 2),
                "sample_rate": 16000,
                "max_amplitude": round(float(max_amplitude), 3),
                "rms_level": round(float(rms), 3),
                "quality": "good" if max_amplitude > 0.01 and duration > 0.5 else "poor"
            },
            "text": {
                "expected_text": expected_text,
                "test_words": ["مرحبا", "قلم", "كتاب", "بيت", "بنت", "سمكة"]
            },
            "asr": {
                "transcription": transcription,
                "confidence": round(asr_confidence, 3),
                "matches_expected": transcription == expected_text if transcription else False
            },
            "recommendation": _get_recommendation(duration, max_amplitude, transcription, expected_text, asr_confidence)
        }
    
    except Exception as e:
        return {
            "ok": False,
            "error": "diagnostic_failed",
            "message": str(e)
        }


def _get_recommendation(duration, max_amplitude, transcription, expected, confidence):
    """Generate recommendation based on diagnostic results."""
    issues = []
    
    if duration < 0.5:
        issues.append("Audio too short (< 0.5s)")
    if max_amplitude < 0.01:
        issues.append("Audio volume too low - might be silence")
    if transcription and transcription != expected:
        issues.append(f"ASR heard '{transcription}' but expected '{expected}'")
    if confidence < 0.8:
        issues.append(f"ASR confidence low ({confidence:.2f})")
    
    if not issues:
        return "✅ Audio looks good! Ready for analysis."
    else:
        return f"❌ Issues found: {'; '.join(issues)}"


@app.post("/test-ctc")
async def test_ctc_alignment(
    audio: UploadFile = File(...),
    expected_text: str = Form(...),
):
    """
    🧪 TEST ENDPOINT: Detailed CTC alignment testing for Postman validation.
    
    Returns comprehensive alignment data including:
    - CTC alignment with frame-accurate timing
    - Character-level alignment for comparison
    - Phoneme-by-phoneme breakdown
    - Confidence scores and error detection
    - Timing visualization
    
    Args:
        audio: Audio file (WebM, WAV, MP3, etc.)
        expected_text: REQUIRED - The text you're pronouncing (Arabic)
    
    Example Postman Request:
        POST http://localhost:8000/test-ctc
        Body: form-data
            - audio: [your .wav/.webm file]
            - expected_text: "مرحبا"
    """
    start_time = time.time()
    
    try:
        # Validate inputs
        if not expected_text or expected_text.strip() == "":
            raise HTTPException(
                status_code=400,
                detail="expected_text is required for CTC testing. Example: 'مرحبا'"
            )
        
        # Read audio
        audio_bytes = await audio.read()
        print(f"\n{'='*60}")
        print(f"[TEST-CTC] Testing CTC Alignment")
        print(f"[TEST-CTC] Audio: {audio.filename} ({len(audio_bytes)} bytes)")
        print(f"[TEST-CTC] Expected: '{expected_text}'")
        print(f"{'='*60}")
        
        if ASR_MODEL is None:
            raise HTTPException(
                status_code=503,
                detail="ASR model not loaded. Please restart the service."
            )
        
        # Convert audio
        audio_array = convert_audio_to_wav(audio_bytes, target_sr=16000, filename=audio.filename)
        audio_duration = len(audio_array) / 16000
        print(f"[TEST-CTC] Audio duration: {audio_duration:.2f}s ({len(audio_array)} samples)")
        
        # Validate duration
        validate_audio_duration(audio_array, sr=16000)
        
        # Transcribe with CTC data
        print(f"[TEST-CTC] Running ASR transcription...")
        transcription, confidence_scores, logits, predicted_ids = transcribe_audio(
            audio_array, 
            sr=16000,
            return_ctc_data=True
        )
        asr_confidence = float(np.mean(confidence_scores)) if len(confidence_scores) > 0 else 0.0
        print(f"[TEST-CTC] Transcription: '{transcription}' (confidence: {asr_confidence:.3f})")
        
        # Get vocab for CTC
        vocab = ASR_PROCESSOR.tokenizer.get_vocab()
        
        # Run CTC alignment
        print(f"[TEST-CTC] Running CTC alignment...")
        ctc_success = True
        ctc_error_msg = None
        try:
            ctc_phonemes = align_phonemes_ctc(
                audio_array=audio_array,
                expected_text=expected_text,
                transcribed_text=transcription,
                logits=logits,
                predicted_ids=predicted_ids,
                vocab=vocab,
                sr=16000
            )
            print(f"[TEST-CTC] ✓ CTC alignment successful: {len(ctc_phonemes)} phonemes")
        except Exception as e:
            ctc_success = False
            ctc_error_msg = str(e)
            ctc_phonemes = []
            print(f"[TEST-CTC] ✗ CTC alignment failed: {e}")
        
        # Run character-level alignment for comparison
        print(f"[TEST-CTC] Running character-level alignment (for comparison)...")
        char_phonemes = align_phonemes_character_level(
            audio_array=audio_array,
            expected_text=expected_text,
            transcribed_text=transcription,
            confidence_scores=confidence_scores,
            sr=16000
        )
        print(f"[TEST-CTC] ✓ Character-level alignment: {len(char_phonemes)} phonemes")
        
        # Calculate scores
        ctc_score = calculate_overall_score(ctc_phonemes) if ctc_success else 0
        char_score = calculate_overall_score(char_phonemes)
        
        # Create timing comparison visualization
        timing_comparison = []
        for i in range(min(len(ctc_phonemes), len(char_phonemes))):
            ctc_ph = ctc_phonemes[i]
            char_ph = char_phonemes[i]
            
            # Get timing info (alignment returns 'timestamp' and 'duration')
            ctc_start = ctc_ph.get('timestamp', 0)
            ctc_duration = ctc_ph.get('duration', 0)
            ctc_end = ctc_start + ctc_duration
            
            char_start = char_ph.get('timestamp', 0)
            char_duration = char_ph.get('duration', 0)
            char_end = char_start + char_duration
            
            timing_diff = abs(ctc_start - char_start)
            
            timing_comparison.append({
                "phonemeIndex": i,
                "character": ctc_ph.get('symbol', '?'),
                "ctc": {
                    "start": ctc_start,
                    "end": ctc_end,
                    "duration": ctc_duration,
                    "confidence": ctc_ph.get('confidence', 0),
                    "status": "correct" if ctc_ph.get('expected', False) else ctc_ph.get('errorType', 'unknown')
                },
                "characterLevel": {
                    "start": char_start,
                    "end": char_end,
                    "duration": char_duration,
                    "confidence": char_ph.get('confidence', 0),
                    "status": "correct" if char_ph.get('expected', False) else char_ph.get('errorType', 'unknown')
                },
                "timingDifference": round(timing_diff, 3),
                "accuracy": "frame-accurate" if timing_diff < 0.05 else "estimated"
            })
        
        # Calculate accuracy metrics
        correct_ctc = sum(1 for p in ctc_phonemes if p.get('expected', False))
        correct_char = sum(1 for p in char_phonemes if p.get('expected', False))
        
        processing_time = time.time() - start_time
        
        print(f"[TEST-CTC] Processing complete in {processing_time:.3f}s")
        print(f"[TEST-CTC] CTC Score: {ctc_score:.1f}/100 ({correct_ctc}/{len(ctc_phonemes)} correct)")
        print(f"[TEST-CTC] Char Score: {char_score:.1f}/100 ({correct_char}/{len(char_phonemes)} correct)")
        print(f"{'='*60}\n")
        
        return {
            "testStatus": "success",
            "expectedText": expected_text,
            "transcription": transcription,
            "asrConfidence": round(asr_confidence, 3),
            
            "ctcAlignment": {
                "success": ctc_success,
                "error": ctc_error_msg,
                "phonemes": ctc_phonemes,
                "score": round(ctc_score, 2),
                "correctPhonemes": correct_ctc,
                "totalPhonemes": len(ctc_phonemes),
                "accuracy": round((correct_ctc / len(ctc_phonemes) * 100), 1) if ctc_phonemes else 0
            },
            
            "characterLevelAlignment": {
                "phonemes": char_phonemes,
                "score": round(char_score, 2),
                "correctPhonemes": correct_char,
                "totalPhonemes": len(char_phonemes),
                "accuracy": round((correct_char / len(char_phonemes) * 100), 1) if char_phonemes else 0
            },
            
            "comparison": {
                "timingComparison": timing_comparison[:10],  # First 10 phonemes for readability
                "scoreDifference": round(ctc_score - char_score, 2),
                "betterMethod": "CTC" if ctc_score > char_score else "Character-Level",
                "avgTimingDifference": round(
                    np.mean([t['timingDifference'] for t in timing_comparison]),
                    3
                ) if timing_comparison else 0
            },
            
            "audioInfo": {
                "filename": audio.filename,
                "sizeBytes": len(audio_bytes),
                "durationSeconds": round(audio_duration, 2),
                "sampleRate": 16000,
                "samples": len(audio_array)
            },
            
            "performanceMetrics": {
                "processingTimeSeconds": round(processing_time, 3),
                "asrLatency": "~30-50ms (estimated)",
                "ctcLatency": "~75-100ms" if ctc_success else "N/A",
                "totalLatency": f"{round(processing_time * 1000)}ms"
            },
            
            "modelInfo": {
                "asrModel": "jonatasgrosman/wav2vec2-large-xlsr-53-arabic",
                "device": DEVICE,
                "alignmentMethod": "CTC segmentation + fallback",
                "mlScorer": USE_ML_SCORER
            },
            
            "instructions": {
                "message": "✅ CTC alignment test complete!",
                "interpretation": {
                    "score": "0-100 range (higher is better)",
                    "status": "correct/substitution/deletion/insertion",
                    "confidence": "0-1 range (ASR confidence per phoneme)",
                    "timing": "Seconds from start of audio"
                },
                "nextSteps": [
                    "Check 'ctcAlignment.phonemes' for frame-accurate timing",
                    "Compare with 'characterLevelAlignment.phonemes'",
                    "Review 'comparison.timingComparison' for accuracy",
                    "CTC should show ±5-15% timing accuracy vs ±30-50% for character-level"
                ]
            }
        }
    
    except HTTPException:
        raise  # Re-raise FastAPI exceptions
    
    except ValueError as e:
        raise HTTPException(status_code=400, detail=f"Validation error: {str(e)}")
    
    except Exception as e:
        print(f"[TEST-CTC ERROR] {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Test failed: {str(e)}")


@app.post("/analyze")
async def analyze(
    audio: UploadFile = File(...),
    expected_text: str = Form(""),
):
    """
    Analyze Arabic speech pronunciation with ASR + Character-Level Alignment + Phoneme Scoring.
    
    Args:
        audio: Audio file (WebM, WAV, MP3, etc.)
        expected_text: Text the child should pronounce
    
    Returns:
        JSON with score, feedbackTone, phonemes array, and metadata
    """
    start_time = time.time()
    
    try:
        # Read audio file
        audio_bytes = await audio.read()
        print(f"[INFO] Received audio: {audio.filename}, {len(audio_bytes)} bytes")
        
        # Fallback if ASR not loaded
        if ASR_MODEL is None:
            return analyze_fallback(expected_text, audio_bytes)
        
        # Step 1: Convert audio to 16kHz mono WAV
        audio_array = convert_audio_to_wav(audio_bytes, target_sr=16000, filename=audio.filename)
        print(f"[INFO] Audio converted: {len(audio_array)} samples, duration={len(audio_array)/16000:.2f}s")
        validate_audio_duration(audio_array, sr=16000)
        
        # Step 2: ASR Transcription (with CTC data for better alignment)
        transcription, confidence_scores, logits, predicted_ids = transcribe_audio(
            audio_array, 
            sr=16000,
            return_ctc_data=True
        )
        asr_confidence = float(np.mean(confidence_scores)) if len(confidence_scores) > 0 else 0.0
        print(f"[INFO] ASR transcription: '{transcription}' (confidence: {asr_confidence:.3f})")
        
        # Step 3: CTC-Based Alignment (Frame-Accurate)
        print(f"[INFO] Starting CTC-based alignment for expected text: '{expected_text}'")
        try:
            # Get vocab from processor for blank token detection
            vocab = ASR_PROCESSOR.tokenizer.get_vocab()
            
            phoneme_results = align_phonemes_ctc(
                audio_array=audio_array,
                expected_text=expected_text,
                transcribed_text=transcription,
                logits=logits,
                predicted_ids=predicted_ids,
                vocab=vocab,
                sr=16000
            )
            alignment_method = "ctc-segmentation"
            print(f"[INFO] CTC alignment complete: {len(phoneme_results)} phonemes")
        except Exception as ctc_error:
            # Fallback to character-level if CTC fails
            print(f"[WARNING] CTC alignment failed: {ctc_error}, falling back to character-level")
            phoneme_results = align_phonemes_character_level(
                audio_array=audio_array,
                expected_text=expected_text,
                transcribed_text=transcription,
                confidence_scores=confidence_scores,
                sr=16000
            )
            alignment_method = "character-level (fallback)"
        
        # Step 4: Score phonemes with ML model (if available)
        if USE_ML_SCORER:
            phoneme_results = score_phonemes_ml(phoneme_results, audio_array)
        
        # Step 5: Calculate overall score
        overall_score = calculate_overall_score(phoneme_results)
        
        # Step 6: Prepare response
        processing_time = time.time() - start_time
        
        return {
            "score": overall_score,
            "feedbackTone": _tone_from_score(overall_score),
            "phonemes": phoneme_results,
            "meta": {
                "transcription": transcription,
                "expectedText": expected_text,
                "processingTime": round(processing_time, 3),
                "audioBytes": len(audio_bytes),
                "audioDuration": round(len(audio_array) / 16000, 2),
                "modelVersion": "wav2vec2-xlsr-53-arabic",
                "alignmentMethod": alignment_method,
                "useMLScorer": USE_ML_SCORER,
                "device": DEVICE
            }
        }
    
    except ValueError as e:
        # Validation errors (audio too short/long, etc.)
        raise HTTPException(status_code=400, detail=str(e))
    
    except RuntimeError as e:
        # Critical system errors
        print(f"[ERROR] Critical error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"System error: {str(e)}")
    
    except Exception as e:
        # Other unexpected errors - log and fail in production
        print(f"[ERROR] Unexpected analysis error: {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")