"""
MedBuddy: Transcrição Médica em Tempo Real com OpenAI e Gradio

Versão: 2.0
Data: 20/07/2025

Este aplicativo Gradio, projetado para rodar em um Hugging Face Space, demonstra
um pipeline de transcrição e sumarização de consultas médicas em tempo real.

Funcionalidades:
  - Captura de áudio do microfone via streaming.
  - Transcrição ao vivo usando a Realtime API da OpenAI (gpt-4o-realtime).
  - Geração de resumos periódicos em "bullet points" (gpt-4o-mini).
  - Elaboração de uma nota final no formato SOAP (gpt-4o).
  - Interface limpa com componentes nativos para copiar texto e baixar o áudio.

Requisitos:
  - gradio
  - openai
  - websockets
  - soundfile
  - numpy

⚠️ AVISO: Este código é um protótipo de referência. Para uso em produção,
é mandatório tratar informações de saúde protegidas (PHI) com o máximo rigor,
utilizar o sistema de Secrets do Hugging Face para chaves de API e implementar
um tratamento de erros mais abrangente.
"""

import asyncio
import json
import os
import tempfile
import time

import gradio as gr
import numpy as np
import openai
import soundfile as sf
import websockets

# -------------------------------------------------------------------
# Configuração
# -------------------------------------------------------------------
# Chave de API da OpenAI (carregada a partir dos Secrets do HF Space)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise ValueError("A variável de ambiente OPENAI_API_KEY não foi definida.")
openai.api_key = OPENAI_API_KEY

# Modelos da OpenAI (verificar documentação para os nomes mais recentes)
STT_MODEL = "gpt-4o-realtime-preview-2025-06-03"
SUMMARY_MODEL = "gpt-4o-mini"
SOAP_MODEL = "gpt-4o"

# Configurações de áudio
SAMPLE_RATE = 16000  # Hz
SUMMARY_EVERY_SEC = 60  # Gerar resumo a cada 60 segundos

# -------------------------------------------------------------------
# Gerenciamento de Estado da Sessão
# -------------------------------------------------------------------
# Nota: Para um ambiente multiusuário, o ideal é usar `gr.State`
# para gerenciar o estado de cada sessão individualmente.
# Para este demo de usuário único, uma classe global simplifica o código.
class SessionState:
    """Encapsula o estado de uma sessão de gravação ativa."""
    def __init__(self):
        self.ws: websockets.WebSocketClientProtocol | None = None
        self.running: bool = False
        self.transcript_full: str = ""
        self.bullets: list[str] = []
        self.last_summary_ts: float = 0.0
        self.audio_chunks: list[np.ndarray] = []
        self.contexto: str = ""

state = SessionState()

# -------------------------------------------------------------------
# Funções de Comunicação com a API OpenAI
# -------------------------------------------------------------------
async def open_realtime_ws() -> websockets.WebSocketClientProtocol:
    """Abre e valida a conexão WebSocket com a Realtime API."""
    uri = f"wss://api.openai.com/v1/realtime?model={STT_MODEL}&sample_rate={SAMPLE_RATE}"
    try:
        ws = await websockets.connect(
            uri,
            extra_headers={"Authorization": f"Bearer {OPENAI_API_KEY}"},
            subprotocols=["realtime"],
            max_size=2 * 1024 * 1024,  # 2 MB
        )
        # O primeiro evento deve ser a confirmação da criação da sessão
        evt = json.loads(await asyncio.wait_for(ws.recv(), timeout=10))
        if evt.get("type") != "session.created":
            raise ConnectionRefusedError(f"Falha ao criar sessão: {evt}")
        return ws
    except Exception as e:
        print(f"Erro ao conectar ao WebSocket: {e}")
        raise

async def pcm_from_numpy(chunk: np.ndarray) -> bytes:
    """Converte um array numpy float32 para bytes no formato PCM 16-bit little-endian."""
    # Garante que o array é float32 e está no range [-1, 1]
    if chunk.dtype != np.float32:
        chunk = chunk.astype(np.float32)
    chunk = np.clip(chunk, -1.0, 1.0)
    
    # Converte para int16
    pcm16 = (chunk * 32767).astype(np.int16)
    return pcm16.tobytes()

# -------------------------------------------------------------------
# Funções de Geração de Texto (Sumário e SOAP)
# -------------------------------------------------------------------
async def summarize_block(text: str) -> str:
    """Gera um resumo conciso (bullet points) para um trecho da transcrição."""
    if not text.strip():
        return ""
    try:
        response = await openai.chat.completions.create(
            model=SUMMARY_MODEL,
            messages=[
                {"role": "system", "content": "Você é um escriba clínico. Resuma o texto a seguir em até 5 bullet points concisos, em português do Brasil."},
                {"role": "user", "content": text},
            ],
            temperature=0.3,
            max_tokens=200,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"Erro ao gerar sumário: {e}")
        return "[Erro ao gerar sumário parcial]"

async def generate_soap(full_txt: str, bullets: list[str], contexto: str) -> str:
    """Gera a nota final no formato SOAP a partir do contexto e da transcrição."""
    if not full_txt.strip():
        return "Nenhuma transcrição foi gerada para criar a nota SOAP."
        
    bullet_summary = "\n".join(bullets)
    prompt_context = f"Contexto prévio da consulta: {contexto if contexto else 'Nenhum'}"
    
    try:
        response = await openai.chat.completions.create(
            model=SOAP_MODEL,
            messages=[
                {"role": "system", "content": "Você é um assistente médico sênior especializado em documentação clínica. Sua tarefa é criar uma nota no formato SOAP (Subjetivo, Objetivo, Avaliação, Plano) baseada na transcrição da consulta."},
                {"role": "user", "content": f"{prompt_context}\n\nResumo dos pontos chave (para guia):\n{bullet_summary}\n\nUse a transcrição completa a seguir para elaborar a nota SOAP final em português do Brasil, de forma estruturada e profissional:\n\n---\n{full_txt}"},
            ],
            temperature=0.2,
            max_tokens=1500,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"Erro ao gerar nota SOAP: {e}")
        return f"[Erro ao gerar nota SOAP final]\n\nTranscrição completa:\n{full_txt}"

