"""
Gradio interface for Speech Disfluency Detection
-------------------------------------------------
Wraps ChildFluencyNet (stage2_fyp_best.pt).
Part of a multi-modal neurodevelopmental assessment system.

Same inference logic as FastAPI version, but deployed on HF Spaces via Gradio.
"""

import io
import json
import os
import re
import sys
import time
import yaml
import tempfile
import traceback
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import torch
import numpy as np
import librosa
import gradio as gr
from huggingface_hub import hf_hub_download
from dotenv import load_dotenv

load_dotenv()

# ── Path setup ────────────────────────────────────────────────────────────
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

# Import your model
try:
    from childfluency import ChildFluencyNet
except ImportError:
    print("[startup] Warning: Could not import ChildFluencyNet from local childfluency.py")
    print("[startup] Make sure childfluency.py is in the same directory as app.py")

# ── Load params ───────────────────────────────────────────────────────────
with open(ROOT / 'params.yaml') as f:
    params = yaml.safe_load(f)

TARGET_SR  = params['data']['target_sr']   # 16000
WINDOW_SEC = params['data']['window_sec']  # 4.0
STRIDE_SEC = params['data']['stride_sec']  # 2.0
MIN_RMS    = params['data']['min_rms']     # 0.001

LPE_OUTPUT_DIM     = params['model']['lpe_dim']  # 32
LPE_INPUT_FEATURES = 6                           # 6

DEFAULT_CKPT = os.environ.get(
    'DISFLUENCY_CHECKPOINT',
    str(ROOT / 'stage2_fyp_best.pt')
)

# ── Gemini & Groq ─────────────────────────────────────────────────────────
GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')
GROK_API_KEY   = os.environ.get('GROK_API_KEY')
gemini_model = None

if GEMINI_API_KEY:
    try:
        import google.generativeai as genai
        genai.configure(api_key=GEMINI_API_KEY)
        gemini_model = genai.GenerativeModel('gemini-2.0-flash')
        print("[startup] Gemini API initialised.")
    except Exception as e:
        print(f"[startup] Warning: Gemini init failed: {e}")
else:
    print("[startup] GEMINI_API_KEY not set — will rely on Groq fallback.")

if GROK_API_KEY:
    print("[startup] Groq API key found — available as fallback.")
else:
    print("[startup] GROK_API_KEY not set — Groq fallback disabled.")

# ── Load Model at Startup ──────────────────────────────────────────────────
print("[startup] Loading model...")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"[startup] Device: {device}")

# Download model from HF Hub
print("[startup] Downloading model from HuggingFace Hub...")
try:
    model_path = hf_hub_download(
        repo_id="Shubham2004/developmental-stuttering-detection",
        filename="stage2_fyp_best.pt",
        cache_dir="./models"
    )
except Exception as e:
    print(f"[startup] Warning: HF Hub download failed: {e}")
    model_path = Path(DEFAULT_CKPT)

print(f"[startup] Model path: {model_path}")

# Initialize and load model
model = ChildFluencyNet(
    wavlm_name   = params['model'].get('wavlm_name', 'microsoft/wavlm-large'),
    acoustic_dim = params['model']['acoustic_dim'],
    lpe_dim      = LPE_OUTPUT_DIM,
    hidden_dim   = params['model']['hidden_dim'],
    dropout      = params['model']['dropout']
).to(device)

if Path(model_path).exists():
    ckpt = torch.load(str(model_path), map_location=device)
    model.load_state_dict(ckpt['model'])
    epoch = ckpt.get('epoch', '?')
    val_f1 = round(float(ckpt.get('val_f1', 0)), 4)
    print(f"[startup] Loaded → epoch {epoch}, val_f1={val_f1}")
else:
    print(f"[startup] WARNING: checkpoint not found at {model_path}")

model.eval()
print("[startup] Model ready for inference.")


# ── Inference Logic (kept from FastAPI version) ────────────────────────────

