import os
import sys
import json
import numpy as np
import gradio as gr
import spaces
from pathlib import Path

sys.path.append(str(Path(__file__).resolve().parent.parent))
from src.config import TOTAL_AUDIO_FEATURES

text_clf = None
audio_clf = None
fusion_engine = None
transcriber = None
text_explainer = None
audio_explainer = None

def ensure_loaded():
    global text_clf, audio_clf, fusion_engine, transcriber, text_explainer, audio_explainer
    if text_clf is not None: return
    from src.models.text_classifier import LinguisticStressClassifier
    from src.models.audio_classifier import AudioEnsemblePipeline
    from src.models.fusion_engine import LateDecisionFusion
    from src.speech.transcriber import SpeechTranscriber
    from src.xai.text_explainer import TextExplainerLIME
    from src.xai.audio_explainer import AudioExplainerSHAP
    
    text_clf = LinguisticStressClassifier()
    text_clf.load_model()
    audio_clf = AudioEnsemblePipeline()
    audio_clf.load_model()
    fusion_engine = LateDecisionFusion(text_pipeline=text_clf, audio_pipeline=audio_clf)
    transcriber = SpeechTranscriber()
    text_explainer = TextExplainerLIME(text_classifier=text_clf)
    audio_explainer = AudioExplainerSHAP(audio_classifier=audio_clf)

@spaces.GPU
def process_audio(audio_filepath: str, external_text: str = None):
    if not audio_filepath or not os.path.exists(audio_filepath):
        return {"error": "No audio"}
    
    # Step 1: Extract acoustic features from the audio file using librosa
    features = None
    try:
        from src.data_prep.audio_processor import extract_195_features_from_audio
        extracted = extract_195_features_from_audio(audio_filepath)
        if extracted is not None and not np.all(extracted == 0):
            features = list(extracted)
            print(f"[Audio] Successfully extracted {len(features)} acoustic features.")
        else:
            print("[Audio] Feature extraction returned zeros - librosa may have failed. Using audio-only fallback with zero features skipped.")
    except Exception as e:
        print(f"[Audio] Feature extraction exception: {e}")
    
    # Step 2: Transcribe spoken words using Whisper (if not provided externally)
    text = external_text if external_text else ""
    if not text and transcriber:
        try:
            text = transcriber.transcribe(audio_filepath).get("transcribed_text", "")
            print(f"[Audio] Whisper transcription: '{text}'")
        except Exception as e:
            print(f"[Audio] Whisper transcription failed: {e}")
    
    # Step 3: Run fusion — only pass audio features if they are valid (not zero/failed)
    # If features failed, fall back to text-only (Whisper transcription) analysis
    # This prevents random noise from contaminating the final clinical score
    if features is not None:
        res = fusion_engine.analyze_multimodal(
            text_input=text if text else None,
            audio_features_195=features
        )
    else:
        # Feature extraction failed — run text-only analysis from transcription
        print("[Audio] Falling back to text-only analysis using Whisper transcription.")
        res = fusion_engine.analyze_multimodal(
            text_input=text if text else None,
            audio_features_195=None
        )
    
    if features is not None and audio_explainer and res.get("audio_analysis"):
        try:
            axai = audio_explainer.explain_instance(features)
            res["audio_xai"] = axai
        except Exception as e:
            print(f"[XAI] Audio SHAP error: {e}")
            
    if text and text_explainer and res.get("text_analysis"):
        try:
            txai = text_explainer.explain_instance(text)
            res["text_xai"] = txai
        except Exception as e:
            print(f"[XAI] Text LIME error: {e}")
            
    return {"transcription": {"text": text}, "fusion_result": res}

@spaces.GPU
def process_text(text: str):
    # Text-only: do NOT pass random audio noise — use pure text analysis
    res = fusion_engine.analyze_multimodal(text_input=text, audio_features_195=None)
    
    if text and text_explainer and res.get("text_analysis"):
        try:
            txai = text_explainer.explain_instance(text)
            res["text_xai"] = txai
        except Exception as e:
            print(f"[XAI] Text LIME error: {e}")
            
    return res

# Pre-load all models into memory before starting Gradio to prevent ZeroGPU timeouts
ensure_loaded()

with gr.Blocks() as demo:
    gr.Markdown("# NeuroSense AI GPU API")
    
    with gr.Tab("Audio"):
        audio_in = gr.Audio(type="filepath")
        external_text_in = gr.Textbox(visible=False)
        audio_out = gr.JSON()
        audio_btn = gr.Button("Analyze Audio")
        audio_btn.click(fn=process_audio, inputs=[audio_in, external_text_in], outputs=audio_out, api_name="analyze_audio")
        
    with gr.Tab("Text"):
        text_in = gr.Textbox()
        text_out = gr.JSON()
        text_btn = gr.Button("Analyze Text")
        text_btn.click(fn=process_text, inputs=text_in, outputs=text_out, api_name="analyze_text")

if __name__ == "__main__":
    demo.launch()