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)


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 medis dalam bentuk paragraf menggunakan 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_file):
    """Transkripsi audio menggunakan Whisper tanpa koreksi model Hugging Face."""
    # segments, _ = model.transcribe(audio_file)
    # raw_transcription = " ".join(segment.text for segment in segments)
    with open(audio_file, "rb") as file:
        res = gclient.audio.transcriptions.create(
            file=(audio_file, file.read()),
            model="whisper-large-v3-turbo",
            language="id",
        )
        print(res)
        raw_transcription = res.text
        soap_output, download_soap = generate_soap_summary(raw_transcription)
        tags_output, download_tags = detect_medical_tags(raw_transcription)
    return (
        save_to_file(raw_transcription, "raw_transcription.txt"),
        audio_file,
        soap_output,
        download_soap,
        tags_output,
        download_tags,
    )


def generate_soap_summary(transcription_text):
    template = """Buat ringkasan SOAP berdasarkan percakapan dokter dan pasien dalam format berikut:

    Subjective:
    ICD10:
    Objective:
    Assessment:
    Plan:

    ### Percakapan:
    {dialogue}
    """

    soap = chat_with_groq(template.format(dialogue=transcription_text))
    return soap, save_to_file(soap, "soap_summary.txt")


def detect_medical_tags(transcription_text):
    """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}
    """
    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 Whisper Large"
) as app:
    gr.Markdown("## Medical SOAP Summarization and Tag Detection with Whisper Large")

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio("microphone", type="filepath", label="🎙️ Rekam Suara")
            transcribe_button = gr.Button("🎧 Tulis Rekam Medis")

        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=[
            download_transcription,
            download_audio,
            soap_output,
            download_soap,
            tags_output,
            download_tags,
        ],
    )

# Jalankan aplikasi
app.launch(share=True)