def segment_audio(audio: np.ndarray):
    window_samples = int(WINDOW_SEC * TARGET_SR)
    stride_samples = int(STRIDE_SEC * TARGET_SR)
    windows, start = [], 0
    while start + window_samples <= len(audio):
        seg = audio[start: start + window_samples]
        if float(np.sqrt(np.mean(seg ** 2))) >= MIN_RMS:
            windows.append((
                start,
                start + window_samples,
                round(start / TARGET_SR, 3),
                round((start + window_samples) / TARGET_SR, 3),
            ))
        start += stride_samples
    return windows


def predict_windows(audio: np.ndarray, threshold: float = 0.5, batch_size: int = 16) -> List[dict]:
    """Segment audio and run SLD inference."""
    windows = segment_audio(audio)
    if not windows:
        return []

    results = []
    for b in range(0, len(windows), batch_size):
        batch = windows[b: b + batch_size]
        audio_t = torch.from_numpy(
            np.stack([audio[s:e].astype(np.float32) for s, e, _, _ in batch])
        ).to(device)

        lpe_t = torch.full(
            (len(batch), LPE_INPUT_FEATURES), 0.5, dtype=torch.float32
        ).to(device)

        with torch.no_grad():
            if device == 'cuda':
                with torch.cuda.amp.autocast():
                    out = model(audio_t, lpe_t)
            else:
                out = model(audio_t, lpe_t)

        probs = torch.sigmoid(out['sld_logit']).cpu().numpy().flatten()

        for i, (_, _, s_sec, e_sec) in enumerate(batch):
            p = float(probs[i])
            results.append({
                'window_idx': b + i,
                'start_sec' : s_sec,
                'end_sec'   : e_sec,
                'sld_prob'  : round(p, 4),
                'prediction': 'SLD' if p >= threshold else 'fluent',
            })

    return results


def compute_confidence(positive_rate: float, n_windows: int) -> str:
    if n_windows < 5:
        return 'low'
    consensus = max(positive_rate, 1 - positive_rate)
    if consensus >= 0.80:
        return 'high'
    elif consensus >= 0.60:
        return 'medium'
    return 'low'


# ── Gradio Main Function ───────────────────────────────────────────────────

def analyze_speech(audio_file, threshold: float = 0.5) -> Tuple[str, str]:
    """
    Main Gradio function: analyze audio and return results + optional analysis.
    
    Args:
        audio_file: Path to uploaded audio file (Gradio provides this)
        threshold: SLD probability threshold (default 0.5)
    
    Returns:
        (results_json, analysis_text)
    """
    try:
        t_start = time.time()
        
        # Load audio
        audio, _ = librosa.load(audio_file, sr=TARGET_SR, mono=True)
        duration_sec = round(len(audio) / TARGET_SR, 2)
        
        if duration_sec < WINDOW_SEC:
            error_msg = f"Audio too short ({duration_sec}s). Need ≥{WINDOW_SEC}s."
            return (
                json.dumps({"error": error_msg, "status": "failed"}, indent=2),
                f"❌ {error_msg}"
            )
        
        # Run inference
        window_scores = predict_windows(audio, threshold=threshold)
        
        if not window_scores:
            error_msg = "No valid speech windows found."
            return (
                json.dumps({"error": error_msg, "status": "failed"}, indent=2),
                f"❌ {error_msg}"
            )
        
        # Compute metrics
        n_total = len(window_scores)
        n_positive = sum(1 for w in window_scores if w['prediction'] == 'SLD')
        pos_rate = round(n_positive / n_total, 4)
        
        pos_probs = [w['sld_prob'] for w in window_scores if w['prediction'] == 'SLD']
        sld_prob = round(
            float(np.mean(pos_probs)) if pos_probs
            else float(np.mean([w['sld_prob'] for w in window_scores])),
            4
        )
        
        elapsed = round(time.time() - t_start, 2)
        
        # Build results dict
        results = {
            'filename': Path(audio_file).name,
            'duration_sec': duration_sec,
            'sld_probability': sld_prob,
            'prediction': 'stuttering_detected' if n_positive > 0 else 'fluent',
            'confidence': compute_confidence(pos_rate, n_total),
            'n_windows_total': n_total,
            'n_windows_positive': n_positive,
            'positive_rate': pos_rate,
            'processing_time_sec': elapsed,
            'note': "LPE features set to neutral (0.5). Ablation study: ΔF1=+0.002 (negligible)."
        }
        
        results_json = json.dumps(results, indent=2)
        
        # Generate text summary
        severity = (
            'Severe'   if pos_rate >= 0.50 else
            'Moderate' if pos_rate >= 0.25 else
            'Mild'     if pos_rate >= 0.10 else
            'Minimal'
        )
        
        summary = f"""
✅ ANALYSIS COMPLETE

📊 Results:
- File: {Path(audio_file).name}
- Duration: {duration_sec}s
- SLD Rate: {pos_rate * 100:.1f}%
- Severity: {severity}
- Confidence: {results['confidence'].upper()}
- Windows Analyzed: {n_total}
- Windows with SLD: {n_positive}

🔍 Prediction: {'🚨 STUTTERING DETECTED' if results['prediction'] == 'stuttering_detected' else '✅ FLUENT (No significant stuttering)'}

⏱️ Processing time: {elapsed}s

📝 Note: This is a screening aid only. Consult a Speech-Language Pathologist for formal evaluation.
"""
        
        return (results_json, summary)
    
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return (
            json.dumps({"error": error_msg, "status": "failed"}, indent=2),
            f"❌ {error_msg}"
        )


