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

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

if not GROQ_API_KEY:
    raise ValueError("GROQ API NOT FOUND!")
gclient = Groq(api_key=GROQ_API_KEY)

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")


def chat_with_groq(message):
    """Handles conversation with Groq LLM."""
    response = gclient.chat.completions.create(
        model="gemma2-9b-it",
        messages=[
            {
                "role": "system",
                "content": """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.""",
            },
            {"role": "user", "content": message},
        ],
        temperature=0.0,
        max_tokens=248,
    )
    return response.choices[0].message.content  # Extract response text


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 = """Harap buat ringkasan dalam format berikut:

    Subjective:
    ICD10:
    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()
    soap = chat_with_groq(template.format(dialogue=transcription_text))
    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 saran dalam bahasa Indonesia tindakan logis selanjutnya dalam format:
    ICD10:
    Obat:
    Laboratorium:
    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()
    tags = chat_with_groq(template.format(dialogue=transcription_text))
    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():
            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=3,
                interactive=True,
            )
            update_transcription_button = gr.Button("💾 Simpan Hasil Edit")
            soap_button = gr.Button("📝 Buat SOAP")
            tags_button = gr.Button("🏷️ Deteksi Saran Tags")

        with gr.Column():
            soap_output = gr.Textbox(label="📃 Hasil SOAP", lines=10, interactive=False)
            tags_output = gr.Textbox(
                label="🏷️ Hasil  Saran Tags ICD 10, Obat, Laboratorium, 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],
        outputs=[soap_output, download_soap],
    )

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

# Jalankan aplikasi
app.launch(share=True)