utils.py

import base64
import json
from datetime import datetime
from io import BytesIO
from typing import Any, Dict, List, Optional

import requests
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from PIL import Image

# ============================================================
# CONFIG
# ============================================================

ENDPOINT_URL = "https://hariharansuthan05--medgemma-1-5-deployment-medgemmaservi-1b5255.modal.run"


def query_medgemma(prompt: str, image_data: str = None):
    """Sends a payload to the hosted MedGemma API endpoint."""
    payload = {"prompt": prompt}
    if image_data:
        payload["image_data"] = image_data

    try:
        response = requests.post(ENDPOINT_URL, json=payload, timeout=300)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return {"error": f"API request failed: {e}"}

# ============================================================
# SYSTEM PROMPT
# ============================================================

MEDICAL_SYSTEM_PROMPT = """
You are MedGemma AI, an expert medical imaging assistant.

ROLE:
- Radiology support
- Differential diagnosis
- Clinical reasoning

OUTPUT FORMAT:

1. FINDINGS
2. IMPRESSION
3. DIFFERENTIAL DIAGNOSES
4. RECOMMENDATIONS

Use medical terminology.
Recommend clinical correlation.
"""

# ============================================================
# PROMPTS
# ============================================================

image_analysis_template = PromptTemplate.from_template(
    """
Please analyze this {image_type} image.

Clinical Question:
{clinical_question}

Provide:

1. Findings
2. Impression
3. Differential Diagnoses
4. Recommendations
"""
)

followup_template = PromptTemplate.from_template(
    """
Previous Context:

{context}

Question:

{question}

Provide a focused clinical answer.
"""
)

ddx_template = PromptTemplate.from_template(
    """
Imaging Findings:

{findings}

Provide:

1. Differential Diagnoses
2. Ranking by likelihood
3. Distinguishing features
4. Additional tests
"""
)

# ============================================================
# IMAGE UTIL
# ============================================================


def image_to_data_url(image: Image.Image) -> str:
    buffer = BytesIO()
    image.convert("RGB").save(buffer, format="JPEG")
    encoded = base64.b64encode(buffer.getvalue()).decode()
    return f"data:image/jpeg;base64,{encoded}"


def _prompt_to_text(prompt_value: Any) -> str:
    if hasattr(prompt_value, "to_string"):
        return prompt_value.to_string()
    return str(prompt_value)


# ============================================================
# MEDGEMMA RUNNABLE
# ============================================================


class MedGemmaRunnable:
    def __init__(self):
        self.system_prompt = MEDICAL_SYSTEM_PROMPT
        self.current_image: Optional[Image.Image] = None
        self.conversation_history: List[Dict[str, Any]] = []

    def set_image(self, image: Optional[Image.Image]):
        self.current_image = image

    def reset(self):
        self.current_image = None
        self.conversation_history = []

    def _prepare_prompt(self, prompt: str) -> str:
        if not self.conversation_history:
            return f"{self.system_prompt.strip()}\n\n{prompt}"
        return prompt

    def invoke(self, prompt: str, max_tokens: int = 1500) -> str:
        prompt = self._prepare_prompt(prompt)
        image = self.current_image

        image_url = None
        if image is not None:
            image_url = image_to_data_url(image)

        self.conversation_history.append({"role": "user", "content": prompt})

        result = query_medgemma(prompt=prompt, image_data=image_url)

        if "error" in result:
            raise RuntimeError(result["error"])

        response = result.get("response", "")
        if not response:
            raise RuntimeError("API returned an empty response.")

        self.conversation_history.append({"role": "assistant", "content": response})
        self.current_image = None
        return response


# Backward-compatible alias for app imports
MedGemmaLLM = MedGemmaRunnable

# ============================================================
# CHATBOT
# ============================================================