# ── Optional: Generate Gemini Analysis ─────────────────────────────────────

def generate_analysis(
    audio_file,
    child_name: str = "Child",
    child_age: str = "?",
    threshold: float = 0.5
) -> str:
    """Generate a detailed Gemini-powered analysis of the results."""
    
    try:
        # First, run the prediction
        audio, _ = librosa.load(audio_file, sr=TARGET_SR, mono=True)
        duration_sec = round(len(audio) / TARGET_SR, 2)
        
        window_scores = predict_windows(audio, threshold=threshold)
        if not window_scores:
            return "❌ No valid speech windows found."
        
        n_total = len(window_scores)
        n_positive = sum(1 for w in window_scores if w['prediction'] == 'SLD')
        pos_rate = round(n_positive / n_total, 4)
        sld_pct = round(pos_rate * 100, 1)
        
        detected = n_positive > 0
        severity = (
            'Severe'   if pos_rate >= 0.50 else
            'Moderate' if pos_rate >= 0.25 else
            'Mild'     if pos_rate >= 0.10 else
            'Minimal'
        )
        
        confidence = compute_confidence(pos_rate, n_total)
        
        # Fallback text
        fallback = (
            f"OVERALL ASSESSMENT\n"
            f"{child_name} ({child_age} years) showed {severity.lower()} stutter-like disfluencies "
            f"in {sld_pct}% of the recorded speech "
            f"({'stuttering detected' if detected else 'no significant stuttering detected'}).\n\n"
            f"ACOUSTIC DETAILS\n"
            f"Total windows analysed: {n_total}\n"
            f"Windows with SLD: {n_positive}\n"
            f"SLD rate: {sld_pct}%\n"
            f"Confidence: {confidence}\n"
            f"Recording duration: {duration_sec:.1f}s\n\n"
            f"RECOMMENDATIONS\n"
            f"1. {'Consult a Speech-Language Pathologist for a formal evaluation.' if detected else 'Continue monitoring speech fluency over time.'}\n"
            f"2. Use slow, relaxed speech when speaking with the child.\n"
            f"3. Create a low-pressure environment to reduce communication anxiety.\n\n"
            f"IMPORTANT DISCLAIMER\n"
            f"This is an automated screening aid and does not constitute a clinical diagnosis. "
            f"A qualified Speech-Language Pathologist must evaluate the child formally."
        )
        
        prompt = f"""You are a Speech-Language Pathologist writing a screening summary for a parent.

YOUR INSTRUCTIONS:
1. DO NOT recalculate any scores.
2. DO NOT change the classification.
3. Your ONLY job is to format the provided data into a warm, empathetic, easy-to-read report.

CHILD INFORMATION:
- Name: {child_name}
- Age: {child_age} years

ACOUSTIC SCREENING RESULTS (Do not alter):
- SLD (Stutter-Like Disfluency) Rate: {sld_pct}%
- Severity Category: {severity}
- System Classification: {'STUTTERING DETECTED' if detected else 'FLUENT (No significant stuttering)'}
- Confidence: {confidence}
- Recording Duration: {duration_sec:.1f} seconds
- Windows Analysed: {n_total}, Windows with SLD: {n_positive}

Please write a comprehensive, plain-text summary (no markdown, no **, no ##, no special characters).

Structure your response with these exact section headings on their own lines:

OVERALL ASSESSMENT
2-3 sentences summarising the result in plain, reassuring language.

WHAT THIS MEANS
Explain what stutter-like disfluencies are and how this result relates to typical speech development for this age.

RECOMMENDATIONS
Simple numbered points (1., 2., 3.) — at least three actionable suggestions for parents.

IMPORTANT DISCLAIMER
State clearly this is a screening tool, not a clinical diagnosis, and that a Speech-Language Pathologist evaluation is required.

Keep the tone warm, empathetic, and parent-friendly. Plain text only."""

        # Tier 1: Gemini
        if gemini_model is not None:
            try:
                response = gemini_model.generate_content(prompt)
                return response.text
            except Exception as e:
                print(f"[generate-analysis] Gemini error: {e}. Falling back to Groq.")

        # Tier 2: Groq
        if GROK_API_KEY:
            try:
                from groq import Groq
                groq_client = Groq(api_key=GROK_API_KEY)
                groq_response = groq_client.chat.completions.create(
                    model="llama-3.3-70b-versatile",
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.65,
                    max_tokens=1024,
                )
                return groq_response.choices[0].message.content
            except Exception as e:
                print(f"[generate-analysis] Groq error: {e}. Using fallback.")

        # Tier 3: Fallback
        return fallback
    
    except Exception as e:
        return f"❌ Error generating analysis: {str(e)}"


