import streamlit as st
import whisper
from transformers import pipeline
from tempfile import NamedTemporaryFile
import torch
import os
from huggingface_hub import login

# Configuration and setup
st.set_page_config(
    page_title="Médical SOAP Transcription",
    page_icon="🩺",
    layout="wide"
)

# App title and introduction
st.title("🩺 MédiScribe: Transcription Médicale SOAP")
st.markdown("""
    Convertissez vos dictées médicales en notes structurées au format SOAP 
    (Subjectif, Objectif, Analyse, Plan) en quelques clics. 
    Fonctionne entièrement hors ligne pour préserver la confidentialité des données patients.
""")

# Sidebar for app information and configuration
with st.sidebar:
    st.header("À propos")
    st.info("""
        **MédiScribe** transforme vos enregistrements audio en notes médicales structurées.
        
        **Fonctionnalités:**
        - Transcription audio → texte
        - Structuration automatique SOAP
        - Traitement 100% local et sécurisé
        - Compatible avec le vocabulaire médical français
    """)
    
    st.header("Configuration")
    st.markdown("**Statut des modèles:**")
    device_info = "GPU 🚀" if torch.cuda.is_available() else "CPU 💻"
    st.code(f"Appareil: {device_info}")

# Check for GPU availability
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Authentication for Hugging Face (if token available)
token = os.getenv("HF_AUTH_TOKEN")
if token:
    login(token=token)
    st.sidebar.success("✅ Authentifié avec Hugging Face")
else:
    st.sidebar.warning("⚠️ Token Hugging Face non configuré")

# Load models with caching
@st.cache_resource
def load_models():
    with st.spinner("Chargement des modèles..."):
        # Speech-to-text model
        stt_model = whisper.load_model("medium")
        
        # Text structuring model
        struct_model = pipeline(
            "text-generation",
            model="meta-llama/Llama-3.2-3B-Instruct",
            device=DEVICE,
        )
        return stt_model, struct_model

# Function to structure text using local model
def structure_text_local(text):
    prompt = f"""Transforme ce texte médical en format SOAP de manière concise et stricte.:
Texte: {text}
Format de sortie attendu:
- S (Subjectif) : 
- O (Objectif) : 
- A (Analyse) : 
- P (Plan) :
Ne rien ajouter d'autre que ce format exact."""
    
    response = struct_model(
        prompt,
        max_length=512,
        num_return_sequences=1,
        temperature=0.4,
        repetition_penalty=1.2
    )
    text = response[0]['generated_text']
    text = text.replace(prompt, "")
    text = text.replace("Réponse:", "")
    
    return text.strip()

# Main app flow
try:
    stt_model, struct_model = load_models()
    
    # File uploader section
    st.header("1. Importer un fichier audio")
    uploaded_file = st.file_uploader(
        "Télécharger un enregistrement audio", 
        type=["wav", "mp3", "m4a"],
        help="Formats supportés: WAV, MP3, M4A"
    )

    # Processing section
    if uploaded_file is not None:
        st.header("2. Écouter l'enregistrement")
        
        with NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
            temp_file.write(uploaded_file.getvalue())
            st.audio(temp_file.name, format='audio/wav')
            
            col1, col2 = st.columns(2)
            process_button = col1.button("🔄 Transcrire et Structurer", use_container_width=True)
            
            if process_button:
                st.toast("Démarrage du traitement...")
                
                # Transcription step
                with st.spinner("🎙️ Transcription en cours..."):
                    result = stt_model.transcribe(temp_file.name, language='fr')
                    transcription = result["text"]
                
                st.toast("Structuration en cours...")
                
                # Structuring step
                with st.spinner("🔄 Structuration en format SOAP..."):
                    soap_output = structure_text_local(transcription)
                
                st.toast("Traitement terminé!")
                
                # Display results
                st.header("3. Résultats")
                
                tabs = st.tabs(["Note SOAP Structurée", "Transcription Brute"])
                
                with tabs[0]:
                    st.subheader("Note Médicale Structurée (SOAP)")
                    st.code(soap_output)
                    st.download_button(
                        "💾 Télécharger la note SOAP",
                        soap_output,
                        file_name="note_medicale_soap.txt",
                        mime="text/plain"
                    )
                
                with tabs[1]:
                    st.subheader("Transcription Brute")
                    st.text_area("Texte transcrit", transcription, height=200)
    
    # Instructions section
    else:
        st.info("👆 Commencez par importer un fichier audio pour générer une note médicale SOAP")
        
except Exception as e:
    st.error(f"Une erreur s'est produite: {str(e)}")
    st.info("Veuillez vérifier que tous les modèles et dépendances sont correctement installés.")

# Footer
st.divider()
st.markdown("""
    **Configuration système recommandée:**
    - CPU: 8+ cœurs (pour traitement rapide)
    - RAM: 16GB+ (pour les modèles locaux)
    - GPU: Recommandé pour performances optimales
    - Stockage: 10GB+ d'espace libre
""")