class LangChainMedicalChatbot:
    def __init__(self, llm: MedGemmaRunnable):
        self.llm = llm
        self.chat_history = InMemoryChatMessageHistory()
        self.stats = {
            "images_analyzed": 0,
            "queries_processed": 0,
            "findings_detected": [],
            "session_start": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }
        self._create_chains()

    def _create_chains(self):
        self.image_analysis_chain = image_analysis_template | RunnableLambda(
            lambda prompt: self.llm.invoke(_prompt_to_text(prompt))
        )
        self.followup_chain = followup_template | RunnableLambda(
            lambda prompt: self.llm.invoke(_prompt_to_text(prompt))
        )
        self.ddx_chain = ddx_template | RunnableLambda(
            lambda prompt: self.llm.invoke(_prompt_to_text(prompt))
        )

    def analyze_image(
        self,
        image: Image.Image,
        image_type: str,
        clinical_question: str,
    ) -> Dict:
        self.llm.set_image(image)

        response = self.image_analysis_chain.invoke(
            {
                "image_type": image_type,
                "clinical_question": clinical_question,
            }
        )

        self.chat_history.add_message(HumanMessage(content=clinical_question))
        self.chat_history.add_message(AIMessage(content=response))

        self.stats["images_analyzed"] += 1
        self.stats["queries_processed"] += 1

        findings = self._detect_findings(response)
        if findings:
            self.stats["findings_detected"].extend(findings)

        return {
            "response": response,
            "findings": findings,
            "has_image": True,
            "chain_used": "image_analysis_chain",
            "timestamp": datetime.now().strftime("%H:%M:%S"),
        }

    def ask_followup(self, question: str) -> Dict:
        context = "\n".join([msg.content for msg in self.chat_history.messages][-10:])

        response = self.followup_chain.invoke(
            {
                "question": question,
                "context": context,
            }
        )

        self.chat_history.add_message(HumanMessage(content=question))
        self.chat_history.add_message(AIMessage(content=response))

        self.stats["queries_processed"] += 1

        return {
            "response": response,
            "has_image": False,
            "chain_used": "followup_chain",
            "timestamp": datetime.now().strftime("%H:%M:%S"),
        }

    def get_differential_diagnosis(self, findings: str) -> Dict:
        response = self.ddx_chain.invoke({"findings": findings})

        self.stats["queries_processed"] += 1

        return {
            "response": response,
            "has_image": False,
            "chain_used": "ddx_chain",
            "timestamp": datetime.now().strftime("%H:%M:%S"),
        }

    def chat(self, message: str, image: Optional[Image.Image] = None) -> Dict:
        if image:
            self.llm.set_image(image)

        response = self.llm.invoke(message)

        self.chat_history.add_message(HumanMessage(content=message))
        self.chat_history.add_message(AIMessage(content=response))

        self.stats["queries_processed"] += 1
        if image is not None:
            self.stats["images_analyzed"] += 1

        findings = self._detect_findings(response)
        if findings:
            self.stats["findings_detected"].extend(findings)

        return {
            "response": response,
            "findings": findings,
            "has_image": image is not None,
            "chain_used": "direct_llm",
            "timestamp": datetime.now().strftime("%H:%M:%S"),
        }

    def _detect_findings(self, text: str) -> List[str]:
        keywords = [
            "fracture",
            "lesion",
            "mass",
            "opacity",
            "consolidation",
            "pneumonia",
            "atelectasis",
            "effusion",
            "pneumothorax",
            "cardiomegaly",
            "edema",
            "nodule",
            "ischemia",
        ]
        text_lower = text.lower()
        return list(set(k.capitalize() for k in keywords if k in text_lower))

    def get_memory(self) -> str:
        memory = [
            {"type": type(msg).__name__, "content": msg.content}
            for msg in self.chat_history.messages
        ]
        return json.dumps(memory, indent=2)

    def get_stats(self) -> Dict:
        return {
            "Session Start": self.stats["session_start"],
            "Images Analyzed": self.stats["images_analyzed"],
            "Queries Processed": self.stats["queries_processed"],
            "Unique Findings": list(set(self.stats["findings_detected"])),
            "Total Findings": len(self.stats["findings_detected"]),
        }

    def reset(self):
        self.chat_history.clear()
        self.llm.reset()
        self.stats = {
            "images_analyzed": 0,
            "queries_processed": 0,
            "findings_detected": [],
            "session_start": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }