import os
import faster_whisper
import gradio as gr
from dotenv import load_dotenv
from huggingface_hub import InferenceClient

# Load API key dari .env
load_dotenv()
HF_API_KEY = os.getenv("HF_API_KEY")

if not HF_API_KEY:
    raise ValueError("API Key Hugging Face tidak ditemukan. Pastikan file .env berisi HF_API_KEY.")

# Inisialisasi klien API Hugging Face
huggingface_client = InferenceClient(api_key=HF_API_KEY)

# Load Faster Whisper model versi large
model = faster_whisper.WhisperModel("turbo", device="cpu", compute_type="int8")

# Daftar model yang dapat dipilih
MODEL_OPTIONS = [
    "mistralai/Mistral-7B-Instruct-v0.3",
    "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "Qwen/Qwen2.5-Coder-32B-Instruct"
]

def save_to_file(content, filename):
    with open(filename, 'w', encoding='utf-8') as file:
        file.write(content)
    return filename

def transcribe_audio(audio_path):
    """Transkripsi audio menggunakan Faster Whisper tanpa koreksi model Hugging Face."""
    segments, _ = model.transcribe(audio_path)
    raw_transcription = " ".join(segment.text for segment in segments)
    return raw_transcription, save_to_file(raw_transcription, 'transcription_large.txt'), audio_path

def generate_soap_summary(transcription_text, selected_model):
    """Membuat ringkasan SOAP dari teks transkripsi menggunakan model yang dipilih."""
    template = """
    Anda adalah asisten medis yang membantu dokter dalam menyusun catatan SOAP berdasarkan percakapan dokter dan pasien.
    Ringkaskan dalam bentuk paragraf tanpa adanya bullet point dan gunakan bahasa Indonesia.
    Harap buat ringkasan dalam format berikut:

    Subjective:
    Objective:
    Assessment:
    Plan:

    ### Percakapan:
    {dialogue}

    ### Catatan SOAP:
    """
    messages = [{"role": "user", "content": template.format(dialogue=transcription_text)}]
    response = huggingface_client.chat.completions.create(
        model=selected_model,
        messages=messages,
        max_tokens=1000,
        stream=False
    )
    soap = response.choices[0].message.content.strip()
    return soap, save_to_file(soap, 'soap_summary.txt')

def detect_medical_tags(transcription_text, selected_model):
    """Mendeteksi tags Diagnosis, Obat, Hasil Lab, dan Radiologi menggunakan model yang dipilih."""
    template = """
    Identifikasi dan berikan luaran dalam bahasa indonesia tags berikut dari percakapan:
    Diagnosis:
    Obat:
    Hasil Lab:
    Radiologi:

    ### Percakapan:
    {dialogue}
    """
    messages = [{"role": "user", "content": template.format(dialogue=transcription_text)}]
    response = huggingface_client.chat.completions.create(
        model=selected_model,
        messages=messages,
        max_tokens=500,
        stream=False
    )
    tags = response.choices[0].message.content.strip()
    return tags, save_to_file(tags, 'medical_tags.txt')

# Antarmuka Gradio
with gr.Blocks(title="AI-based Medical SOAP Summarization and Tag Detection with Faster Whisper Large") as app:
    gr.Markdown("## Medical SOAP Summarization and Tag Detection with Faster Whisper Large")

    with gr.Row():
        with gr.Column():
            model_selector = gr.Dropdown(
                choices=MODEL_OPTIONS,
                value="mistralai/Mixtral-8x7B-Instruct-v0.1",
                label="🔍 Pilih Model AI"
            )
            audio_input = gr.Audio("microphone", type="filepath", label="🎙️ Rekam Suara")
            transcribe_button = gr.Button("🎧 Transkripsi dengan Whisper Large")
            transcription_edit_box = gr.Textbox(label="📄 Hasil Transkripsi (Faster Whisper Large) - Bisa Diedit", lines=12, interactive=True)
            update_transcription_button = gr.Button("💾 Simpan Hasil Edit")
            soap_button = gr.Button("📝 Buat SOAP")
            tags_button = gr.Button("🏷️ Deteksi Tags")

        with gr.Column():
            soap_output = gr.Textbox(label="📃 Hasil SOAP", lines=10, interactive=False)
            tags_output = gr.Textbox(label="🏷️ Hasil Tags Diagnosis, Obat, Hasil Lab, Radiologi", lines=10, interactive=False)
            download_audio = gr.File(label="⬇️ Download Rekaman")
            download_transcription = gr.File(label="⬇️ Download Transkripsi")
            download_soap = gr.File(label="⬇️ Download SOAP")
            download_tags = gr.File(label="⬇️ Download Tags")

    # Tombol Transkripsi
    transcribe_button.click(
        transcribe_audio,
        inputs=[audio_input],
        outputs=[transcription_edit_box, download_transcription, download_audio]
    )

    # Tombol Simpan Hasil Edit
    update_transcription_button.click(
        lambda text: (text, save_to_file(text, 'user_edited_transcription.txt')),
        inputs=[transcription_edit_box],
        outputs=[transcription_edit_box, download_transcription]
    )

    # Tombol SOAP
    soap_button.click(
        generate_soap_summary,
        inputs=[transcription_edit_box, model_selector],
        outputs=[soap_output, download_soap]
    )

    # Tombol Tags
    tags_button.click(
        detect_medical_tags,
        inputs=[transcription_edit_box, model_selector],
        outputs=[tags_output, download_tags]
    )

# Jalankan aplikasi
app.launch(share=True)