# ── Gradio Interface (Simplified) ──────────────────────────────────────────

def create_demo():
    with gr.Blocks(title="Developmental Stuttering Detection") as demo:
        gr.Markdown("# 🎤 Developmental Stuttering Detection")
        gr.Markdown("Upload a recording of a child speaking to screen for signs of stutter-like disfluencies.")
        
        with gr.Tabs():
            # TAB 1: Prediction
            with gr.Tab("Analysis"):
                with gr.Row():
                    audio_input = gr.Audio(
                        type="filepath",
                        label="Upload Audio (WAV recommended, min 4 seconds)"
                    )
                    threshold_slider = gr.Slider(
                        minimum=0.1,
                        maximum=0.9,
                        value=0.5,
                        step=0.05,
                        label="SLD Threshold"
                    )
                
                analyze_btn = gr.Button("🔍 Analyze", variant="primary")
                
                with gr.Row():
                    results_output = gr.JSON(label="Raw Results")
                    summary_output = gr.Textbox(label="Summary", lines=10)
                
                analyze_btn.click(
                    fn=analyze_speech,
                    inputs=[audio_input, threshold_slider],
                    outputs=[results_output, summary_output]
                )
            
            # TAB 2: Detailed Analysis
            with gr.Tab("Detailed Report"):
                gr.Markdown("Generate a detailed clinical-style report powered by AI.")
                
                audio_input_2 = gr.Audio(type="filepath", label="Upload Audio")
                child_name_input = gr.Textbox(value="Child", label="Child's Name")
                child_age_input = gr.Textbox(value="?", label="Child's Age")
                threshold_slider_2 = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.5, step=0.05, label="SLD Threshold"
                )
                
                generate_btn = gr.Button("📄 Generate Report", variant="primary")
                report_output = gr.Textbox(label="Generated Report", lines=15)
                
                generate_btn.click(
                    fn=generate_analysis,
                    inputs=[audio_input_2, child_name_input, child_age_input, threshold_slider_2],
                    outputs=report_output
                )
        
        gr.Markdown("""
---
⚠️ **IMPORTANT DISCLAIMER**

This tool is a **screening aid only** and does not constitute a clinical diagnosis.
All results must be reviewed by a qualified Speech-Language Pathologist.
""")
    
    return demo


# ── Launch ──────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(share=False)