# -------------------------------------------------------------------
# Callbacks da Interface Gradio
# -------------------------------------------------------------------
async def cb_start(contexto: str):
    """Callback: Inicia a gravação."""
    if state.running:
        return
    
    state.__init__()  # Reseta o estado da sessão
    state.contexto = contexto
    state.running = True
    
    try:
        state.ws = await open_realtime_ws()
        state.last_summary_ts = time.time()
        print("Sessão de gravação iniciada.")
        # Limpa os campos da UI
        return "", "", "", None
    except Exception as e:
        state.running = False
        gr.Warning(f"Não foi possível iniciar a gravação: {e}")
        return "", "", "", None

async def cb_stream(audio_stream, live_txt, live_sum):
    """Callback: Processa o stream de áudio em tempo real."""
    if not state.running or audio_stream is None or not state.ws:
        return live_txt, live_sum

    try:
        await state.ws.send(await pcm_from_numpy(audio_stream))
        state.audio_chunks.append(audio_stream)

        # Processa mensagens recebidas do WebSocket de forma não-bloqueante
        while True:
            try:
                msg = await asyncio.wait_for(state.ws.recv(), timeout=0.01)
                evt = json.loads(msg)
                if evt.get("type") == "transcript" and (text := evt.get("transcript", {}).get("text")):
                    state.transcript_full += text + " "
            except asyncio.TimeoutError:
                break  # Não há mais mensagens no buffer, sai do loop

        # Gera resumo periódico
        if (time.time() - state.last_summary_ts) >= SUMMARY_EVERY_SEC:
            # Pega os últimos ~4000 caracteres para o resumo parcial
            transcript_slice = state.transcript_full[-4000:]
            bullet = await summarize_block(transcript_slice)
            if bullet:
                state.bullets.append(bullet)
                state.last_summary_ts = time.time()
                
        live_summary_md = "\n\n".join(state.bullets)
        return state.transcript_full, live_summary_md

    except (websockets.exceptions.ConnectionClosed, Exception) as e:
        print(f"Erro durante o streaming: {e}")
        await cb_stop()  # Tenta finalizar a sessão de forma limpa
        return live_txt, live_sum

async def cb_stop(audio_filepath):
    """Callback: Finaliza a gravação, gera a nota SOAP e prepara o download."""
    if not state.running:
        return "", None  # Retorna valores para sumário final e botão de download

    print("Finalizando a gravação...")
    state.running = False
    if state.ws:
        await state.ws.close()

    # Gera um último resumo se houver transcrição nova
    if state.transcript_full and (not state.bullets or (time.time() - state.last_summary_ts) > 15):
        bullet = await summarize_block(state.transcript_full[-4000:])
        if bullet:
            state.bullets.append(bullet)

    soap_note = await generate_soap(state.transcript_full, state.bullets, state.contexto)

    # O `audio_filepath` já é fornecido pelo Gradio quando `type="filepath"`
    # e `stop_recording` é acionado. Não precisamos mais montar o áudio manualmente.
    print(f"Áudio final salvo em: {audio_filepath}")
    
    return soap_note, audio_filepath

# -------------------------------------------------------------------
# Definição da Interface Gráfica (Gradio)
# -------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="MedBuddy – Transcrição Médica") as demo:
    gr.Markdown("# MedBuddy")
    gr.Markdown("### Um Modelo Open-Source de Transcrição Inteligente de Consultas Médicas")

    with gr.Tabs():
        with gr.TabItem("Gravação e Transcrição"):
            contexto_txt = gr.Textbox(
                label="Contexto da Consulta (opcional)",
                lines=3,
                placeholder="Ex.: Paciente com histórico de dispneia crônica, fumante há 20 anos, apresentando tosse persistente."
            )
            
            mic_audio = gr.Audio(
                sources=["microphone"],
                type="filepath",  # 'filepath' é ideal para o botão de download
                label="Microfone (16kHz)",
                streaming=True,
            )

            with gr.Row():
                transcricao_txt = gr.Textbox(label="Transcrição em Tempo Real", lines=15, interactive=False)
                sumario_basico_txt = gr.Textbox(label="Resumo (Bullet Points)", lines=15, interactive=False)
            
            with gr.Accordion("Resultados Finais", open=False):
                sumario_final_txt = gr.Textbox(
                    label="Nota SOAP Final", 
                    lines=15, 
                    interactive=False, 
                    show_copy_button=True
                )
                baixar_btn = gr.DownloadButton(
                    "Baixar Áudio (.wav)",
                    interactive=True
                )

    # ---------------------------------------------------------------
    # Lógica de Eventos da Interface
    # ---------------------------------------------------------------
    mic_audio.start_recording(
        fn=cb_start,
        inputs=[contexto_txt],
        outputs=[transcricao_txt, sumario_basico_txt, sumario_final_txt, baixar_btn]
    )
    
    mic_audio.stream(
        fn=cb_stream,
        inputs=[mic_audio, transcricao_txt, sumario_basico_txt],
        outputs=[transcricao_txt, sumario_basico_txt]
    )
    
    mic_audio.stop_recording(
        fn=cb_stop,
        inputs=[mic_audio],
        outputs=[sumario_final_txt, baixar_btn]
    )

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