import os
import nltk
import uvicorn
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
from groq import Groq
from dotenv import load_dotenv

load_dotenv()

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
groq_client = Groq(api_key=GROQ_API_KEY)

def transcribe_audio(audio_path: str):
    with open(audio_path, "rb") as audio_file:
        response = groq_client.audio.transcriptions.create(
            model="whisper-large-v3",
            file=audio_file,
            response_format="text"
        )
    return response

def summarize_soap(dialogue: str):
    prompt_soap = f"""
    Anda adalah asisten medis yang membantu dokter dalam menyusun catatan SOAP berdasarkan percakapan dokter dan pasien.
    Ringkaskan dalam bentuk paragraf tanpa bullet point dan gunakan bahasa Indonesia.
    Harap buat ringkasan dalam format berikut:
    Subjective:
    Objective:
    Assessment:
    Plan:
    ### Percakapan:
    {dialogue}
    Tolong jangan tambahkan informasi tambahan selain yang berkaitan dengan diagnosis, obat, hasil lab, dan radiologi.
    """
    response_soap = groq_client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[{"role": "user", "content": prompt_soap}]
    )
    return response_soap.choices[0].message.content

def detect_medical_tags(dialogue: str):
    prompt_tags = f"""
    Identifikasi dan berikan luaran dalam bahasa Indonesia tags berikut dari percakapan dengan format:
    Diagnosis:
    Obat:
    Hasil Lab:
    Radiologi:
    ### Percakapan:
    {dialogue}
    Tolong jangan tambahkan informasi tambahan selain yang berkaitan dengan diagnosis, obat, hasil lab, dan radiologi.
    """
    response_tags = groq_client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[{"role": "user", "content": prompt_tags}]
    )
    return response_tags.choices[0].message.content

app = FastAPI(title="Medical Transcription Pipeline (Groq API)")

@app.get("/")
async def root():
    return {
        "message": "🚀 SOAP AI FastAPI is running. Use /full_process or /soap_tags to interact"
    }

@app.post("/full_process")
async def full_process(audio: UploadFile = File(...)):
    try:
        filename = audio.filename
        temp_audio_path = f"/tmp/temp_{filename}"
        with open(temp_audio_path, "wb") as f:
            f.write(await audio.read())

        transcription = transcribe_audio(temp_audio_path)

        soap_content = summarize_soap(transcription)
        tags_content = detect_medical_tags(transcription)

        os.remove(temp_audio_path)

        return {
            "transcription": transcription,
            "soap_content": soap_content,
            "tags_content": tags_content
        }
    except Exception as e:
        return {"error": str(e)}

class TranscriptionInput(BaseModel):
    dialogue: str

@app.post("/soap_tags")
async def soap_tags(data: TranscriptionInput):
    transcript_text = data.dialogue

    soap_content = summarize_soap(transcript_text)
    tags_content = detect_medical_tags(transcript_text)

    return {
        "soap_content": soap_content,
        "tags_content": tags_content
    }

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)