src/main.py

from __future__ import annotations

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


for stream_name in ("stdout", "stderr"):
    stream = getattr(sys, stream_name, None)
    if hasattr(stream, "reconfigure"):
        try:
            stream.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass

import edge_tts
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
import numpy as np
from utils import set_asr_globals, convert_audio_to_wav, validate_audio_duration, transcribe_audio, reduce_noise
try:
    from services.asr_service import transcribe_with_model
except ImportError:
    from .services.asr_service import transcribe_with_model

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

# ML models and alignment
try:
    try:
        # Package-style imports when loaded as src.main
        from .phoneme_features import extract_phoneme_features, get_feature_names
        from .phoneme_scorer import PronunciationScorer
    except ImportError:
        # Script-style imports when loaded as main.py from src working dir
        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)
try:
    from .character_aligner import (
        align_phonemes_character_level,
        align_phonemes_ctc,
        calculate_pronunciation_accuracy
    )
except ImportError:
    from character_aligner import (
        align_phonemes_character_level,
        align_phonemes_ctc,
        calculate_pronunciation_accuracy
    )

BASE_DIR = Path(__file__).resolve().parent.parent
MODELS_DIR = BASE_DIR / "models"

# ============================================================================
# GLOBAL MODEL CACHE (Loaded at startup for low latency)
# ============================================================================
ASR_MODEL: Optional[Any] = None
ASR_PROCESSOR: Optional[Any] = None
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
ANALYZE_ASR_MODEL_ID = "nvidia/stt_ar_fastconformer_hybrid_large_pcd_v1.0"

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

# Startup warmup state. Loading the model is not enough: PyTorch, tokenizer
# paths, CTC softmax, and the optional ML scorer all pay one-time costs on
# their first real use.
MODEL_WARMED = False
MODEL_WARMUP_SECONDS: Optional[float] = None
MODEL_WARMUP_ERROR: Optional[str] = None


def _env_flag(name: str, default: str = "true") -> bool:
    return os.environ.get(name, default).strip().lower() in {"1", "true", "yes", "on"}


def warmup_runtime() -> None:
    """Run a tiny synthetic request path so the first child request is not cold."""
    global MODEL_WARMED, MODEL_WARMUP_SECONDS, MODEL_WARMUP_ERROR

    if ASR_MODEL is None or ASR_PROCESSOR is None:
        return

    warmup_started = time.time()
    sample_rate = 16000
    warmup_seconds = float(os.environ.get("AI_SERVICE_WARMUP_AUDIO_SECONDS", "1.0"))
    warmup_samples = max(sample_rate // 2, int(sample_rate * warmup_seconds))

    # Low-amplitude tone avoids an all-silence edge case while staying harmless.
    t = np.linspace(0, warmup_samples / sample_rate, warmup_samples, endpoint=False)
    synthetic_audio = (0.01 * np.sin(2 * np.pi * 220 * t)).astype(np.float32)

    try:
        # Warm soundfile/librosa conversion with the same 16 kHz WAV shape that
        # the child assessment frontend normally uploads.
        with io.BytesIO() as wav_buffer:
            import soundfile as sf

            sf.write(wav_buffer, synthetic_audio, sample_rate, format="WAV")
            convert_audio_to_wav(wav_buffer.getvalue(), target_sr=sample_rate, filename="warmup.wav")

        transcription, confidence_scores, logits, predicted_ids = transcribe_audio(
            synthetic_audio,
            sr=sample_rate,
            return_ctc_data=True,
        )
        vocab = ASR_PROCESSOR.tokenizer.get_vocab()
        align_phonemes_ctc(
            audio_array=synthetic_audio,
            expected_text="ماما",
            transcribed_text=transcription,
            logits=logits,
            predicted_ids=predicted_ids,
            vocab=vocab,
            sr=sample_rate,
        )

        if USE_ML_SCORER and PHONEME_SCORER is not None:
            dummy_phoneme = {
                "symbol": "م",
                "expected": True,
                "confidence": float(np.mean(confidence_scores)) if len(confidence_scores) else 0.85,
                "duration": 0.12,
                "timestamp": 0.0,
            }
            score_phonemes_ml([dummy_phoneme], synthetic_audio)

        if torch.cuda.is_available():
            torch.cuda.synchronize()

        MODEL_WARMED = True
        MODEL_WARMUP_SECONDS = round(time.time() - warmup_started, 3)
        MODEL_WARMUP_ERROR = None
        print(f"[STARTUP] ✓ Runtime warmup complete in {MODEL_WARMUP_SECONDS:.3f}s")
    except Exception as e:
        MODEL_WARMED = False
        MODEL_WARMUP_SECONDS = round(time.time() - warmup_started, 3)
        MODEL_WARMUP_ERROR = str(e)
        print(f"[STARTUP] ⚠ Runtime warmup failed after {MODEL_WARMUP_SECONDS:.3f}s: {e}")


# ============================================================================
# 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
    
    if _env_flag("AI_SERVICE_LOAD_LEGACY_ASR", "false"):
        from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor

        print(f"[STARTUP] Loading legacy 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] ✓ Legacy ASR model loaded successfully ({sum(p.numel() for p in ASR_MODEL.parameters()) / 1e6:.1f}M parameters)")
        except Exception as e:
            ASR_MODEL = None
            ASR_PROCESSOR = None
            print(f"[STARTUP] ⚠ Failed to load legacy startup model: {e}")
    else:
        ASR_MODEL = None
        ASR_PROCESSOR = None
        print("[STARTUP] Legacy wav2vec2 startup model skipped (AI_SERVICE_LOAD_LEGACY_ASR=false)")

    # Shared utility functions must see the loaded model before startup warmup.
    set_asr_globals(ASR_MODEL, ASR_PROCESSOR, DEVICE)
    
    # Load phoneme scoring model (PyTorch-based - optional). Spaces default to
    # rule-based scoring so /analyze can boot quickly and avoid CPU cold-starts.
    if ML_SCORER_AVAILABLE and _env_flag("AI_SERVICE_ENABLE_ML_SCORER", "false"):
        try:
            PHONEME_SCORER = PronunciationScorer(models_dir=str(MODELS_DIR))
            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
    else:
        PHONEME_SCORER = None
        USE_ML_SCORER = False
        print("[STARTUP] ML scorer skipped (AI_SERVICE_ENABLE_ML_SCORER=false)")
    
    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}")
    
    if _env_flag("AI_SERVICE_WARMUP", "false"):
        warmup_runtime()
    else:
        print("[STARTUP] Runtime warmup skipped (AI_SERVICE_WARMUP=false)")

    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",
        "https://nateq.online",
        "https://www.nateq.online",
        "https://smart-arabic-speech-therapist-production.up.railway.app",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)



# ============================================================================
# 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
    return "".join(text.split()).strip()


def _build_phoneme_result(
    symbol: str,
    is_correct: bool,
    confidence: float,
    error_type: Optional[str],
    duration: float,
    timestamp: float,
    substituted_with: Optional[str] = None
) -> Dict[str, Any]:
    """Create a phoneme response entry with optional substitution metadata."""
    phoneme = {
        "symbol": symbol,
        "expected": is_correct,
        "confidence": confidence,
        "errorType": error_type,
        "duration": duration,
        "timestamp": timestamp
    }

    if error_type == "substitution" and substituted_with is not None:
        phoneme["substitutedWith"] = substituted_with

    return phoneme

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(_build_phoneme_result(
            symbol=expected_char,
            is_correct=is_correct,
            confidence=round(final_confidence, 3),
            error_type=error_type,
            duration=round(duration, 3),
            timestamp=round(start_time, 3),
            substituted_with=transcribed_char if error_type == "substitution" else None
        ))
    
    # 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(_build_phoneme_result(
                symbol=transcribed_clean[i],
                is_correct=False,
                confidence=0.3,
                error_type="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(_build_phoneme_result(
            symbol=ch,
            is_correct=ok,
            confidence=0.80 if ok else 0.50,
            error_type=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("/")
def root():
    return {
        "ok": True,
        "service": "SmartArabicSpeechTherapy AI",
        "focus": "/analyze",
        "docs": "/docs",
    }


@app.get("/health")
def health():
    """Health check endpoint with system status."""
    return {
        "ok": True,
        "analyze_model": ANALYZE_ASR_MODEL_ID,
        "legacy_asr_loaded": ASR_MODEL is not None,
        "alignment_method": "character-level (nvidia-fastconformer)",
        "ml_scorer_available": USE_ML_SCORER,
        "device": DEVICE,
        "warmed": MODEL_WARMED,
        "warmupSeconds": MODEL_WARMUP_SECONDS,
        "warmupError": MODEL_WARMUP_ERROR,
        "production_ready": True
    }


@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)"
        },
        "warmup": {
            "enabled": _env_flag("AI_SERVICE_WARMUP", "true"),
            "complete": MODEL_WARMED,
            "seconds": MODEL_WARMUP_SECONDS,
            "error": MODEL_WARMUP_ERROR
        }
    }


@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")

        # 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 using NVIDIA FastConformer
        asr_result = transcribe_with_model(
            model_id=ANALYZE_ASR_MODEL_ID,
            audio_bytes=audio_bytes,
            filename=audio.filename,
        )
        if asr_result.get("error"):
            raise RuntimeError(asr_result["error"])

        transcription = (asr_result.get("transcription") or "").strip()
        asr_confidence = 0.0
        print(f"[INFO] ASR transcription: '{transcription}'")

        # Step 3: Character-level alignment using a uniform confidence curve.
        # NeMo does not expose the same CTC logits shape as the legacy wav2vec2
        # pipeline, so we align the transcript at the character level here.
        confidence_scores = np.ones(max(1, len(audio_array) // 320), dtype=np.float32)
        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 (nvidia-fastconformer)"
        
        # 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": ANALYZE_ASR_MODEL_ID,
                "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)}")


# ============================================================================
# TTS ENDPOINT (edge-tts: Microsoft Edge TTS, Egyptian Arabic)
# ============================================================================
@app.get("/tts")
async def text_to_speech(text: str = "", voice: str = "ar-EG-SalmaNeural"):
    """
    Convert Arabic text to speech using Microsoft Edge TTS.
    Strips diacritics before synthesis for cleaner pronunciation.
    Returns audio (MP3 by default).

    Query params:
        text  (required) - Arabic text to speak
        voice (optional) - TTS voice (default: ar-EG-SalmaNeural,
                           also available: ar-EG-ShakirNeural)
    """
    if not text.strip():
        raise HTTPException(status_code=400, detail="Text is required")

    clean_text = ARABIC_DIACRITICS.sub("", text).strip()
    start = time.time()

    try:
        communicate = edge_tts.Communicate(clean_text, voice)
        audio_data = b""
        async for chunk in communicate.stream():
            if chunk["type"] == "audio":
                audio_data += chunk["data"]

        if not audio_data:
            raise HTTPException(status_code=500, detail="No audio generated")

        elapsed = round((time.time() - start) * 1000)
        return Response(
            content=audio_data,
            media_type="audio/mpeg",
            headers={
                "X-TTS-Model": f"edge-tts/{voice}",
                "X-Response-Time-Ms": str(elapsed),
            },
        )
    except HTTPException:
        raise
    except Exception as e:
        print(f"[TTS] Error: {e}")
        raise HTTPException(status_code=500, detail=f"TTS failed: {str(e)}")